Context fixes

6 files changed, 77 insertions, 26 deletions

diff --git a/forged/internal/incoming/hooks/hooks.go b/forged/internal/incoming/hooks/hooks.go
index 18b9742..c021406 100644
--- a/forged/internal/incoming/hooks/hooks.go
+++ b/forged/internal/incoming/hooks/hooks.go
@@ -1,6 +1,7 @@
 package hooks
 
 import (
+	"context"
 	"fmt"
 	"net"
 
@@ -40,15 +41,25 @@ func New(config Config) (server *Server) {
 	}
 }
 
-func (server *Server) Run() error {
+func (server *Server) Run(ctx context.Context) error {
 	listener, _, err := misc.ListenUnixSocket(server.socketPath)
 	if err != nil {
 		return fmt.Errorf("listen unix socket for hooks: %w", err)
 	}
+	defer func() {
+		_ = listener.Close()
+	}()
+
+	go func() {
+		<-ctx.Done()
+		_ = listener.Close()
+		// TODO: Log the error
+	}()
 
 	for {
 		conn, err := listener.Accept()
 		if err != nil {
+			// TODO: Handle errors caused by context cancel
 			return fmt.Errorf("accept conn: %w", err)
 		}
 
diff --git a/forged/internal/incoming/lmtp/config.go b/forged/internal/incoming/lmtp/config.go
index e6db5a6..c622648 100644
--- a/forged/internal/incoming/lmtp/config.go
+++ b/forged/internal/incoming/lmtp/config.go
@@ -1,6 +1,7 @@
 package lmtp
 
 import (
+	"context"
 	"fmt"
 	"net"
 
@@ -33,11 +34,20 @@ func New(config Config) (server *Server) {
 	}
 }
 
-func (server *Server) Run() error {
+func (server *Server) Run(ctx context.Context) error {
 	listener, _, err := misc.ListenUnixSocket(server.socket)
 	if err != nil {
 		return fmt.Errorf("listen unix socket for LMTP: %w", err)
 	}
+	defer func() {
+		_ = listener.Close()
+	}()
+
+	go func() {
+		<-ctx.Done()
+		_ = listener.Close()
+		// TODO: Log the error
+	}()
 
 	for {
 		conn, err := listener.Accept()
diff --git a/forged/internal/incoming/ssh/ssh.go b/forged/internal/incoming/ssh/ssh.go
index b0269e5..efa24f1 100644
--- a/forged/internal/incoming/ssh/ssh.go
+++ b/forged/internal/incoming/ssh/ssh.go
@@ -1,6 +1,7 @@
 package ssh
 
 import (
+	"context"
 	"fmt"
 	"os"
 
@@ -53,8 +54,19 @@ func New(config Config) (server *Server, err error) {
 	return
 }
 
-func (server *Server) Run() (err error) {
+func (server *Server) Run(ctx context.Context) (err error) {
 	listener, err := misc.Listen(server.net, server.addr)
+	defer func() {
+		_ = listener.Close()
+	}()
+
+	go func() {
+		<-ctx.Done()
+		_ = server.gliderServer.Close()
+		_ = listener.Close() // unnecessary?
+		// TODO: Log the error
+	}()
+
 	if err = server.gliderServer.Serve(listener); err != nil {
 		return fmt.Errorf("serve SSH: %w", err)
 	}
diff --git a/forged/internal/incoming/web/web.go b/forged/internal/incoming/web/web.go
index e4675f0..a95b650 100644
--- a/forged/internal/incoming/web/web.go
+++ b/forged/internal/incoming/web/web.go
@@ -1,6 +1,7 @@
 package web
 
 import (
+	"context"
 	"fmt"
 	"net/http"
 	"time"
@@ -48,8 +49,19 @@ func New(config Config) (server *Server) {
 	}
 }
 
-func (server *Server) Run() (err error) {
+func (server *Server) Run(ctx context.Context) (err error) {
 	listener, err := misc.Listen(server.net, server.addr)
+	defer func() {
+		_ = listener.Close()
+	}()
+
+	go func() {
+		<-ctx.Done()
+		_ = server.httpServer.Close()
+		_ = listener.Close() // unnecessary?
+		// TODO: Log the error
+	}()
+
 	if err = server.httpServer.Serve(listener); err != nil {
 		return fmt.Errorf("serve web: %w", err)
 	}
diff --git a/forged/internal/server/server.go b/forged/internal/server/server.go
index b3cd92a..0b14e62 100644
--- a/forged/internal/server/server.go
+++ b/forged/internal/server/server.go
@@ -3,7 +3,6 @@ package server
 import (
 	"context"
 	"fmt"
-	"log"
 
 	"go.lindenii.runxiyu.org/forge/forged/internal/config"
 	"go.lindenii.runxiyu.org/forge/forged/internal/database"
@@ -29,7 +28,7 @@ type Server struct {
 	}
 }
 
-func New(ctx context.Context, configPath string) (server *Server, err error) {
+func New(configPath string) (server *Server, err error) {
 	server = &Server{}
 
 	server.config, err = config.Open(configPath)
@@ -37,11 +36,6 @@ func New(ctx context.Context, configPath string) (server *Server, err error) {
 		return server, fmt.Errorf("open config: %w", err)
 	}
 
-	server.database, err = database.Open(ctx, server.config.DB)
-	if err != nil {
-		return server, fmt.Errorf("open database: %w", err)
-	}
-
 	server.hookServer = hooks.New(server.config.Hooks)
 	server.lmtpServer = lmtp.New(server.config.LMTP)
 	server.webServer = web.New(server.config.Web)
@@ -53,33 +47,45 @@ func New(ctx context.Context, configPath string) (server *Server, err error) {
 	return server, nil
 }
 
-func (s *Server) Run() error {
+func (server *Server) Run(ctx context.Context) (err error) {
 	// TODO: Not running git2d because it should be run separately.
 	// This needs to be documented somewhere, hence a TODO here for now.
 
+	subCtx, cancel := context.WithCancel(ctx)
+	defer cancel()
+
+	server.database, err = database.Open(subCtx, server.config.DB)
+	if err != nil {
+		return fmt.Errorf("open database: %w", err)
+	}
+
+	errCh := make(chan error)
+
 	go func() {
-		if err := s.hookServer.Run(); err != nil {
-			log.Fatalf("run hook server: %v", err)
-		}
+		err := server.hookServer.Run(subCtx)
+		errCh <- err
 	}()
 
 	go func() {
-		if err := s.lmtpServer.Run(); err != nil {
-			log.Fatalf("run LMTP server: %v", err)
-		}
+		err := server.lmtpServer.Run(subCtx)
+		errCh <- err
 	}()
 
 	go func() {
-		if err := s.webServer.Run(); err != nil {
-			log.Fatalf("run web server: %v", err)
-		}
+		err := server.webServer.Run(subCtx)
+		errCh <- err
 	}()
 
 	go func() {
-		if err := s.sshServer.Run(); err != nil {
-			log.Fatalf("run SSH server: %v", err)
-		}
+		err := server.sshServer.Run(subCtx)
+		errCh <- err
 	}()
 
+	select {
+	case err := <-errCh:
+		return fmt.Errorf("server error: %w", err)
+	case <-ctx.Done():
+	}
+
 	return nil
 }
diff --git a/forged/main.go b/forged/main.go
index e9609f3..38e22ff 100644
--- a/forged/main.go
+++ b/forged/main.go
@@ -19,10 +19,10 @@ func main() {
 	)
 	flag.Parse()
 
-	s, err := server.New(context.Background(), *configPath)
+	s, err := server.New(*configPath)
 	if err != nil {
 		panic(err)
 	}
 
-	panic(s.Run())
+	panic(s.Run(context.Background()))
 }