From f1f1c854f6c39c62fb76f43bba2978763ae24c2b Mon Sep 17 00:00:00 2001
From: miko <michko@gmail.com>
Date: Thu, 23 Feb 2023 10:29:09 +0100
Subject: [PATCH] Add unmount handler (#63)

---
 engine.go  | 12 ++++++++++++
 handler.go | 19 +++++++++++++++++++
 2 files changed, 31 insertions(+)

diff --git a/engine.go b/engine.go
index b5269af..3d084a0 100644
--- a/engine.go
+++ b/engine.go
@@ -26,6 +26,10 @@ type Engine interface {
 	// is called on initial GET request and later when the websocket connects.
 	// Data to render the handler should be fetched here and returned.
 	Mount() MountHandler
+	// UnmountHandler the func that is called by a handler to report that a connection
+	// is closed. This is called on websocket close. Can be used to track number of
+	// connected users.
+	Unmount() UnmountHandler
 	// Params called to handle any incoming paramters after mount.
 	Params() []EventHandler
 	// Render is called to generate the HTML of a Socket. It is defined
@@ -108,6 +112,10 @@ func (e *BaseEngine) Mount() MountHandler {
 	return e.handler.getMount()
 }
 
+func (e *BaseEngine) Unmount() UnmountHandler {
+	return e.handler.getUnmount()
+}
+
 func (e *BaseEngine) Params() []EventHandler {
 	return e.handler.getParams()
 }
@@ -180,6 +188,10 @@ func (e *BaseEngine) DeleteSocket(sock Socket) {
 	e.socketsMu.Lock()
 	defer e.socketsMu.Unlock()
 	delete(e.socketMap, sock.ID())
+	err := e.Unmount()(sock)
+	if err != nil {
+		log.Println("socket unmount error", err)
+	}
 }
 
 // CallEvent route an event to the correct handler.
diff --git a/handler.go b/handler.go
index b434804..ce0b671 100644
--- a/handler.go
+++ b/handler.go
@@ -18,6 +18,11 @@ type HandlerConfig func(h Handler) error
 // in the socket.
 type MountHandler func(ctx context.Context, c Socket) (interface{}, error)
 
+// UnmountHandler the func that is called by a handler to report that a connection
+// is closed. This is called on websocket close. Can be used to track number of
+// connected users.
+type UnmountHandler func(c Socket) error
+
 // RenderHandler the func that is called to render the current state of the
 // data for the socket.
 type RenderHandler func(ctx context.Context, rc *RenderContext) (io.Reader, error)
@@ -39,6 +44,8 @@ type Handler interface {
 	// HandleMount handles initial setup on first request, and then later when
 	// the socket first connets.
 	HandleMount(handler MountHandler)
+	// HandleUnmount used to track webcocket disconnections.
+	HandleUnmount(handler UnmountHandler)
 	// HandleRender used to set the render method for the handler.
 	HandleRender(handler RenderHandler)
 	// HandleError for when an error occurs.
@@ -54,6 +61,7 @@ type Handler interface {
 	HandleParams(handler EventHandler)
 
 	getMount() MountHandler
+	getUnmount() UnmountHandler
 	getRender() RenderHandler
 	getError() ErrorHandler
 	getEvent(t string) (EventHandler, error)
@@ -67,6 +75,8 @@ type BaseHandler struct {
 	// is called on initial GET request and later when the websocket connects.
 	// Data to render the handler should be fetched here and returned.
 	mountHandler MountHandler
+	// unmountHandler used to track webcocket disconnections.
+	unmountHandler UnmountHandler
 	// Render is called to generate the HTML of a Socket. It is defined
 	// by default and will render any template provided.
 	renderHandler RenderHandler
@@ -90,6 +100,9 @@ func NewHandler(configs ...HandlerConfig) *BaseHandler {
 		mountHandler: func(ctx context.Context, s Socket) (interface{}, error) {
 			return nil, nil
 		},
+		unmountHandler: func(s Socket) error {
+			return nil
+		},
 		renderHandler: func(ctx context.Context, rc *RenderContext) (io.Reader, error) {
 			return nil, ErrNoRenderer
 		},
@@ -112,6 +125,9 @@ func NewHandler(configs ...HandlerConfig) *BaseHandler {
 func (h *BaseHandler) HandleMount(f MountHandler) {
 	h.mountHandler = f
 }
+func (h *BaseHandler) HandleUnmount(f UnmountHandler) {
+	h.unmountHandler = f
+}
 func (h *BaseHandler) HandleRender(f RenderHandler) {
 	h.renderHandler = f
 }
@@ -140,6 +156,9 @@ func (h *BaseHandler) HandleParams(handler EventHandler) {
 func (h *BaseHandler) getMount() MountHandler {
 	return h.mountHandler
 }
+func (h *BaseHandler) getUnmount() UnmountHandler {
+	return h.unmountHandler
+}
 func (h *BaseHandler) getRender() RenderHandler {
 	return h.renderHandler
 }