Refactors context usage at server entity

server/server.go

@@ -3,8 +3,8 @@ package server
 import (
 	"context"
 	"fmt"
+	"net"
 	"net/http"
-	"strconv"
 	"sync"
 	"time"
 
@@ -18,21 +18,31 @@ const (
 )
 
 type Server struct {
-	mutex     *sync.Mutex
-	channels  []wasabi.Channel
-	port      uint16
-	isRunning bool
+	mutex    *sync.Mutex
+	channels []wasabi.Channel
+	addr     string
+	http     *http.Server
+	baseCtx  context.Context
 }
 
+type ServerOption func(*Server)
+
 // NewServer creates new instance of Wasabi server
 // port - port to listen on
 // returns new instance of Server
-func NewServer(port uint16) *Server {
-	return &Server{
-		port:     port,
+func NewServer(addr string, opts ...ServerOption) *Server {
+	server := &Server{
+		addr:     addr,
 		channels: make([]wasabi.Channel, 0, 1),
 		mutex:    &sync.Mutex{},
+		baseCtx:  context.Background(),
+	}
+
+	for _, opt := range opts {
+		opt(server)
 	}
+
+	return server
 }
 
 // AddChannel adds new channel to server
@@ -43,63 +53,45 @@ func (s *Server) AddChannel(channel wasabi.Channel) {
 // Run starts server
 // ctx - context
 // returns error if any
-func (s *Server) Run(ctx context.Context) error {
-	s.mutex.Lock()
-	defer s.mutex.Unlock()
-
-	if s.isRunning {
+func (s *Server) Run() error {
+	if !s.mutex.TryLock() {
 		return fmt.Errorf("server is already running")
 	}
 
-	listen := ":" + strconv.Itoa(int(s.port))
-
-	execCtx, cancel := context.WithCancel(ctx)
-
-	defer cancel()
+	defer s.mutex.Unlock()
 
 	mux := http.NewServeMux()
 
 	for _, channel := range s.channels {
-		channel.SetContext(execCtx)
 		mux.Handle(
 			channel.Path(),
 			channel.Handler(),
 		)
 	}
 
-	slog.Info("Starting app server on " + listen)
+	slog.Info("Starting app server on " + s.addr)
 
-	server := &http.Server{
-		Addr:              listen,
+	s.http = &http.Server{
+		Addr:              s.addr,
 		ReadHeaderTimeout: ReadHeaderTimeout,
 		ReadTimeout:       ReadTimeout,
 		Handler:           mux,
+		BaseContext: func(_ net.Listener) context.Context {
+			return s.baseCtx
+		},
 	}
 
-	wg := &sync.WaitGroup{}
-	wg.Add(1)
-
-	go func() {
-		<-execCtx.Done()
-
-		slog.Info("Shutting down app server on " + listen)
-
-		if err := server.Shutdown(context.Background()); err != nil {
-			slog.Error("Failed to shutdown app server", "error", err)
-		}
-
-		wg.Done()
-	}()
-
-	err := server.ListenAndServe()
-
-	cancel()
-
-	wg.Wait()
+	return s.http.ListenAndServe()
+}
 
-	if err != http.ErrServerClosed {
-		return err
+// BaseContext optionally specifies based context that will be used for all connections.
+// If not specified, context.Background() will be used.
+func WithBaseContext(ctx context.Context) ServerOption {
+	if ctx == nil {
+		panic("nil context")
 	}
 
-	return nil
+	return func(s *Server) {
+		s.baseCtx = ctx
+	}
 }