fix single-instance orchestration purge requests

client/client_grpc.go

@@ -271,7 +271,8 @@ func (c *TaskHubGrpcClient) ResumeOrchestration(ctx context.Context, id api.Inst
 // [api.api.ErrInstanceNotFound] is returned if the specified orchestration instance doesn't exist.
 func (c *TaskHubGrpcClient) PurgeOrchestrationState(ctx context.Context, id api.InstanceID, opts ...api.PurgeOptions) error {
 	req := &protos.PurgeInstancesRequest{
-		Request: &protos.PurgeInstancesRequest_InstanceId{InstanceId: string(id)},
+		Request:         &protos.PurgeInstancesRequest_InstanceId{InstanceId: string(id)},
+		IsOrchestration: true,
 	}
 	for _, configure := range opts {
 		if err := configure(req); err != nil {

client/client_grpc_management_test.go

@@ -2,6 +2,7 @@ package client
 
 import (
 	"context"
+	"errors"
 	"testing"
 	"time"
 
@@ -21,6 +22,7 @@ type managementServer struct {
 
 	queryErr error
 	rewind   func(context.Context, *protos.RewindInstanceRequest) (*protos.RewindInstanceResponse, error)
+	purge    func(context.Context, *protos.PurgeInstancesRequest) (*protos.PurgeInstancesResponse, error)
 }
 
 func (s *managementServer) RewindInstance(ctx context.Context, req *protos.RewindInstanceRequest) (*protos.RewindInstanceResponse, error) {
@@ -30,6 +32,13 @@ func (s *managementServer) RewindInstance(ctx context.Context, req *protos.Rewin
 	return s.rewind(ctx, req)
 }
 
+func (s *managementServer) PurgeInstances(ctx context.Context, req *protos.PurgeInstancesRequest) (*protos.PurgeInstancesResponse, error) {
+	if s.purge == nil {
+		return nil, status.Error(codes.Unimplemented, "purge is not implemented")
+	}
+	return s.purge(ctx, req)
+}
+
 func (s *managementServer) QueryInstances(
 	_ context.Context,
 	req *protos.QueryInstancesRequest,
@@ -91,14 +100,76 @@ func TestTaskHubGrpcManagementErrorsRoundTrip(t *testing.T) {
 		},
 	} {
 		t.Run(test.name, func(t *testing.T) {
-			client := startQueryClient(t, &managementServer{queryErr: test.err})
+			client := startQueryClient(t, &managementServer{
+				queryErr: test.err,
+				purge: func(context.Context, *protos.PurgeInstancesRequest) (*protos.PurgeInstancesResponse, error) {
+					return nil, test.err
+				},
+			})
 
 			ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
 			defer cancel()
 
 			_, err := client.QueryInstances(ctx, api.OrchestrationQuery{PageSize: 10})
 			require.ErrorIs(t, err, test.want)
 			require.Equal(t, status.Code(test.err), status.Code(err))
+
+			err = client.PurgeOrchestrationState(ctx, "instance")
+			require.ErrorIs(t, err, test.want)
+			require.Equal(t, status.Code(test.err), status.Code(err))
 		})
 	}
 }
+
+func TestPurgeOrchestrationStateRequest(t *testing.T) {
+	for _, test := range []struct {
+		name      string
+		options   []api.PurgeOptions
+		recursive bool
+	}{
+		{name: "default"},
+		{name: "recursive", options: []api.PurgeOptions{api.WithRecursivePurge(true)}, recursive: true},
+		{name: "nonrecursive", options: []api.PurgeOptions{api.WithRecursivePurge(false)}},
+	} {
+		t.Run(test.name, func(t *testing.T) {
+			requests := make(chan *protos.PurgeInstancesRequest, 1)
+			client := startQueryClient(t, &managementServer{
+				purge: func(_ context.Context, req *protos.PurgeInstancesRequest) (*protos.PurgeInstancesResponse, error) {
+					requests <- req
+					return &protos.PurgeInstancesResponse{DeletedInstanceCount: 1}, nil
+				},
+			})
+			ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+			defer cancel()
+
+			require.NoError(t, client.PurgeOrchestrationState(ctx, "completed-instance", test.options...))
+			req := <-requests
+			require.IsType(t, &protos.PurgeInstancesRequest_InstanceId{}, req.Request)
+			require.Equal(t, "completed-instance", req.GetInstanceId())
+			require.Equal(t, test.recursive, req.GetRecursive())
+			require.True(t, req.GetIsOrchestration(), "single-instance purge must identify an orchestration")
+		})
+	}
+}
+
+func TestPurgeOrchestrationStateMissingInstance(t *testing.T) {
+	client := startQueryClient(t, &managementServer{
+		purge: func(context.Context, *protos.PurgeInstancesRequest) (*protos.PurgeInstancesResponse, error) {
+			return &protos.PurgeInstancesResponse{}, nil
+		},
+	})
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+
+	require.ErrorIs(t, client.PurgeOrchestrationState(ctx, "missing-instance"), api.ErrInstanceNotFound)
+}
+
+func TestPurgeOrchestrationStateOptionError(t *testing.T) {
+	client := &TaskHubGrpcClient{}
+	optionErr := errors.New("invalid purge option")
+	err := client.PurgeOrchestrationState(context.Background(), "instance", func(*protos.PurgeInstancesRequest) error {
+		return optionErr
+	})
+	require.ErrorIs(t, err, api.ErrInvalidArgument)
+	require.ErrorIs(t, err, optionErr)
+}

tests/durabletaskscheduler/dts_emulator_test.go

@@ -559,6 +559,58 @@ func TestDTSEmulatorSequenceMetadataAndPurge(t *testing.T) {
 	require.ErrorIs(t, err, api.ErrInstanceNotFound)
 }
 
+func TestDTSSingleInstancePurge(t *testing.T) {
+	name := string(uniqueInstanceID("DTSSingleInstancePurge"))
+	registry := task.NewTaskRegistry()
+	require.NoError(t, registry.AddOrchestratorN(name, func(*task.OrchestrationContext) (any, error) {
+		return "done", nil
+	}))
+	managementClient, _, _ := startEmulatorClientAndWorker(t, registry, durabletaskclient.WithAutoWorkItemFilters())
+
+	ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
+	defer cancel()
+	targetID := uniqueInstanceID("go-single-purge")
+	controlID := uniqueInstanceID("go-single-purge-control")
+	for _, id := range []api.InstanceID{targetID, controlID} {
+		t.Cleanup(func() {
+			cleanupCtx, cleanupCancel := context.WithTimeout(context.Background(), 30*time.Second)
+			defer cleanupCancel()
+			metadata, err := managementClient.FetchOrchestrationMetadata(cleanupCtx, id)
+			if errors.Is(err, api.ErrInstanceNotFound) {
+				return
+			}
+			require.NoError(t, err)
+			if !metadata.IsComplete() {
+				require.NoError(t, managementClient.TerminateOrchestration(cleanupCtx, id, api.WithRecursiveTerminate(false)))
+				_, err = managementClient.WaitForOrchestrationCompletion(cleanupCtx, id)
+				require.NoError(t, err)
+			}
+			_, err = managementClient.PurgeInstances(cleanupCtx, api.PurgeInstancesRequest{InstanceIDs: []api.InstanceID{id}})
+			require.NoError(t, err)
+			metadata, err = managementClient.FetchOrchestrationMetadata(cleanupCtx, id)
+			require.ErrorIs(t, err, api.ErrInstanceNotFound)
+			require.Nil(t, metadata)
+			t.Logf("cleaned up owned orchestration %s", id)
+		})
+		t.Logf("owned orchestration %s", id)
+		_, err := managementClient.ScheduleNewOrchestration(ctx, name, api.WithInstanceID(id))
+		require.NoError(t, err)
+		metadata, err := managementClient.WaitForOrchestrationCompletion(ctx, id)
+		require.NoError(t, err)
+		require.Equal(t, api.RUNTIME_STATUS_COMPLETED, metadata.RuntimeStatus)
+	}
+
+	require.NoError(t, managementClient.PurgeOrchestrationState(ctx, targetID))
+	metadata, err := managementClient.FetchOrchestrationMetadata(ctx, targetID)
+	require.ErrorIs(t, err, api.ErrInstanceNotFound)
+	require.Nil(t, metadata)
+	control, err := managementClient.FetchOrchestrationMetadata(ctx, controlID)
+	require.NoError(t, err)
+	require.Equal(t, api.RUNTIME_STATUS_COMPLETED, control.RuntimeStatus)
+	require.Equal(t, controlID, control.InstanceID)
+	t.Logf("purged %s; completed control %s remains", targetID, controlID)
+}
+
 func TestDTSEmulatorDurableEntities(t *testing.T) {
 	registry := task.NewTaskRegistry()
 	require.NoError(t, registry.AddEntityN("counter", func(ctx *task.EntityContext) (any, error) {