Update storage driver store context (temporalio#2268)

Author: jmaeagle99

contrib/aws/s3driver/driver.go

@@ -5,6 +5,7 @@ import (
 	"encoding/hex"
 	"errors"
 	"fmt"
+	"net/url"
 
 	commonpb "go.temporal.io/api/common/v1"
 	"go.temporal.io/sdk/converter"
@@ -144,7 +145,7 @@ func (d *s3StorageDriver) Store(
 	g, gctx := errgroup.WithContext(ctx.Context)
 	for i, pp := range prepared {
 		g.Go(func() error {
-			key := objectKey(pp.hexDigest)
+			key := objectKey(ctx.Target, pp.hexDigest)
 			exists, err := d.client.ObjectExists(gctx, pp.bucket, key)
 			if err != nil {
 				return fmt.Errorf("existence check failed [bucket=%s, key=%s]: %w", pp.bucket, key, err)
@@ -231,8 +232,33 @@ func (d *s3StorageDriver) Retrieve(
 	return payloads, nil
 }
 
-func objectKey(hexDigest string) string {
-	return keyVersion + "/d/" + hashAlgorithm + "/" + hexDigest
+func objectKey(target converter.StorageDriverTargetInfo, hexDigest string) string {
+	digestSegment := "/d/" + hashAlgorithm + "/" + hexDigest
+	switch t := target.(type) {
+	case converter.StorageDriverWorkflowInfo:
+		return keyVersion +
+			"/ns/" + pathEscape(t.Namespace) +
+			"/wt/" + pathEscape(t.WorkflowType) +
+			"/wi/" + pathEscape(t.WorkflowID) +
+			"/ri/" + pathEscape(t.RunID) +
+			digestSegment
+	case converter.StorageDriverActivityInfo:
+		return keyVersion +
+			"/ns/" + pathEscape(t.Namespace) +
+			"/at/" + pathEscape(t.ActivityType) +
+			"/ai/" + pathEscape(t.ActivityID) +
+			"/ri/" + pathEscape(t.RunID) +
+			digestSegment
+	default:
+		return keyVersion + digestSegment
+	}
+}
+
+func pathEscape(s string) string {
+	if s == "" {
+		return "null"
+	}
+	return url.PathEscape(s)
 }
 
 func sha256Hex(data []byte) string {

contrib/aws/s3driver/driver_test.go

@@ -81,6 +81,10 @@ func storeCtx() converter.StorageDriverStoreContext {
 	return converter.StorageDriverStoreContext{Context: context.Background()}
 }
 
+func storeCtxWithTarget(target converter.StorageDriverTargetInfo) converter.StorageDriverStoreContext {
+	return converter.StorageDriverStoreContext{Context: context.Background(), Target: target}
+}
+
 func retrieveCtx() converter.StorageDriverRetrieveContext {
 	return converter.StorageDriverRetrieveContext{Context: context.Background()}
 }
@@ -459,8 +463,72 @@ func TestRetrieve_ClaimMissingHashValue(t *testing.T) {
 
 // --- Key generation tests ---
 
-func TestObjectKey(t *testing.T) {
-	assert.Equal(t, "v0/d/sha256/abc123", objectKey("abc123"))
+func TestObjectKey_NoTarget(t *testing.T) {
+	assert.Equal(t, "v0/d/sha256/abc123", objectKey(nil, "abc123"))
+}
+
+func TestObjectKey_WorkflowInfo(t *testing.T) {
+	target := converter.StorageDriverWorkflowInfo{
+		Namespace:    "default",
+		WorkflowType: "MyWorkflow",
+		WorkflowID:   "wf-123",
+		RunID:        "run-456",
+	}
+	assert.Equal(t,
+		"v0/ns/default/wt/MyWorkflow/wi/wf-123/ri/run-456/d/sha256/abc123",
+		objectKey(target, "abc123"),
+	)
+}
+
+func TestObjectKey_ActivityInfo(t *testing.T) {
+	target := converter.StorageDriverActivityInfo{
+		Namespace:    "default",
+		ActivityType: "MyActivity",
+		ActivityID:   "act-789",
+		RunID:        "run-abc",
+	}
+	assert.Equal(t,
+		"v0/ns/default/at/MyActivity/ai/act-789/ri/run-abc/d/sha256/abc123",
+		objectKey(target, "abc123"),
+	)
+}
+
+func TestObjectKey_WorkflowInfo_EmptyFields(t *testing.T) {
+	// Empty strings should fall back to "null" in each segment.
+	target := converter.StorageDriverWorkflowInfo{
+		Namespace: "my-ns",
+		// WorkflowType, WorkflowID, RunID intentionally empty
+	}
+	assert.Equal(t,
+		"v0/ns/my-ns/wt/null/wi/null/ri/null/d/sha256/abc123",
+		objectKey(target, "abc123"),
+	)
+}
+
+func TestObjectKey_ActivityInfo_EmptyFields(t *testing.T) {
+	target := converter.StorageDriverActivityInfo{
+		Namespace: "my-ns",
+		// ActivityType, ActivityID, RunID intentionally empty
+	}
+	assert.Equal(t,
+		"v0/ns/my-ns/at/null/ai/null/ri/null/d/sha256/abc123",
+		objectKey(target, "abc123"),
+	)
+}
+
+func TestObjectKey_WorkflowInfo_SpecialChars(t *testing.T) {
+	// Slashes, spaces, and other special characters must be percent-encoded.
+	target := converter.StorageDriverWorkflowInfo{
+		Namespace:    "my namespace",
+		WorkflowType: "my/workflow",
+		WorkflowID:   "wf id+1",
+		RunID:        "run=abc",
+	}
+	key := objectKey(target, "abc123")
+	assert.Equal(t,
+		"v0/ns/my%20namespace/wt/my%2Fworkflow/wi/wf%20id+1/ri/run=abc/d/sha256/abc123",
+		key,
+	)
 }
 
 func TestSha256Hex(t *testing.T) {
@@ -469,3 +537,81 @@ func TestSha256Hex(t *testing.T) {
 	expected := hex.EncodeToString(h[:])
 	assert.Equal(t, expected, sha256Hex(data))
 }
+
+// --- Store with target context tests ---
+
+func TestStore_WithWorkflowTarget(t *testing.T) {
+	mc := newMemClient()
+	d := newDriver(t, mc)
+	p := testPayload("workflow payload")
+
+	target := converter.StorageDriverWorkflowInfo{
+		Namespace:    "default",
+		WorkflowType: "MyWorkflow",
+		WorkflowID:   "wf-123",
+		RunID:        "run-456",
+	}
+	claims, err := d.Store(storeCtxWithTarget(target), []*commonpb.Payload{p})
+	require.NoError(t, err)
+	require.Len(t, claims, 1)
+
+	key := claims[0].ClaimData["key"]
+	assert.Contains(t, key, "v0/ns/default/wt/MyWorkflow/wi/wf-123/ri/run-456/d/sha256/")
+}
+
+func TestStore_WithActivityTarget(t *testing.T) {
+	mc := newMemClient()
+	d := newDriver(t, mc)
+	p := testPayload("activity payload")
+
+	target := converter.StorageDriverActivityInfo{
+		Namespace:    "default",
+		ActivityType: "MyActivity",
+		ActivityID:   "act-789",
+		RunID:        "run-abc",
+	}
+	claims, err := d.Store(storeCtxWithTarget(target), []*commonpb.Payload{p})
+	require.NoError(t, err)
+	require.Len(t, claims, 1)
+
+	key := claims[0].ClaimData["key"]
+	assert.Contains(t, key, "v0/ns/default/at/MyActivity/ai/act-789/ri/run-abc/d/sha256/")
+}
+
+func TestStore_RoundTrip_WithWorkflowTarget(t *testing.T) {
+	mc := newMemClient()
+	d := newDriver(t, mc)
+	original := testPayload("round-trip with target")
+
+	target := converter.StorageDriverWorkflowInfo{
+		Namespace:    "default",
+		WorkflowType: "MyWorkflow",
+		WorkflowID:   "wf-123",
+		RunID:        "run-456",
+	}
+	claims, err := d.Store(storeCtxWithTarget(target), []*commonpb.Payload{original})
+	require.NoError(t, err)
+
+	restored, err := d.Retrieve(retrieveCtx(), claims)
+	require.NoError(t, err)
+	require.Len(t, restored, 1)
+	assert.True(t, proto.Equal(original, restored[0]))
+}
+
+func TestStore_DifferentTargets_SamePayload_DifferentKeys(t *testing.T) {
+	// The same payload stored under different targets produces different keys.
+	mc := newMemClient()
+	d := newDriver(t, mc)
+	p := testPayload("shared payload")
+
+	wfTarget := converter.StorageDriverWorkflowInfo{Namespace: "ns", WorkflowID: "wf-1", RunID: "run-1"}
+	actTarget := converter.StorageDriverActivityInfo{Namespace: "ns", ActivityID: "act-1", RunID: "run-1"}
+
+	wfClaims, err := d.Store(storeCtxWithTarget(wfTarget), []*commonpb.Payload{p})
+	require.NoError(t, err)
+
+	actClaims, err := d.Store(storeCtxWithTarget(actTarget), []*commonpb.Payload{p})
+	require.NoError(t, err)
+
+	assert.NotEqual(t, wfClaims[0].ClaimData["key"], actClaims[0].ClaimData["key"])
+}

converter/extstore.go

@@ -6,6 +6,53 @@ import (
 	commonpb "go.temporal.io/api/common/v1"
 )
 
+// StorageDriverTargetInfo identifies the workflow or activity on whose behalf
+// a payload is being stored. Use a type switch on [StorageDriverWorkflowInfo]
+// and [StorageDriverActivityInfo] to access the concrete values.
+//
+// NOTE: Experimental
+type StorageDriverTargetInfo interface {
+	isStorageDriverTargetInfo()
+}
+
+// StorageDriverWorkflowInfo carries workflow identity for a storage operation.
+//
+// NOTE: Experimental
+type StorageDriverWorkflowInfo struct {
+	// Namespace is the Temporal namespace of the workflow execution.
+	Namespace string
+	// WorkflowType is the type name of the workflow.
+	WorkflowType string
+	// WorkflowID is the ID of the workflow execution.
+	WorkflowID string
+	// RunID is the run ID of the workflow execution.
+	RunID string
+}
+
+func (StorageDriverWorkflowInfo) isStorageDriverTargetInfo() {}
+
+var _ StorageDriverTargetInfo = StorageDriverWorkflowInfo{}
+
+// StorageDriverActivityInfo carries activity identity for a storage operation.
+// This is only used for standalone (non-workflow-bound) activities; activities
+// started by a workflow use [StorageDriverWorkflowInfo] as the target.
+//
+// NOTE: Experimental
+type StorageDriverActivityInfo struct {
+	// Namespace is the Temporal namespace of the activity execution.
+	Namespace string
+	// ActivityType is the type name of the activity.
+	ActivityType string
+	// ActivityID is the ID of the activity execution.
+	ActivityID string
+	// RunID is the run ID of the activity execution.
+	RunID string
+}
+
+func (StorageDriverActivityInfo) isStorageDriverTargetInfo() {}
+
+var _ StorageDriverTargetInfo = StorageDriverActivityInfo{}
+
 // StorageDriverStoreContext carries context passed to StorageDriver.Store and
 // StorageDriverSelector.SelectDriver operations.
 //
@@ -15,6 +62,10 @@ type StorageDriverStoreContext struct {
 	// Drivers should use it to respect cancellation and to propagate deadlines
 	// and trace information to downstream calls (e.g. cloud storage SDKs).
 	Context context.Context
+	// Target identifies the workflow or activity on whose behalf payloads are
+	// being stored. Use a type switch on [StorageDriverWorkflowInfo] and
+	// [StorageDriverActivityInfo] to access the concrete values.
+	Target StorageDriverTargetInfo
 }
 
 // StorageDriverRetrieveContext carries context passed to StorageDriver.Retrieve

internal/extstore.go

@@ -83,6 +83,7 @@ type storageOperationCallback interface {
 }
 
 const storageOperationCallbackContextKey contextKey = "storageOperationCallback"
+const storageTargetContextKey contextKey = "storageTarget"
 
 // metadataEncodingStorageRef is the metadata encoding value used to identify
 // payloads that are storage references rather than actual data.
@@ -247,7 +248,11 @@ func (v *externalStorageVisitor) Visit(ctx *proxy.VisitPayloadsContext, payloads
 	driverBatches := map[string]*driverBatch{}
 
 	result := make([]*commonpb.Payload, len(payloads))
-	driverCtx := converter.StorageDriverStoreContext{Context: ctx.Context}
+	var target converter.StorageDriverTargetInfo
+	if t, ok := ctx.Context.Value(storageTargetContextKey).(converter.StorageDriverTargetInfo); ok {
+		target = t
+	}
+	driverCtx := converter.StorageDriverStoreContext{Context: ctx.Context, Target: target}
 
 	for i, p := range payloads {
 		if proto.Size(p) < v.params.payloadSizeThreshold {
@@ -287,7 +292,7 @@ func (v *externalStorageVisitor) Visit(ctx *proxy.VisitPayloadsContext, payloads
 	// Fan out to each driver concurrently. The errgroup context is used as the
 	// StorageDriverStoreContext so a failing driver cancels in-flight siblings.
 	eg, egCtx := errgroup.WithContext(ctx.Context)
-	storeDrCtx := converter.StorageDriverStoreContext{Context: egCtx}
+	storeDrCtx := converter.StorageDriverStoreContext{Context: egCtx, Target: target}
 	sizes := make([]int64, len(driverOrder))
 
 	externalCount := 0

internal/internal_activity_client.go

@@ -587,7 +587,12 @@ func (w *workflowClientInterceptor) ExecuteActivity(
 		return nil, err
 	}
 
-	if err := visitProtoPayloads(ctx, w.client.outboundPayloadVisitor, request); err != nil {
+	storeCtx := context.WithValue(ctx, storageTargetContextKey, converter.StorageDriverActivityInfo{
+		Namespace:    w.client.namespace,
+		ActivityID:   request.ActivityId,
+		ActivityType: in.ActivityType,
+	})
+	if err := visitProtoPayloads(storeCtx, w.client.outboundPayloadVisitor, request); err != nil {
 		return nil, err
 	}
 

internal/internal_schedule_client.go

@@ -131,6 +131,15 @@ func (w *workflowClientInterceptor) CreateSchedule(ctx context.Context, in *Sche
 		SearchAttributes: searchAttr,
 	}
 
+	storeCtx := context.WithValue(ctx, storageTargetContextKey, converter.StorageDriverWorkflowInfo{
+		Namespace:    w.client.namespace,
+		WorkflowID:   action.GetStartWorkflow().GetWorkflowId(),
+		WorkflowType: action.GetStartWorkflow().GetWorkflowType().GetName(),
+	})
+	if err := visitProtoPayloads(storeCtx, w.client.outboundPayloadVisitor, startRequest); err != nil {
+		return nil, err
+	}
+
 	grpcCtx, cancel := newGRPCContext(ctx, defaultGrpcRetryParameters(ctx))
 	defer cancel()
 
@@ -283,15 +292,26 @@ func (scheduleHandle *scheduleHandleImpl) Update(ctx context.Context, options Sc
 		}
 	}
 
-	_, err = scheduleHandle.client.workflowService.UpdateSchedule(grpcCtx, &workflowservice.UpdateScheduleRequest{
+	updateRequest := &workflowservice.UpdateScheduleRequest{
 		Namespace:        scheduleHandle.client.namespace,
 		ScheduleId:       scheduleHandle.ID,
 		Schedule:         newSchedulePB,
 		ConflictToken:    nil,
 		Identity:         scheduleHandle.client.identity,
 		RequestId:        uuid.NewString(),
 		SearchAttributes: newSA,
+	}
+
+	storeCtx := context.WithValue(ctx, storageTargetContextKey, converter.StorageDriverWorkflowInfo{
+		Namespace:    scheduleHandle.client.namespace,
+		WorkflowID:   newSchedulePB.GetAction().GetStartWorkflow().GetWorkflowId(),
+		WorkflowType: newSchedulePB.GetAction().GetStartWorkflow().GetWorkflowType().GetName(),
 	})
+	if err := visitProtoPayloads(storeCtx, scheduleHandle.client.outboundPayloadVisitor, updateRequest); err != nil {
+		return err
+	}
+
+	_, err = scheduleHandle.client.workflowService.UpdateSchedule(grpcCtx, updateRequest)
 	return err
 }