From 8b189e069b713137ac49f83f32863f6c569310fe Mon Sep 17 00:00:00 2001
From: luke-lombardi <33990301+luke-lombardi@users.noreply.github.com>
Date: Tue, 16 Apr 2024 18:37:45 -0400
Subject: [PATCH] Serve exit fixes & debug logs (#141)

Co-authored-by: Luke Lombardi <luke@beam.cloud>
---
 internal/abstractions/endpoint/buffer.go | 7 +++++++
 sdk/src/beta9/runner/serve.py            | 7 ++++---
 2 files changed, 11 insertions(+), 3 deletions(-)

diff --git a/internal/abstractions/endpoint/buffer.go b/internal/abstractions/endpoint/buffer.go
index eb32378e1..384f7c103 100644
--- a/internal/abstractions/endpoint/buffer.go
+++ b/internal/abstractions/endpoint/buffer.go
@@ -6,6 +6,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"io"
+	"log"
 	"net/http"
 	"sort"
 	"sync"
@@ -168,6 +169,8 @@ func (rb *RequestBuffer) discoverContainers() {
 						return
 					}
 
+					log.Printf("<%s> Retrieved container address: %s\n", cs.ContainerId, containerAddress)
+
 					inFlight, err := rb.requestsInFlight(cs.ContainerId)
 					if err != nil {
 						return
@@ -179,7 +182,11 @@ func (rb *RequestBuffer) discoverContainers() {
 							address:  containerAddress,
 							inFlight: inFlight,
 						}
+						return
 					}
+
+					log.Printf("<%s> Container not ready\n", cs.ContainerId)
+
 				}(containerState)
 			}
 
diff --git a/sdk/src/beta9/runner/serve.py b/sdk/src/beta9/runner/serve.py
index c5de0e24f..44ec34d8a 100644
--- a/sdk/src/beta9/runner/serve.py
+++ b/sdk/src/beta9/runner/serve.py
@@ -1,4 +1,3 @@
-import atexit
 import os
 import signal
 import subprocess
@@ -47,6 +46,9 @@ def __init__(self) -> None:
         self.restart_event = Event()
         self.exit_event = Event()
 
+        # Register signal handlers
+        signal.signal(signal.SIGTERM, self.shutdown)
+
         # Set up the file change event handler & observer
         self.event_handler = SyncEventHandler(self.trigger_restart)
         self.observer = PollingObserver()
@@ -56,7 +58,7 @@ def __init__(self) -> None:
     def shutdown(self, signum=None, frame=None) -> None:
         self.kill_subprocess()
         self.observer.stop()
-        self.observer.join()
+        self.observer.join(timeout=0.1)
         self.exit_event.set()
         self.restart_event.set()
 
@@ -117,7 +119,6 @@ def _command() -> List[str]:
 
 if __name__ == "__main__":
     sg = ServeGateway()
-    atexit.register(sg.shutdown)
 
     try:
         sg.run(command=_command())