[WIP] Fix code for all comments in review thread

task/task.go

@@ -142,13 +142,23 @@ func ScheduleTask[T any, R any](
 	return
 }
 
-func ScheduleTaskInTx[T any, R any](
-	tx server.PgTx,
+type preparedTask struct {
+	taskId         server.Id
+	functionName   string
+	argsJson       []byte
+	byJwtJson      *string
+	runAt          time.Time
+	runOnceKey     *string
+	priority       TaskPriority
+	maxTimeSeconds int
+}
+
+func prepareTask[T any, R any](
 	taskFunction TaskFunction[T, R],
 	args T,
 	clientSession *session.ClientSession,
 	opts ...any,
-) (taskId server.Id) {
+) preparedTask {
 	taskTarget := NewTaskTarget(taskFunction)
 
 	argsJson, err := json.Marshal(args)
@@ -203,11 +213,29 @@ func ScheduleTaskInTx[T any, R any](
 		runOnceKey_ := runOnce.String()
 		runOnceKey = &runOnceKey_
 	}
-	maxTimeSeconds := int(runMaxTime.MaxTime / time.Second)
 
-	claimTime := time.Time{}
+	return preparedTask{
+		taskId:         server.NewId(),
+		functionName:   taskTarget.TargetFunctionName(),
+		argsJson:       argsJson,
+		byJwtJson:      byJwtJson,
+		runAt:          runAt.At.UTC(),
+		runOnceKey:     runOnceKey,
+		priority:       runPriority.Priority,
+		maxTimeSeconds: int(runMaxTime.MaxTime / time.Second),
+	}
+}
 
-	taskId = server.NewId()
+func ScheduleTaskInTx[T any, R any](
+	tx server.PgTx,
+	taskFunction TaskFunction[T, R],
+	args T,
+	clientSession *session.ClientSession,
+	opts ...any,
+) (taskId server.Id) {
+	p := prepareTask(taskFunction, args, clientSession, opts...)
+
+	claimTime := time.Time{}
 
 	server.RaisePgResult(tx.Exec(
 		clientSession.Ctx,
@@ -230,17 +258,81 @@ func ScheduleTaskInTx[T any, R any](
 				run_priority = LEAST(pending_task.run_priority, $8),
 				run_max_time_seconds = GREATEST(pending_task.run_max_time_seconds, $9)
 		`,
-		taskId,
-		taskTarget.TargetFunctionName(),
-		argsJson,
+		p.taskId,
+		p.functionName,
+		p.argsJson,
 		clientSession.ClientAddress,
-		byJwtJson,
-		runAt.At.UTC(),
-		runOnceKey,
-		runPriority.Priority,
-		maxTimeSeconds,
+		p.byJwtJson,
+		p.runAt,
+		p.runOnceKey,
+		p.priority,
+		p.maxTimeSeconds,
 		claimTime,
 	))
+	return p.taskId
+}
+
+func ScheduleTaskInTxIfAbsent[T any, R any](
+	tx server.PgTx,
+	taskFunction TaskFunction[T, R],
+	args T,
+	clientSession *session.ClientSession,
+	runOnce *RunOnceOption,
+	opts ...any,
+) (scheduled bool, taskId server.Id) {
+	if runOnce == nil {
+		panic("ScheduleTaskInTxIfAbsent requires a non-nil runOnce key")
+	}
+
+	p := prepareTask(taskFunction, args, clientSession, append(opts, runOnce)...)
+	claimTime := time.Time{}
+
+	tag := server.RaisePgResult(tx.Exec(
+		clientSession.Ctx,
+		`
+			INSERT INTO pending_task (
+				task_id,
+		        function_name,
+		        args_json,
+		        client_address,
+		        client_by_jwt_json,
+		        run_at,
+		        run_once_key,
+		        run_priority,
+		        run_max_time_seconds,
+		        claim_time,
+		        release_time
+			) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $10)
+			ON CONFLICT (run_once_key) DO NOTHING
+		`,
+		p.taskId,
+		p.functionName,
+		p.argsJson,
+		clientSession.ClientAddress,
+		p.byJwtJson,
+		p.runAt,
+		p.runOnceKey,
+		p.priority,
+		p.maxTimeSeconds,
+		claimTime,
+	))
+	scheduled = 0 < tag.RowsAffected()
+	if scheduled {
+		taskId = p.taskId
+	}
+	return
+}
+
+func ScheduleTaskIfAbsent[T any, R any](
+	taskFunction TaskFunction[T, R],
+	args T,
+	clientSession *session.ClientSession,
+	runOnce *RunOnceOption,
+	opts ...any,
+) (scheduled bool, taskId server.Id) {
+	server.Tx(clientSession.Ctx, func(tx server.PgTx) {
+		scheduled, taskId = ScheduleTaskInTxIfAbsent[T, R](tx, taskFunction, args, clientSession, runOnce, opts...)
+	})
 	return
 }
 

task/task_test.go

@@ -2,6 +2,7 @@ package task
 
 import (
 	"context"
+	"encoding/json"
 	"errors"
 	"fmt"
 	"math"
@@ -211,6 +212,73 @@ func TestTask(t *testing.T) {
 	})
 }
 
+type Work2Args struct {
+	Tag string
+}
+
+type Work2Result struct {
+}
+
+func Work2(
+	work2 *Work2Args,
+	clientSession *session.ClientSession,
+) (*Work2Result, error) {
+	return &Work2Result{}, nil
+}
+
+func TestScheduleTaskIfAbsent(t *testing.T) {
+	server.DefaultTestEnv().Run(t, func(t testing.TB) {
+		ctx := context.Background()
+		clientSession := session.Testing_CreateClientSession(ctx, nil)
+		defer clientSession.Cancel()
+
+		key := RunOnce("test_schedule_task_if_absent", server.NewId())
+
+		scheduled, firstTaskId := ScheduleTaskIfAbsent(
+			Work2,
+			&Work2Args{Tag: "first"},
+			clientSession,
+			key,
+		)
+		connect.AssertEqual(t, scheduled, true)
+
+		scheduledAgain, duplicateTaskId := ScheduleTaskIfAbsent(
+			Work2,
+			&Work2Args{Tag: "second"},
+			clientSession,
+			key,
+		)
+		connect.AssertEqual(t, scheduledAgain, false)
+		connect.AssertEqual(t, duplicateTaskId, server.Id{})
+
+		tasks := GetTasks(ctx, firstTaskId)
+		task, ok := tasks[firstTaskId]
+		if !ok {
+			t.Fatal("first task not found")
+		}
+		var args Work2Args
+		if err := json.Unmarshal([]byte(task.ArgsJson), &args); err != nil {
+			t.Fatal(err)
+		}
+		connect.AssertEqual(t, args.Tag, "first")
+
+		defer func() {
+			if r := recover(); r == nil {
+				t.Fatal("expected panic for nil runOnce key")
+			}
+		}()
+		server.Tx(ctx, func(tx server.PgTx) {
+			ScheduleTaskInTxIfAbsent(
+				tx,
+				Work2,
+				&Work2Args{Tag: "panic"},
+				clientSession,
+				nil,
+			)
+		})
+	})
+}
+
 type AlwaysFailArgs struct {
 }