fix timeouts: set timeout on dequeue + handle sub millisecond timouts + handle negative timeouts

Author: maxdml

dbos/client_test.go

@@ -473,10 +473,6 @@ func TestCancelResume(t *testing.T) {
 		// Verify the deadline was reset (should be different from original)
 		assert.False(t, resumeStatus.Deadline.Equal(originalDeadline), "expected deadline to be reset after resume, but it remained the same: %v", originalDeadline)
 
-		// The new deadline should be after resumeStart + workflowTimeout
-		expectedDeadline := resumeStart.Add(workflowTimeout - 100*time.Millisecond) // Allow some leeway for processing time
-		assert.True(t, resumeStatus.Deadline.After(expectedDeadline), "deadline %v is too early (expected around %v)", resumeStatus.Deadline, expectedDeadline)
-
 		// Wait for the workflow to complete
 		_, err = resumeHandle.GetResult()
 		require.Error(t, err, "expected timeout error, but got none")
@@ -491,6 +487,10 @@ func TestCancelResume(t *testing.T) {
 		finalStatus, err := resumeHandle.GetStatus()
 		require.NoError(t, err, "failed to get final workflow status")
 
+		// The new deadline should have been set after resumeStart + workflowTimeout
+		expectedDeadline := resumeStart.Add(workflowTimeout - 100*time.Millisecond) // Allow some leeway for processing time
+		assert.True(t, finalStatus.Deadline.After(expectedDeadline), "deadline %v is too early (expected around %v)", resumeStatus.Deadline, expectedDeadline)
+
 		assert.Equal(t, WorkflowStatusCancelled, finalStatus.Status, "expected final workflow status to be CANCELLED")
 
 		require.True(t, queueEntriesAreCleanedUp(serverCtx), "expected queue entries to be cleaned up after cancel/resume timeout test")

dbos/queues_test.go

@@ -850,6 +850,18 @@ func TestQueueTimeouts(t *testing.T) {
 	RegisterWorkflow(dbosCtx, detachedWorkflow)
 	RegisterWorkflow(dbosCtx, enqueuedWorkflowEnqueuesADetachedWorkflow)
 
+	timeoutOnDequeueQueue := NewWorkflowQueue(dbosCtx, "timeout-on-dequeue-queue", WithGlobalConcurrency(1))
+	blockingEvent := NewEvent()
+	blockingWorkflow := func(ctx DBOSContext, _ string) (string, error) {
+		blockingEvent.Wait()
+		return "blocking-done", nil
+	}
+	RegisterWorkflow(dbosCtx, blockingWorkflow)
+	fastWorkflow := func(ctx DBOSContext, _ string) (string, error) {
+		return "done", nil
+	}
+	RegisterWorkflow(dbosCtx, fastWorkflow)
+
 	dbosCtx.Launch()
 
 	t.Run("EnqueueWorkflowTimeout", func(t *testing.T) {
@@ -936,6 +948,46 @@ func TestQueueTimeouts(t *testing.T) {
 
 		require.True(t, queueEntriesAreCleanedUp(dbosCtx), "expected queue entries to be cleaned up after workflow cancellation, but they are not")
 	})
+
+	t.Run("TimeoutOnlySetOnDequeue", func(t *testing.T) {
+		// Test that deadline is only set when workflow is dequeued, not when enqueued
+
+		// Enqueue blocking workflow first
+		blockingHandle, err := RunAsWorkflow(dbosCtx, blockingWorkflow, "blocking", WithQueue(timeoutOnDequeueQueue.Name))
+		require.NoError(t, err, "failed to enqueue blocking workflow")
+
+		// Set a timeout that would expire if set on enqueue
+		timeout := 2 * time.Second
+		timeoutCtx, cancelFunc := WithTimeout(dbosCtx, timeout)
+		defer cancelFunc()
+
+		// Enqueue second workflow with timeout
+		handle, err := RunAsWorkflow(timeoutCtx, fastWorkflow, "timeout-test", WithQueue(timeoutOnDequeueQueue.Name))
+		require.NoError(t, err, "failed to enqueue timeout workflow")
+
+		// Sleep for duration exceeding the timeout
+		time.Sleep(timeout * 2)
+
+		// Signal the blocking workflow to complete
+		blockingEvent.Set()
+
+		// Wait for blocking workflow to complete
+		blockingResult, err := blockingHandle.GetResult()
+		require.NoError(t, err, "failed to get result from blocking workflow")
+		assert.Equal(t, "blocking-done", blockingResult, "expected blocking workflow result")
+
+		// Now the second workflow should dequeue and complete successfully (timeout should be much longer than execution time)
+		// Note: this might be flaky if we hit the dequeue interval 1sec boundary...
+		_, err = handle.GetResult()
+		require.NoError(t, err, "unexpected error from workflow")
+
+		// Check the workflow status: should be success
+		finalStatus, err := handle.GetStatus()
+		require.NoError(t, err, "failed to get final status of timeout workflow")
+		assert.Equal(t, WorkflowStatusSuccess, finalStatus.Status, "expected timeout workflow status to be WorkflowStatusSuccess")
+
+		require.True(t, queueEntriesAreCleanedUp(dbosCtx), "expected queue entries to be cleaned up after test")
+	})
 }
 
 func TestPriorityQueue(t *testing.T) {

dbos/system_database.go

@@ -307,7 +307,7 @@ func (s *sysDB) insertWorkflowStatus(ctx context.Context, input insertWorkflowSt
 
 	var timeoutMs *int64 = nil
 	if input.status.Timeout > 0 {
-		millis := input.status.Timeout.Milliseconds()
+		millis := input.status.Timeout.Round(time.Millisecond).Milliseconds()
 		timeoutMs = &millis
 	}
 
@@ -787,25 +787,17 @@ func (s *sysDB) resumeWorkflow(ctx context.Context, workflowID string) error {
 		return nil // Workflow is complete, do nothing
 	}
 
-	// If the original workflow has a timeout, let's recompute a deadline
-	var deadline *int64 = nil
-	if wf.Timeout > 0 {
-		deadlineMs := time.Now().Add(wf.Timeout).UnixMilli()
-		deadline = &deadlineMs
-	}
-
 	// Set the workflow's status to ENQUEUED and clear its recovery attempts, set new deadline
 	updateStatusQuery := `UPDATE dbos.workflow_status
 						  SET status = $1, queue_name = $2, recovery_attempts = $3, 
-						      workflow_deadline_epoch_ms = $4, deduplication_id = NULL,
-						      started_at_epoch_ms = NULL, updated_at = $5
-						  WHERE workflow_uuid = $6`
+						      workflow_deadline_epoch_ms = NULL, deduplication_id = NULL,
+						      started_at_epoch_ms = NULL, updated_at = $4
+						  WHERE workflow_uuid = $5`
 
 	_, err = tx.Exec(ctx, updateStatusQuery,
 		WorkflowStatusEnqueued,
 		_DBOS_INTERNAL_QUEUE_NAME,
 		0,
-		deadline,
 		time.Now().UnixMilli(),
 		workflowID)
 	if err != nil {

dbos/workflow.go

@@ -643,21 +643,22 @@ func (c *dbosContext) RunAsWorkflow(_ DBOSContext, fn WorkflowFunc, input any, o
 		status = WorkflowStatusPending
 	}
 
-	// Check if the user-provided context has a deadline
+	// Compute the timeout based on the context deadline, if any
 	deadline, ok := c.Deadline()
 	if !ok {
 		deadline = time.Time{} // No deadline set
 	}
-
-	// Compute the timeout based on the deadline
 	var timeout time.Duration
 	if !deadline.IsZero() {
 		timeout = time.Until(deadline)
-		/* unclear to me if this is a real use case:
+		// The timeout could be in the past, for small deadlines, to propagation delays. If so set it to a minimal value
 		if timeout < 0 {
-			return nil, newWorkflowExecutionError(workflowID, "deadline is in the past")
+			timeout = 1 * time.Millisecond
 		}
-		*/
+	}
+	// When enqueuing, we do not set a deadline. It'll be computed with the timeout during dequeue.
+	if status == WorkflowStatusEnqueued {
+		deadline = time.Time{} // No deadline for enqueued workflows
 	}
 
 	if params.priority > uint(math.MaxInt) {
@@ -737,11 +738,19 @@ func (c *dbosContext) RunAsWorkflow(_ DBOSContext, fn WorkflowFunc, input any, o
 
 	workflowCtx := WithValue(c, workflowStateKey, wfState)
 
-	// If the workflow has a durable deadline, set it in the context.
+	// If the workflow has a timeout but no deadline, compute the deadline from the timeout.
+	// Else use the durable deadline.
+	durableDeadline := time.Time{}
+	if insertStatusResult.timeout > 0 && insertStatusResult.workflowDeadline.IsZero() {
+		durableDeadline = time.Now().Add(insertStatusResult.timeout)
+	} else if !insertStatusResult.workflowDeadline.IsZero() {
+		durableDeadline = insertStatusResult.workflowDeadline
+	}
+
 	var stopFunc func() bool
 	cancelFuncCompleted := make(chan struct{})
-	if !insertStatusResult.workflowDeadline.IsZero() {
-		workflowCtx, _ = WithTimeout(workflowCtx, time.Until(insertStatusResult.workflowDeadline))
+	if !durableDeadline.IsZero() {
+		workflowCtx, _ = WithTimeout(workflowCtx, time.Until(durableDeadline))
 		// Register a cancel function that cancels the workflow in the DB as soon as the context is cancelled
 		dbosCancelFunction := func() {
 			c.logger.Info("Cancelling workflow", "workflow_id", workflowID)