Fix server interrupt (#161)

* pass ctx to startServer

Signed-off-by: npolshakova <[email protected]>

* fix sim test to use ctx

Signed-off-by: npolshakova <[email protected]>

---------

Signed-off-by: npolshakova <[email protected]>

Author: npolshakova

pkg/llm-d-inference-sim/simulator.go

@@ -152,8 +152,8 @@ func (s *VllmSimulator) Start(ctx context.Context) error {
 		return err
 	}
 
-	// start the http server
-	return s.startServer(listener)
+	// start the http server with context support
+	return s.startServer(ctx, listener)
 }
 
 func (s *VllmSimulator) newListener() (net.Listener, error) {
@@ -166,7 +166,7 @@ func (s *VllmSimulator) newListener() (net.Listener, error) {
 }
 
 // startServer starts http server on port defined in command line
-func (s *VllmSimulator) startServer(listener net.Listener) error {
+func (s *VllmSimulator) startServer(ctx context.Context, listener net.Listener) error {
 	r := fasthttprouter.New()
 
 	// support completion APIs
@@ -195,7 +195,33 @@ func (s *VllmSimulator) startServer(listener net.Listener) error {
 		}
 	}()
 
-	return server.Serve(listener)
+	// Start server in a goroutine
+	serverErr := make(chan error, 1)
+	go func() {
+		s.logger.Info("HTTP server starting")
+		serverErr <- server.Serve(listener)
+	}()
+
+	// Wait for either context cancellation or server error
+	select {
+	case <-ctx.Done():
+		s.logger.Info("Shutdown signal received, shutting down HTTP server gracefully")
+
+		// Gracefully shutdown the server
+		if err := server.Shutdown(); err != nil {
+			s.logger.Error(err, "Error during server shutdown")
+			return err
+		}
+
+		s.logger.Info("HTTP server stopped")
+		return nil
+
+	case err := <-serverErr:
+		if err != nil {
+			s.logger.Error(err, "HTTP server failed")
+		}
+		return err
+	}
 }
 
 // Print prints to a log, implementation of fasthttp.Logger

pkg/llm-d-inference-sim/simulator_test.go

@@ -115,7 +115,7 @@ func startServerWithArgsAndMetrics(ctx context.Context, mode string, args []stri
 
 	// start the http server
 	go func() {
-		if err := s.startServer(listener); err != nil {
+		if err := s.startServer(ctx, listener); err != nil {
 			logger.Error(err, "error starting server")
 		}
 	}()