From b67f4b04aa4a9ce638ef1909b73eb3aeec3f66e4 Mon Sep 17 00:00:00 2001 From: Konstantin Pereiaslov Date: Mon, 23 Sep 2024 01:16:04 -0500 Subject: [PATCH] Add exit points when program was interrupted to functions locking the service mutex Fix #14 and potentially #10 --- main.go | 40 +++++++++++++++++-- main_test.go | 9 +++++ test-server/healthcheck-stuck.json | 2 +- test-server/service-stuck-no-healthcheck.json | 12 ++++++ 4 files changed, 59 insertions(+), 4 deletions(-) create mode 100644 test-server/service-stuck-no-healthcheck.json diff --git a/main.go b/main.go index da592a2..55f6dd8 100644 --- a/main.go +++ b/main.go @@ -176,6 +176,9 @@ func startProxy(serviceConfig ServiceConfig) { }(listener) for { + if interrupted { + return + } clientConnection, err := listener.Accept() if err != nil { log.Printf("[%s] Error accepting connection: %v", serviceConfig.Name, err) @@ -191,6 +194,10 @@ func humanReadableConnection(conn net.Conn) string { return fmt.Sprintf("%s->%s", conn.LocalAddr().String(), conn.RemoteAddr().String()) } func handleConnection(clientConnection net.Conn, serviceConfig ServiceConfig) { + if interrupted { + _ = clientConnection.Close() + return + } defer func(clientConnection net.Conn, serviceConfig ServiceConfig) { log.Printf("[%s] Closing client connection %s on port %s", serviceConfig.Name, humanReadableConnection(clientConnection), serviceConfig.ListenPort) err := clientConnection.Close() @@ -218,6 +225,9 @@ func handleConnection(clientConnection net.Conn, serviceConfig ServiceConfig) { } func startServiceIfNotAlreadyRunningAndConnect(serviceConfig ServiceConfig) net.Conn { + if interrupted { + return nil + } var serviceConnection net.Conn runningService, found := resourceManager.maybeGetRunningService(serviceConfig.Name) if !found { @@ -229,6 +239,9 @@ func startServiceIfNotAlreadyRunningAndConnect(serviceConfig ServiceConfig) net. serviceConnection = serviceConn } else { if !runningService.manageMutex.TryLock() { + if interrupted { + return nil + } //The service could be currently starting or stopping, so let's wait for that to finish and try again runningService.manageMutex.Lock() runningService.manageMutex.Unlock() @@ -273,12 +286,20 @@ func startService(serviceConfig ServiceConfig) (net.Conn, error) { return nil, fmt.Errorf("failed to run command \"%s %s\"", serviceConfig.Command, serviceConfig.Args) } performHealthCheck(serviceConfig) + if interrupted { + return nil, fmt.Errorf("interrupt signal was received") + } var serviceConnection = connectWithWaiting(serviceConfig.ProxyTargetHost, serviceConfig.ProxyTargetPort, serviceConfig.Name, 120*time.Second) - + if interrupted { + return nil, fmt.Errorf("interrupt signal was received") + } runningService.cmd = cmd idleTimeout := getIdleTimeout(serviceConfig) runningService.idleTimer = time.AfterFunc(idleTimeout, func() { + if interrupted { + return + } resourceManager.serviceMutex.Lock() defer resourceManager.serviceMutex.Unlock() @@ -291,6 +312,9 @@ func startService(serviceConfig ServiceConfig) (net.Conn, error) { log.Printf("[%s] Idle timeout %s reached, stopping service", serviceConfig.Name, idleTimeout) stopService(serviceConfig.Name) }) + if interrupted { + return nil, fmt.Errorf("interrupt signal was received") + } resourceManager.storeRunningService(serviceConfig.Name, runningService) return serviceConnection, nil } @@ -350,6 +374,9 @@ func connectToService(serviceConfig ServiceConfig) net.Conn { func connectWithWaiting(serviceHost string, servicePort string, serviceName string, timeout time.Duration) net.Conn { deadline := time.Now().Add(timeout) for time.Now().Before(deadline) { + if interrupted { + return nil + } conn, err := net.DialTimeout("tcp", net.JoinHostPort(serviceHost, servicePort), time.Second) if err == nil { return conn @@ -524,7 +551,12 @@ func forwardConnection(clientConnection net.Conn, serviceConnection net.Conn, se } func stopService(serviceName string) { - resourceManager.runningServices[serviceName].manageMutex.Lock() + if interrupted { + //Shouldn't be necessary, but there might be some locks causing issues + resourceManager.runningServices[serviceName].manageMutex.TryLock() + } else { + resourceManager.runningServices[serviceName].manageMutex.Lock() + } runningService := resourceManager.runningServices[serviceName] if runningService.idleTimer != nil { runningService.idleTimer.Stop() @@ -555,7 +587,9 @@ func stopService(serviceName string) { } releaseResources(runningService.resourceRequirements) - resourceManager.runningServices[serviceName].manageMutex.Unlock() + if !interrupted { + resourceManager.runningServices[serviceName].manageMutex.Unlock() + } delete(resourceManager.runningServices, serviceName) } diff --git a/main_test.go b/main_test.go index e566fcd..fd0f6dc 100644 --- a/main_test.go +++ b/main_test.go @@ -20,6 +20,9 @@ func connectOnly(test *testing.T, proxyAddress string) { test.Error(err) return } + //give large-model-proxy time to start the service, so that it doesn't get killed before it started it + //which can lead to false positive passing tests + time.Sleep(1 * time.Second) } func minimal(test *testing.T, proxyAddress string) { @@ -189,6 +192,12 @@ func TestAppScenarios(test *testing.T) { "2005", connectOnly, }, + { + "service-stuck-no-healthcheck", + "test-server/service-stuck-no-healthcheck.json", + "2006", + connectOnly, + }, } for _, testCase := range tests { diff --git a/test-server/healthcheck-stuck.json b/test-server/healthcheck-stuck.json index ea4d83e..deb52b3 100644 --- a/test-server/healthcheck-stuck.json +++ b/test-server/healthcheck-stuck.json @@ -1,7 +1,7 @@ { "Services": [ { - "Name": "test-server_healthcheck", + "Name": "test-server_healthcheck-stuck", "ListenPort": "2005", "ProxyTargetHost": "localhost", "ProxyTargetPort": "12005", diff --git a/test-server/service-stuck-no-healthcheck.json b/test-server/service-stuck-no-healthcheck.json new file mode 100644 index 0000000..14d33aa --- /dev/null +++ b/test-server/service-stuck-no-healthcheck.json @@ -0,0 +1,12 @@ +{ + "Services": [ + { + "Name": "test-server_service-stuck-no-healthcheck", + "ListenPort": "2006", + "ProxyTargetHost": "localhost", + "ProxyTargetPort": "12006", + "Command": "test-server/test-server", + "Args": "-p 12006 -startup-duration 1d" + } + ] +} \ No newline at end of file