From 2f597ebc0b5d7493b66949712d6418e9f6fbd903 Mon Sep 17 00:00:00 2001
From: "diamondburned (Forefront)" <datutbrus@gmail.com>
Date: Mon, 6 Apr 2020 19:36:06 -0700
Subject: [PATCH] Session: Added an AfterClose handler as a *session.Close
 event

---
 gateway/gateway.go          | 12 +++++++++++-
 gateway/integration_test.go |  3 +++
 session/session.go          | 19 +++++++++++++++++++
 3 files changed, 33 insertions(+), 1 deletion(-)

diff --git a/gateway/gateway.go b/gateway/gateway.go
index a11b32b..1865c7a 100644
--- a/gateway/gateway.go
+++ b/gateway/gateway.go
@@ -85,6 +85,11 @@ type Gateway struct {
 
 	ErrorLog func(err error) // default to log.Println
 
+	// AfterClose is called after each close. Error can be non-nil, as this is
+	// called even when the Gateway is gracefully closed. It's used mainly for
+	// reconnections or any type of connection interruptions.
+	AfterClose func(err error) // noop by default
+
 	// FatalError is where Reconnect errors will go to. When an error is sent
 	// here, the Gateway is already dead, so Close() shouldn't be called.
 	// This channel is buffered once.
@@ -127,6 +132,7 @@ func NewGatewayWithDriver(token string, driver json.Driver) (*Gateway, error) {
 		Identifier: DefaultIdentifier(token),
 		Sequence:   NewSequence(),
 		ErrorLog:   WSError,
+		AfterClose: func(error) {},
 		fatalError: make(chan error, 1),
 	}
 	g.FatalError = g.fatalError
@@ -152,6 +158,8 @@ func (g *Gateway) Close() error {
 	// Check if the WS is already closed:
 	if g.waitGroup == nil && g.paceDeath == nil {
 		WSDebug("Gateway is already closed.")
+
+		g.AfterClose(nil)
 		return nil
 	}
 
@@ -177,7 +185,9 @@ func (g *Gateway) Close() error {
 	g.waitGroup = nil
 
 	// Stop the Websocket
-	return g.WS.Close()
+	err := g.WS.Close()
+	g.AfterClose(err)
+	return err
 }
 
 // Reconnects and resumes.
diff --git a/gateway/integration_test.go b/gateway/integration_test.go
index 6522639..4678dc8 100644
--- a/gateway/integration_test.go
+++ b/gateway/integration_test.go
@@ -52,6 +52,9 @@ func TestIntegration(t *testing.T) {
 	if err != nil {
 		t.Fatal("Failed to make a Gateway:", err)
 	}
+	g.AfterClose = func(err error) {
+		log.Println("Closed.")
+	}
 	gateway = g
 
 	if err := g.Open(); err != nil {
diff --git a/session/session.go b/session/session.go
index bf4768f..35ee8ce 100644
--- a/session/session.go
+++ b/session/session.go
@@ -13,6 +13,18 @@ import (
 	"github.com/pkg/errors"
 )
 
+// Closed is an event that's sent to Session's command handler. This works by
+// using (*Gateway).AfterError. If the user sets this callback, no Closed events
+// would be sent.
+//
+// Usage
+//
+//    ses.AddHandler(func(*session.Closed) {})
+//
+type Closed struct {
+	Error error
+}
+
 var ErrMFA = errors.New("Account has 2FA enabled")
 
 // Session manages both the API and Gateway. As such, Session inherits all of
@@ -92,6 +104,13 @@ func (s *Session) Open() error {
 	s.hstop = stop
 	go s.startHandler(stop)
 
+	// Set the AfterClose's handler.
+	s.Gateway.AfterClose = func(err error) {
+		s.Handler.Call(&Closed{
+			Error: err,
+		})
+	}
+
 	if err := s.Gateway.Open(); err != nil {
 		return errors.Wrap(err, "Failed to start gateway")
 	}