From 3a5af241e76ab9b57c3a59257c07855e17d3fbf8 Mon Sep 17 00:00:00 2001
From: ken <ken8li@126.com>
Date: Fri, 4 Apr 2025 20:33:44 +0800
Subject: [PATCH] client close nicely

---
 src/agent/client.go  | 10 +++++++---
 src/workflow/task.go |  2 +-
 2 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/src/agent/client.go b/src/agent/client.go
index cf0bda4..cadfed2 100644
--- a/src/agent/client.go
+++ b/src/agent/client.go
@@ -49,6 +49,7 @@ type client struct {
 	wantResponseId uint64
 	wantResponse   chan []byte
 	lock           sync.Mutex
+	once           sync.Once
 }
 
 func newClient(conn *websocket.Conn, a *Agent) *client {
@@ -142,9 +143,12 @@ func (c *client) writePump() {
 }
 
 func (c *client) Close() {
-	c.fnCancel()
-	c.agent.removeClient(c)
-	c.conn.Close()
+	once.Do(func(){
+		c.fnCancel()
+		c.agent.removeClient(c)
+		c.conn.Close()
+		log.Info("client close", "remote", c.conn.RemoteAddr().String(), "connected", c.agent.clients.Size())
+	})
 }
 
 func (c *client) WriteJson(data interface{}, block bool) error {
diff --git a/src/workflow/task.go b/src/workflow/task.go
index 2e13ac5..772cd34 100644
--- a/src/workflow/task.go
+++ b/src/workflow/task.go
@@ -47,8 +47,8 @@ func (task *Task) SendExecFail(errMsg string) {
 
 func (task *Task) do() {
 	defer func() {
+		<-task.pool.sem
 		task.pool.remove(task.client)
-		<-task.pool.sem 
 		if err := recover(); err != nil {
 			log.Error("[Task] do panic", "err", err)
 		}