Fix: shutdown by exiting scheduler with context (#676)

cmd/gateway/main.go

@@ -2,9 +2,6 @@ package main
 
 import (
 	"log"
-	"os"
-	"os/signal"
-	"syscall"
 
 	"github.com/beam-cloud/beta9/pkg/gateway"
 )
@@ -15,12 +12,6 @@ func main() {
 		log.Fatalf("Error creating gateway service: %+v\n", err)
 	}
 
-	stop := make(chan os.Signal, 1)
-	signal.Notify(stop, syscall.SIGTERM)
-
 	gw.Start()
-
-	<-stop
-
 	log.Println("Gateway stopped")
 }

pkg/scheduler/scheduler.go

@@ -197,6 +197,14 @@ func (s *Scheduler) getControllers(request *types.ContainerRequest) ([]WorkerPoo
 
 func (s *Scheduler) StartProcessingRequests() {
 	for {
+		select {
+		case <-s.ctx.Done():
+			// Context has been cancelled
+			return
+		default:
+			// Continue processing requests
+		}
+
 		if s.requestBacklog.Len() == 0 {
 			time.Sleep(requestProcessingInterval)
 			continue

pkg/scheduler/scheduler_test.go

@@ -59,6 +59,7 @@ func NewSchedulerForTest() (*Scheduler, error) {
 	}
 
 	return &Scheduler{
+		ctx:               context.Background(),
 		eventBus:          eventBus,
 		workerRepo:        workerRepo,
 		workerPoolManager: workerPoolManager,