Changeset 449 in code for trunk/server.go

Timestamp:

Feb 9, 2021, 4:34:46 PM (4 years ago)

Author:

contact

Message:

Add support for graceful shutdown

Closes: ​https://todo.sr.ht/~emersion/soju/45

File:

• trunk/server.go (modified) (8 diffs)

trunk/server.go

@@ -6 +6 @@
         "net"
         "net/http"
+        "strings"
         "sync"
         "sync/atomic"
@@ -55 +56 @@
         Identd         *Identd // can be nil
 
-        db *DB
-
-        lock  sync.Mutex
-        users map[string]*user
+        db     *DB
+        stopWG sync.WaitGroup
+
+        lock      sync.Mutex
+        listeners map[net.Listener]struct{}
+        users     map[string]*user
 }
 
@@ -65 +68 @@
                 Logger:       log.New(log.Writer(), "", log.LstdFlags),
                 HistoryLimit: 1000,
+                db:           db,
+                listeners:    make(map[net.Listener]struct{}),
                 users:        make(map[string]*user),
-                db:           db,
         }
 }
@@ -74 +78 @@
 }
 
-func (s *Server) Run() error {
+func (s *Server) Start() error {
         users, err := s.db.ListUsers()
         if err != nil {
@@ -86 +90 @@
         s.lock.Unlock()
 
-        select {}
+        return nil
+}
+
+func (s *Server) Shutdown() {
+        s.lock.Lock()
+        for ln := range s.listeners {
+                if err := ln.Close(); err != nil {
+                        s.Logger.Printf("failed to stop listener: %v", err)
+                }
+        }
+        for _, u := range s.users {
+                u.events <- eventStop{}
+        }
+        s.lock.Unlock()
+
+        s.stopWG.Wait()
 }
 
@@ -117 +136 @@
         s.users[u.Username] = u
 
+        s.stopWG.Add(1)
+
         go func() {
                 u.run()
@@ -123 +144 @@
                 delete(s.users, u.Username)
                 s.lock.Unlock()
+
+                s.stopWG.Done()
         }()
 
@@ -146 +169 @@
 
 func (s *Server) Serve(ln net.Listener) error {
+        s.lock.Lock()
+        s.listeners[ln] = struct{}{}
+        s.lock.Unlock()
+
+        s.stopWG.Add(1)
+
+        defer func() {
+                s.lock.Lock()
+                delete(s.listeners, ln)
+                s.lock.Unlock()
+
+                s.stopWG.Done()
+        }()
+
         for {
                 conn, err := ln.Accept()
-                if err != nil {
+                // TODO: use net.ErrClosed when available
+                if err != nil && strings.Contains(err.Error(), "use of closed network connection") {
+                        return nil
+                } else if err != nil {
                         return fmt.Errorf("failed to accept connection: %v", err)
                 }