Changeset 449 in code

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

Location:

trunk

Files:

• cmd/soju/main.go (modified) (7 diffs)
• server.go (modified) (8 diffs)

trunk/cmd/soju/main.go

@@ -8 +8 @@
         "net/http"
         "net/url"
+        "os"
+        "os/signal"
         "strings"
+        "syscall"
 
         "github.com/pires/go-proxyproto"
@@ -90 +93 @@
                         ln = proxyProtoListener(ln, srv)
                         go func() {
-                                log.Fatal(srv.Serve(ln))
+                                if err := srv.Serve(ln); err != nil {
+                                        log.Printf("serving %q: %v", listen, err)
+                                }
                         }()
                 case "irc+insecure":
@@ -103 +108 @@
                         ln = proxyProtoListener(ln, srv)
                         go func() {
-                                log.Fatal(srv.Serve(ln))
+                                if err := srv.Serve(ln); err != nil {
+                                        log.Printf("serving %q: %v", listen, err)
+                                }
                         }()
                 case "wss":
@@ -116 +123 @@
                         }
                         go func() {
-                                log.Fatal(httpSrv.ListenAndServeTLS("", ""))
+                                if err := httpSrv.ListenAndServeTLS("", ""); err != nil {
+                                        log.Fatalf("serving %q: %v", listen, err)
+                                }
                         }()
                 case "ws+insecure":
@@ -128 +137 @@
                         }
                         go func() {
-                                log.Fatal(httpSrv.ListenAndServe())
+                                if err := httpSrv.ListenAndServe(); err != nil {
+                                        log.Fatalf("serving %q: %v", listen, err)
+                                }
                         }()
                 case "ident":
@@ -145 +156 @@
                         ln = proxyProtoListener(ln, srv)
                         go func() {
-                                log.Fatal(srv.Identd.Serve(ln))
+                                if err := srv.Identd.Serve(ln); err != nil {
+                                        log.Printf("serving %q: %v", listen, err)
+                                }
                         }()
                 default:
@@ -153 +166 @@
                 log.Printf("server listening on %q", listen)
         }
-        log.Fatal(srv.Run())
+
+        sigCh := make(chan os.Signal, 1)
+        signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)
+
+        if err := srv.Start(); err != nil {
+                log.Fatal(err)
+        }
+
+        <-sigCh
+        log.Print("shutting down server")
+        srv.Shutdown()
 }
 

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)
                 }