fixup! refactor(container): restart handling

internal/container/container.go

@@ -377,23 +377,30 @@ func (a *Client) runContainerLoop(
 			respC, errC := a.apiClient.ContainerWait(ctx, containerID, container.WaitConditionNextExit)
 			select {
 			case resp := <-respC:
-				if shouldRestartFunc != nil {
+				exit := func(err error) {
-					shouldRestart, err := shouldRestartFunc(resp.StatusCode, restartCount, time.Since(startedWaitingAt))
+					a.logger.Info("Container exited", "id", shortID(containerID), "should_restart", "false", "exit_code", resp.StatusCode, "restart_count", restartCount)
-					if shouldRestart && err != nil {
+					containerRespC <- containerWaitResponse{
-						panic(fmt.Errorf("shouldRestart must return nil error if restarting, but returned: %w", err))
+						WaitResponse: resp,
-					}
+						restarting:   false,
-					if !shouldRestart {
+						restartCount: restartCount,
-						a.logger.Info("Container exited", "id", shortID(containerID), "should_restart", "false", "exit_code", resp.StatusCode, "restart_count", restartCount)
+						err:          err,
-						containerRespC <- containerWaitResponse{
-							WaitResponse: resp,
-							restarting:   false,
-							restartCount: restartCount,
-							err:          err,
-						}
-						return
 					}
 				}
 
+				if shouldRestartFunc == nil {
+					exit(nil)
+					return
+				}
+
+				shouldRestart, err := shouldRestartFunc(resp.StatusCode, restartCount, time.Since(startedWaitingAt))
+				if shouldRestart && err != nil {
+					panic(fmt.Errorf("shouldRestart must return nil error if restarting, but returned: %w", err))
+				}
+				if !shouldRestart {
+					exit(err)
+					return
+				}
+
 				a.logger.Info("Container exited", "id", shortID(containerID), "should_restart", "true", "exit_code", resp.StatusCode, "restart_count", restartCount)
 				timer.Reset(restartInterval)
 

internal/container/container_test.go

@@ -94,9 +94,6 @@ func TestClientRunContainer(t *testing.T) {
 				Mode:    0755,
 			},
 		},
-		ShouldRestart: func(int64, int, time.Duration) (bool, error) {
-			return false, nil
-		},
 	})
 
 	done := make(chan struct{})