Handle client closed request situations

Author: AlexPykavy

tests/helper/stdout.go

@@ -1,6 +1,7 @@
 package helper
 
 import (
+	"fmt"
 	"io"
 	"os"
 )
@@ -17,6 +18,8 @@ func CaptureStdout(f func()) string {
 
 	f()
 
+	fmt.Println() // to flush data
+
 	w.Close()
 	b, _ := io.ReadAll(r)
 

tests/worker/start_task_test.go

@@ -1,6 +1,7 @@
 package worker
 
 import (
+	"context"
 	"dirigeant/task"
 	"dirigeant/tests/helper"
 	"dirigeant/worker"
@@ -9,6 +10,7 @@ import (
 	"net/http/httptest"
 	"sync"
 	"testing"
+	"time"
 
 	"github.com/google/uuid"
 	"github.com/stretchr/testify/assert"
@@ -105,3 +107,41 @@ func TestStartTask__AllButOneRequestsShouldFailIfCreatingTheSameTaskSimultaneous
 	assert.Equal(t, 1, len(api.Worker.Tasks), "Tasks map should contain 1 task")
 	assert.NotNil(t, api.Worker.Tasks[testTask.ID], "Persisted task ID should match the one from request")
 }
+
+func TestStartTask__ShouldHandleClientClosedRequest(t *testing.T) {
+	api := &worker.Api{
+		Worker: &worker.Worker{
+			Tasks: make(map[uuid.UUID]*task.Task),
+		},
+	}
+	testTask := helper.PingTask("ping-task", "127.0.0.1")
+	ctx, cancel := context.WithCancel(context.TODO())
+
+	// 1 - Create a task
+	var wg sync.WaitGroup
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+
+		createRequest := helper.NewTaskPostRequest(testTask).WithContext(ctx)
+		createResponseRecorder := httptest.NewRecorder()
+
+		stdout := helper.CaptureStdout(func() {
+			api.HandleCreateTask(createResponseRecorder, createRequest)
+		})
+
+		assert.Equal(t, 499, createResponseRecorder.Code, "Response status code should be 499 Client Closed Request")
+		assert.Equal(t, "Error when executing the task: client closed request", createResponseRecorder.Body.String(), "Response body should contain error message")
+		assert.NotEmpty(t, stdout, "Task logs shouldn't be empty")
+		assert.Empty(t, api.Worker.Tasks, "Tasks map should be empty")
+	}()
+
+	time.Sleep(1 * time.Second)
+
+	assert.Equal(t, 1, len(api.Worker.Tasks), "Tasks map should contain 1 task")
+
+	// 2 - Cancel a request
+	cancel()
+
+	wg.Wait()
+}

worker/api.go

@@ -69,18 +69,34 @@ func (a *Api) HandleCreateTask(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	if err := a.Worker.StartTask(t); err != nil {
-		if errors.Is(err, task.ErrAlreadyExists) {
-			w.WriteHeader(http.StatusConflict)
-		} else {
-			w.WriteHeader(http.StatusInternalServerError)
+	errCh := make(chan error)
+	defer close(errCh)
+	go func() {
+		errCh <- a.Worker.StartTask(t)
+	}()
+
+	select {
+	case <-r.Context().Done():
+		a.Worker.StopTask(t.ID)
+
+		<-errCh
+
+		w.WriteHeader(499) // client closed request
+		fmt.Fprint(w, "Error when executing the task: client closed request")
+	case err := <-errCh:
+		if err != nil {
+			if errors.Is(err, task.ErrAlreadyExists) {
+				w.WriteHeader(http.StatusConflict)
+			} else {
+				w.WriteHeader(http.StatusInternalServerError)
+			}
+
+			fmt.Fprintf(w, "Error when executing the task: %v", err)
+			return
 		}
 
-		fmt.Fprintf(w, "Error when executing the task: %v", err)
-		return
+		w.WriteHeader(http.StatusCreated)
 	}
-
-	w.WriteHeader(http.StatusCreated)
 }
 
 func (a *Api) HandleDeleteTask(w http.ResponseWriter, r *http.Request) {