fix(engine): handle durable wait_for conditions without base (#5260)

Author: igor-kupczynski

internal/services/dispatcher/dispatcher.go

@@ -346,7 +346,7 @@ func New(fs ...DispatcherOpt) (*DispatcherImpl, error) {
 		streamEventBufferTimeout:            opts.streamEventBufferTimeout,
 		version:                             opts.version,
 		refreshTimeoutBuf:                   newRefreshTimeoutBuffer(),
-		serviceV1:                           newDispatcherService(opts.repov1, opts.mqv1, opts.pubsub, v, opts.l, opts.dispatcherId, opts.analytics, opts.promGate),
+		serviceV1:                           newDispatcherService(opts.repov1, opts.mqv1, opts.pubsub, v, opts.l, opts.dispatcherId, opts.analytics, a, opts.promGate),
 	}, nil
 }
 

internal/services/dispatcher/dispatcher_v1.go

@@ -10,6 +10,7 @@ import (
 	"github.com/hatchet-dev/hatchet/internal/services/shared/streams"
 	"github.com/hatchet-dev/hatchet/internal/syncx"
 	"github.com/hatchet-dev/hatchet/pkg/analytics"
+	hatcheterrors "github.com/hatchet-dev/hatchet/pkg/errors"
 	"github.com/hatchet-dev/hatchet/pkg/integrations/metrics/prometheus"
 	v1 "github.com/hatchet-dev/hatchet/pkg/repository"
 	"github.com/hatchet-dev/hatchet/pkg/validator"
@@ -28,6 +29,7 @@ type DispatcherServiceImpl struct {
 	mq                 msgqueue.MessageQueue
 	v                  validator.Validator
 	analytics          analytics.Analytics
+	alerter            *hatcheterrors.Wrapped
 	triggerWriter      *trigger.TriggerWriter
 	pubBuffer          *msgqueue.MQPubBuffer
 	streamSessions     *streams.Registry
@@ -44,7 +46,7 @@ func (d *DispatcherServiceImpl) CancelStreamSessions() {
 	d.streamSessions.CancelAll()
 }
 
-func newDispatcherService(repo v1.Repository, mq msgqueue.MessageQueue, pubsub msgqueue.PubSub, v validator.Validator, l *zerolog.Logger, dispatcherId uuid.UUID, a analytics.Analytics, promGate *prometheus.Gate) *DispatcherServiceImpl {
+func newDispatcherService(repo v1.Repository, mq msgqueue.MessageQueue, pubsub msgqueue.PubSub, v validator.Validator, l *zerolog.Logger, dispatcherId uuid.UUID, a analytics.Analytics, alerter *hatcheterrors.Wrapped, promGate *prometheus.Gate) *DispatcherServiceImpl {
 	pubBuffer := msgqueue.NewMQPubBuffer(mq)
 	tw := trigger.NewTriggerWriter(mq, pubsub, repo, l, pubBuffer, 0, promGate)
 
@@ -57,6 +59,7 @@ func newDispatcherService(repo v1.Repository, mq msgqueue.MessageQueue, pubsub m
 		pubBuffer:     pubBuffer,
 		dispatcherId:  dispatcherId,
 		analytics:     a,
+		alerter:       alerter,
 
 		streamSessions: streams.NewRegistry(),
 	}

internal/services/dispatcher/server_v1.go

@@ -19,6 +19,7 @@ import (
 
 	"github.com/hatchet-dev/hatchet/internal/msgqueue"
 	contracts "github.com/hatchet-dev/hatchet/internal/services/shared/proto/v1"
+	"github.com/hatchet-dev/hatchet/internal/services/shared/recoveryutils"
 	"github.com/hatchet-dev/hatchet/internal/services/shared/rpcstream"
 	tasktypes "github.com/hatchet-dev/hatchet/internal/services/shared/tasktypes/v1"
 	"github.com/hatchet-dev/hatchet/pkg/analytics"
@@ -885,7 +886,15 @@ func (d *DispatcherServiceImpl) handleDurableTaskRequest(
 	ctx context.Context,
 	invocation *durableTaskInvocation,
 	req *contracts.DurableTaskRequest,
-) error {
+) (err error) {
+	// recover per request, not per goroutine, so a panic doesn't stop the worker-status runner
+	defer func() {
+		if r := recover(); r != nil {
+			_ = recoveryutils.RecoverWithAlert(d.l, d.alerter, r)
+			err = fmt.Errorf("recovered from panic handling %T: %v", req.GetMessage(), r)
+		}
+	}()
+
 	ctx, span := telemetry.NewRootSpan(ctx, "dispatcher.handle-durable-task-request")
 	defer span.End()
 
@@ -1270,7 +1279,7 @@ func (d *DispatcherServiceImpl) handleWaitFor(
 
 	if req.WaitForConditions != nil {
 		for _, condition := range req.WaitForConditions.SleepConditions {
-			orGroupId, parseErr := uuid.Parse(condition.Base.OrGroupId)
+			orGroupId, parseErr := uuid.Parse(condition.GetBase().GetOrGroupId())
 			if parseErr != nil {
 				return connect.NewError(connect.CodeInvalidArgument, fmt.Errorf("or group id is not a valid uuid: %v", parseErr))
 			}
@@ -1283,7 +1292,7 @@ func (d *DispatcherServiceImpl) handleWaitFor(
 		}
 
 		for _, condition := range req.WaitForConditions.UserEventConditions {
-			orGroupId, parseErr := uuid.Parse(condition.Base.OrGroupId)
+			orGroupId, parseErr := uuid.Parse(condition.GetBase().GetOrGroupId())
 			if parseErr != nil {
 				return connect.NewError(connect.CodeInvalidArgument, fmt.Errorf("or group id is not a valid uuid: %v", parseErr))
 			}

internal/services/dispatcher/server_v1_durable_task_test.go

@@ -7,13 +7,16 @@ import (
 	"testing"
 	"time"
 
+	"connectrpc.com/connect"
 	"github.com/google/uuid"
 	"github.com/rs/zerolog"
+	"github.com/stretchr/testify/require"
 
 	contracts "github.com/hatchet-dev/hatchet/internal/services/shared/proto/v1"
 	"github.com/hatchet-dev/hatchet/internal/services/shared/rpcstream"
 	"github.com/hatchet-dev/hatchet/internal/services/shared/streams"
 	"github.com/hatchet-dev/hatchet/pkg/analytics"
+	v1 "github.com/hatchet-dev/hatchet/pkg/repository"
 	"github.com/hatchet-dev/hatchet/pkg/repository/sqlcv1"
 )
 
@@ -201,3 +204,81 @@ func TestDurableTaskReturnsNilOnEOF(t *testing.T) {
 		t.Fatal("DurableTask did not return after EOF")
 	}
 }
+
+type fakeFoundTask struct {
+	v1.Repository
+	v1.TaskRepository
+}
+
+func (f fakeFoundTask) Tasks() v1.TaskRepository { return f }
+func (fakeFoundTask) GetTaskByExternalId(context.Context, uuid.UUID, uuid.UUID, bool) (*sqlcv1.FlattenExternalIdsRow, error) {
+	return &sqlcv1.FlattenExternalIdsRow{}, nil
+}
+
+func TestHandleWaitForNilConditionBaseReturnsInvalidArgument(t *testing.T) {
+	d := newTestDispatcher()
+	d.repo = fakeFoundTask{}
+
+	for _, c := range []*contracts.DurableEventListenerConditions{
+		{SleepConditions: []*contracts.SleepMatchCondition{{}}},
+		{UserEventConditions: []*contracts.UserEventMatchCondition{{}}},
+	} {
+		req := &contracts.DurableTaskWaitForRequest{DurableTaskExternalId: uuid.NewString(), WaitForConditions: c}
+		err := d.handleWaitFor(context.Background(), &durableTaskInvocation{}, req)
+		require.Equal(t, connect.CodeInvalidArgument, connect.CodeOf(err), err)
+	}
+}
+
+type panickingRepo struct {
+	v1.Repository
+	calls chan struct{}
+}
+
+func (r panickingRepo) Tasks() v1.TaskRepository {
+	r.calls <- struct{}{}
+	panic("boom")
+}
+
+func TestDurableTaskRecoversFromHandlerPanics(t *testing.T) {
+	d := newTestDispatcher()
+	calls := make(chan struct{}, 3)
+	d.repo = panickingRepo{calls: calls}
+
+	taskId := uuid.NewString()
+	memo := &contracts.DurableTaskRequest{Message: &contracts.DurableTaskRequest_Memo{
+		Memo: &contracts.DurableTaskMemoRequest{DurableTaskExternalId: taskId},
+	}}
+	status := &contracts.DurableTaskRequest{Message: &contracts.DurableTaskRequest_WorkerStatus{
+		WorkerStatus: &contracts.DurableTaskWorkerStatusRequest{
+			WaitingEntries: []*contracts.DurableTaskAwaitedCompletedEntry{{DurableTaskExternalId: taskId}},
+		},
+	}}
+
+	reqs := []*contracts.DurableTaskRequest{memo, status, status}
+	server := &fakeDurableTaskServer{
+		ctx: tenantContext(),
+		recv: func() (*contracts.DurableTaskRequest, error) {
+			// wait for the previous request's handler to panic, so the second status needs the
+			// worker-status runner to have survived the first
+			if len(reqs) < 3 {
+				<-calls
+			}
+
+			if len(reqs) == 0 {
+				return nil, io.EOF
+			}
+
+			req := reqs[0]
+			reqs = reqs[1:]
+
+			return req, nil
+		},
+	}
+
+	select {
+	case err := <-runDurableTask(d, server):
+		require.NoError(t, err)
+	case <-time.After(5 * time.Second):
+		t.Fatal("a request after a panicking one was never handled")
+	}
+}