qol change

ws/registry.go

@@ -79,22 +79,6 @@ func (r *Registry) RegisterMod(channelID, serverID string, conn *websocket.Conn)
 	// caller should use FlushChannelWithSender to perform actual sends
 }
 
-// UnregisterMod : remove mod for a channel
-func (r *Registry) UnregisterMod(channelID string) {
-	r.mu.RLock()
-	e := r.entries[channelID]
-	r.mu.RUnlock()
-	if e == nil {
-		return
-	}
-	e.mu.Lock()
-	defer e.mu.Unlock()
-	if e.Mod != nil && e.Mod.Conn != nil {
-		_ = e.Mod.Conn.Close()
-	}
-	e.Mod = nil
-}
-
 // RegisterBot : single connection for bot. after registration call FlushAllToBotWithSender
 func (r *Registry) RegisterBot(conn *websocket.Conn) {
 	r.botMu.Lock()
@@ -105,13 +89,45 @@ func (r *Registry) RegisterBot(conn *websocket.Conn) {
 	r.botMu.Unlock()
 }
 
+func (r *Registry) UnregisterMod(channelID string) {
+	r.mu.RLock()
+	e := r.entries[channelID]
+	r.mu.RUnlock()
+	if e == nil {
+		return
+	}
+
+	e.mu.Lock()
+	modConn := e.Mod
+	e.Mod = nil
+	e.mu.Unlock()
+
+	if modConn != nil && modConn.Conn != nil {
+		_ = modConn.Conn.SetWriteDeadline(time.Now().Add(time.Second))
+		_ = modConn.Conn.WriteControl(
+			websocket.CloseMessage,
+			websocket.FormatCloseMessage(websocket.CloseNormalClosure, "Disconnecting."),
+			time.Now().Add(time.Second),
+		)
+		_ = modConn.Conn.Close()
+	}
+}
+
 func (r *Registry) UnregisterBot() {
 	r.botMu.Lock()
-	if r.bot != nil && r.bot.Conn != nil {
-		_ = r.bot.Conn.Close()
-	}
+	botConn := r.bot
 	r.bot = nil
 	r.botMu.Unlock()
+
+	if botConn != nil && botConn.Conn != nil {
+		_ = botConn.Conn.SetWriteDeadline(time.Now().Add(time.Second))
+		_ = botConn.Conn.WriteControl(
+			websocket.CloseMessage,
+			websocket.FormatCloseMessage(websocket.CloseNormalClosure, "Disconnecting."),
+			time.Now().Add(time.Second),
+		)
+		_ = botConn.Conn.Close()
+	}
 }
 
 func (r *Registry) Send(channelID string, out GatewayMessageOut, sendOverConn func(*websocket.Conn, GatewayMessageOut) error) (delivered bool, queued bool, err error) {