orbit/sdk/go/worker_cancel_integration_test.go

package orbit

import (
	"context"
	"encoding/json"
	"errors"
	"sync/atomic"
	"testing"
	"time"
)

func cancelRuntimeTask(t *testing.T, ctx context.Context, c *Client, id string) CancelTaskResult {
	t.Helper()
	for {
		current, err := c.GetTask(ctx, id)
		if err != nil {
			t.Fatal(err)
		}
		out, err := c.CancelTask(ctx, id, "cancel-"+workerTestID(t), CancelTaskRequest{ExpectedRevision: current.Revision, Reason: "controlled cancellation regression"})
		if err == nil {
			return out
		}
		var api *APIError
		if !errors.As(err, &api) || api.Code != "STATE_CONFLICT" {
			t.Fatal(err)
		}
		if ctx.Err() != nil {
			t.Fatal(ctx.Err())
		}
	}
}

func TestIntegrationWorkerRuntimeCancellation(t *testing.T) {
	api, queue, token := integrationWorkerClient(t)
	for _, mode := range []string{"confirmed", "unknown", "success", "lost_notice", "before_handler", "uncooperative"} {
		t.Run(mode, func(t *testing.T) {
			ctx, cancel := context.WithTimeout(context.Background(), 18*time.Second)
			defer cancel()
			g := newWorkerFaultGateway(t, api.baseURL)
			client, err := NewClient(g.server.URL, Options{Token: token, HTTPClient: g.server.Client()})
			if err != nil {
				t.Fatal(err)
			}
			name := "cancel." + mode + "." + workerTestID(t)
			entered := make(chan struct{})
			observed := make(chan error, 1)
			release := make(chan struct{})
			var calls atomic.Int32
			w, err := NewWorker(client, WorkerOptions{QueueID: queue, Capacity: 1, ShutdownTimeout: 2 * time.Second, Handlers: []WorkerHandler{runtimeHandler(name, func(work context.Context, e WorkerExecution) (json.RawMessage, error) {
				calls.Add(1)
				close(entered)
				if mode == "uncooperative" {
					<-release
					return json.RawMessage(`{"late":true}`), nil
				}
				<-work.Done()
				observed <- context.Cause(work)
				switch mode {
				case "unknown":
					return nil, work.Err()
				case "success":
					return json.RawMessage(`{"sequence":9007199254740993}`), nil
				default:
					return nil, &WorkerCancellationConfirmed{Details: json.RawMessage(`{"evidence_ref":"objects/stopped","sequence":9007199254740993}`)}
				}
			})}})
			if err != nil {
				t.Fatal(err)
			}
			timeout := uint64(12000)
			if mode == "uncooperative" {
				timeout = 2000
			}
			task := runtimeTask(t, ctx, api, queue, name, timeout)
			if mode == "before_handler" {
				g.dropOffline.Store(true)
				g.drop.Store("task.start")
			}
			run, stop := context.WithCancel(ctx)
			done := make(chan error, 1)
			go func() { done <- w.Run(run) }()
			defer stop()
			defer close(release)
			if mode == "before_handler" {
				g.expectDrop(t, "task.start")
			} else {
				select {
				case <-entered:
				case err := <-done:
					t.Fatalf("Worker stopped: %v", err)
				case <-ctx.Done():
					t.Fatal(ctx.Err())
				}
			}
			if mode == "lost_notice" {
				g.dropCancel.Store(true)
			}
			if mode == "confirmed" {
				g.drop.Store("task.finish")
			}
			out := cancelRuntimeTask(t, ctx, api, task.ID)
			if out.Task.State != TaskRunning || out.Task.CancelRequestedAt == nil || out.Task.CompletedAt != nil {
				t.Fatal("cancel intent reported completion")
			}
			if mode == "before_handler" {
				g.offline.Store(false)
			}
			if mode == "lost_notice" {
				g.expectDrop(t, "task.cancel")
			}
			if mode != "uncooperative" && mode != "before_handler" {
				select {
				case cause := <-observed:
					if !errors.Is(cause, ErrTaskCancellationRequested) {
						t.Fatalf("wrong handler cancellation cause: %v", cause)
					}
				case <-ctx.Done():
					t.Fatal("handler missed cancellation")
				}
			}
			want := TaskCancelled
			if mode == "unknown" || mode == "uncooperative" {
				want = TaskNeedsReview
			}
			if mode == "success" {
				want = TaskSucceeded
			}
			finished := waitRuntimeState(t, ctx, api, task.ID, want)
			if finished.AttemptCount != 1 {
				t.Fatal("cancel request created retry")
			}
			if mode == "success" && string(finished.Result) != `{"sequence":9007199254740993}` {
				t.Fatal("cancel overwrote success")
			}
			if mode == "before_handler" && calls.Load() != 0 {
				t.Fatal("handler started after cancellation observed")
			}
			if mode == "confirmed" {
				g.expectDrop(t, "task.finish")
			}
			stop()
			select {
			case err := <-done:
				if mode == "uncooperative" {
					if !errors.Is(err, ErrWorkerShutdownTimeout) {
						t.Fatalf("uncooperative handler falsely drained: %v", err)
					}
				} else if err != nil {
					t.Fatal(err)
				}
			case <-ctx.Done():
				t.Fatal("Worker did not drain")
			}
		})
	}
}