add ability to use context with serve

hscontrol/app.go

@@ -555,7 +555,7 @@ func nodesChangedHook(
 }
 
 // Serve launches the HTTP and gRPC server service Headscale and the API.
-func (h *Headscale) Serve() error {
+func (h *Headscale) Serve(ctx context.Context) error {
 	capver.CanOldCodeBeCleanedUp()
 
 	if profilingEnabled {
@@ -631,7 +631,7 @@ func (h *Headscale) Serve() error {
 
 	// Start all scheduled tasks, e.g. expiring nodes, derp updates and
 	// records updates
-	scheduleCtx, scheduleCancel := context.WithCancel(context.Background())
+	scheduleCtx, scheduleCancel := context.WithCancel(ctx)
 	defer scheduleCancel()
 	go h.scheduledTasks(scheduleCtx)
 
@@ -644,7 +644,6 @@ func (h *Headscale) Serve() error {
 	// Prepare group for running listeners
 	errorGroup := new(errgroup.Group)
 
-	ctx := context.Background()
 	ctx, cancel := context.WithCancel(ctx)
 	defer cancel()
 
@@ -837,7 +836,8 @@ func (h *Headscale) Serve() error {
 	sigFunc := func(c chan os.Signal) {
 		// Wait for a SIGINT or SIGKILL:
 		for {
-			sig := <-c
+			select {
+			case sig := <-c:
 				switch sig {
 				case syscall.SIGHUP:
 					log.Info().
@@ -942,8 +942,19 @@ func (h *Headscale) Serve() error {
 
 					return
 				}
+			case <-ctx.Done():
+				// send signal to kill
+				// could be done a lot better
+				select {
+				case sigc <- os.Kill:
+					fmt.Println("sent kill message")
+				default:
+					fmt.Println("no kill message sent")
 				}
 			}
+		}
+
+	}
 	errorGroup.Go(func() error {
 		sigFunc(sigc)