diff --git a/services/socketwrap/client.go b/services/socketwrap/client.go index ce8dbb5e8..7d354453c 100644 --- a/services/socketwrap/client.go +++ b/services/socketwrap/client.go @@ -14,6 +14,11 @@ type Client struct { Send chan *models.Action } +func (c *Client) Close() { + close(c.Send) + c.Conn.Close() +} + func (c *Client) WritePump() { defer func() { diff --git a/services/socketwrap/clientManager.go b/services/socketwrap/clientManager.go index d98a7312e..2330aa134 100644 --- a/services/socketwrap/clientManager.go +++ b/services/socketwrap/clientManager.go @@ -1,7 +1,12 @@ package socketwrap import ( + "os" + "os/signal" + "syscall" + "code.gitea.io/gitea/models" + "code.gitea.io/gitea/modules/log" ) type ClientsManager struct { @@ -22,6 +27,9 @@ var LastActionsQueue = NewSyncQueue(20) func (h *ClientsManager) Run() { initActionQueue() + sig := make(chan os.Signal, 1) + signal.Notify(sig, os.Interrupt, syscall.SIGTERM) + var signalsReceived uint for { select { case client := <-h.Register: @@ -41,7 +49,17 @@ func (h *ClientsManager) Run() { delete(h.Clients, client) } } + case s := <-sig: + log.Info("received signal", s) + signalsReceived++ + if signalsReceived < 2 { + for client, _ := range h.Clients { + delete(h.Clients, client) + client.Close() + } + break + } } } }