[patch] refactored SSE to cleanup and organize better (#55)

Author: bnkamalesh

extensions/sse/client.go

@@ -67,9 +67,20 @@ func (cs *Clients) listener(events <-chan event) {
 			ev.Response <- &eventResponse{
 				Clients: copied,
 			}
+
 		case eTypeRemoveClient:
+			cli := cs.clients[ev.ClientID]
+			if cli == nil {
+				ev.Response <- nil
+				return
+			}
+
+			close(cli.Msg)
+			cli.Ctx.Done()
+
 			delete(cs.clients, ev.ClientID)
 			ev.Response <- nil
+
 		case eTypeClient:
 			ev.Response <- &eventResponse{
 				Client: cs.clients[ev.ClientID],
@@ -86,6 +97,7 @@ func (cs *Clients) New(ctx context.Context, w http.ResponseWriter, clientID stri
 		ResponseWriter: w,
 		Ctx:            ctx,
 	}
+
 	cs.events <- event{
 		Type:   eTypeNewClient,
 		Client: cli,
@@ -149,11 +161,12 @@ func (cs *Clients) Client(clientID string) *Client {
 }
 
 func NewClientManager() ClientManager {
-	events := make(chan event, 10)
+	const buffer = 10
+	events := make(chan event, buffer)
 	cli := &Clients{
 		clients:   make(map[string]*Client),
 		events:    events,
-		MsgBuffer: 10,
+		MsgBuffer: buffer,
 	}
 	go cli.listener(events)
 	return cli

extensions/sse/sse.go

@@ -24,7 +24,7 @@ type SSE struct {
 
 	// OnSend is a hook, which is called *after* a message is sent to a client
 	OnSend func(ctx context.Context, client *Client, err error)
-	// BeforeSend is a hook, which is called right before a message is sent to a client
+	// BeforeSend is a hook, which is called before starting to listen for messages to send
 	BeforeSend func(ctx context.Context, client *Client)
 
 	Clients ClientManager
@@ -101,12 +101,6 @@ func (sse *SSE) ActiveClients() int {
 }
 
 func (sse *SSE) RemoveClient(ctx context.Context, clientID string) {
-	cli := sse.Clients.Client(clientID)
-	if cli != nil {
-		close(cli.Msg)
-		cli.Ctx.Done()
-	}
-
 	sse.OnRemoveClient(
 		ctx,
 		clientID,