1 files changed, 66 insertions, 37 deletions
diff --git a/server/server.go b/server/server.go
index 424f359..92f590c 100644
--- a/server/server.go
+++ b/server/server.go
@@ -1,67 +1,96 @@
-// Copyright 2022 Terin Stock.
-// SPDX-License-Identifier: MPL-2.0
+// Copyright 2022 The Kubernetes Authors.
+// SPDX-License-Identifier: Apache-2.0
+/*
+Copyright 2022 The Kubernetes Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+   http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
 
 package server
 
 import (
 	"context"
+	"errors"
+	"log/slog"
 	"net"
 	"net/http"
-	"strconv"
 	"time"
-
-	"github.com/gorilla/mux"
 )
 
+// Server is a general purpose HTTP server Runnable for a manager.
 type Server struct {
-	mux  *mux.Router
-	host string
-	port int
-}
+	// Name is an optional string that describes the purpose of the server. It is used in logs to distinguish
+	// among multiple servers.
+	Name string
 
-type Options struct {
-	Host string
-	Port int
-}
+	// Server is the HTTP server to run. It is required.
+	Server *http.Server
 
-func New(options Options) *Server {
-	return &Server{
-		mux:  mux.NewRouter(),
-		host: options.Host,
-		port: options.Port,
-	}
-}
+	// Listener is an optional listener to use. If not set, the server start a listener using the server.Addr.
+	// Using a listener is useful when the port reservation needs to happen in advance of this runnable starting.
+	Listener net.Listener
 
-func (s *Server) Register(path string, handler http.Handler) {
-	s.mux.Handle(path, handler)
+	// ShutdownTimeout is an optional duration that indicates how long to wait for the server to shutdown gracefully. If not set,
+	// the server will wait indefinitely for all connections to close.
+	ShutdownTimeout *time.Duration
 }
 
+// Start starts the server. It will block until the server is stopped or an error occurs.
 func (s *Server) Start(ctx context.Context) error {
-	ln, err := net.Listen("tcp", net.JoinHostPort(s.host, strconv.Itoa(s.port)))
-	if err != nil {
-		return err
-	}
+	serverShutdown := make(chan struct{})
 
-	srv := &http.Server{
-		Handler:           s.mux,
-		MaxHeaderBytes:    1 << 20,
-		IdleTimeout:       90 * time.Second,
-		ReadHeaderTimeout: 32 * time.Second,
+	logger := slog.With("addr", s.addr())
+	if s.Name != "" {
+		logger = logger.With("name", s.Name)
 	}
 
-	shutdownCh := make(chan struct{})
 	go func() {
 		<-ctx.Done()
-		if err := srv.Shutdown(context.Background()); err != nil {
-			_ = err
+		logger.Info("shutting down server")
+		shutdownCtx := context.Background()
+		if s.ShutdownTimeout != nil {
+			var shutdownCancel context.CancelFunc
+			shutdownCtx, shutdownCancel = context.WithTimeout(context.Background(), *s.ShutdownTimeout)
+			defer shutdownCancel()
 		}
-		close(shutdownCh)
+
+		if err := s.Server.Shutdown(shutdownCtx); err != nil {
+			logger.Error("error shutting down server", "error", err)
+		}
+		close(serverShutdown)
 	}()
 
-	if err := srv.Serve(ln); err != nil && err != http.ErrServerClosed {
+	logger.Info("starting server")
+	if err := s.serve(); err != nil && !errors.Is(err, http.ErrServerClosed) {
 		return err
 	}
 
-	<-shutdownCh
+	<-serverShutdown
 	return nil
 }
+
+func (s *Server) addr() string {
+	if s.Listener != nil {
+		return s.Listener.Addr().String()
+	}
+
+	return s.Server.Addr
+}
+
+func (s *Server) serve() error {
+	if s.Listener != nil {
+		return s.Server.Serve(s.Listener)
+	}
+
+	return s.Server.ListenAndServe()
+}