Delegate branch token serialization/deserialization to store layer (#…

…3225)

common/metrics/defs.go

@@ -739,6 +739,10 @@ const (
 	PersistenceAppendRawHistoryNodesScope
 	// PersistenceDeleteHistoryNodesScope tracks DeleteHistoryNodes calls made by service to persistence layer
 	PersistenceDeleteHistoryNodesScope
+	// PersistenceParseHistoryBranchInfoScope tracks NewHistoryBranch calls made by service to persistence layer
+	PersistenceParseHistoryBranchInfoScope
+	// PersistenceUpdateHistoryBranchInfoScope tracks NewHistoryBranch calls made by service to persistence layer
+	PersistenceUpdateHistoryBranchInfoScope
 	// PersistenceNewHistoryBranchScope tracks NewHistoryBranch calls made by service to persistence layer
 	PersistenceNewHistoryBranchScope
 	// PersistenceReadHistoryBranchScope tracks ReadHistoryBranch calls made by service to persistence layer
@@ -1408,6 +1412,8 @@ var ScopeDefs = map[ServiceIdx]map[int]scopeDefinition{
 		PersistenceAppendHistoryNodesScope:         {operation: "AppendHistoryNodes"},
 		PersistenceAppendRawHistoryNodesScope:      {operation: "AppendRawHistoryNodes"},
 		PersistenceDeleteHistoryNodesScope:         {operation: "DeleteHistoryNodes"},
+		PersistenceParseHistoryBranchInfoScope:     {operation: "ParseHistoryBranch"},
+		PersistenceUpdateHistoryBranchInfoScope:    {operation: "UpdateHistoryBranch"},
 		PersistenceNewHistoryBranchScope:           {operation: "NewHistoryBranch"},
 		PersistenceReadHistoryBranchScope:          {operation: "ReadHistoryBranch"},
 		PersistenceReadHistoryBranchReverseScope:   {operation: "ReadHistoryBranchReverse"},

common/persistence/cassandra/history_store.go

@@ -162,6 +162,36 @@ func (h *HistoryStore) DeleteHistoryNodes(
 	return nil
 }
 
+// ParseHistoryBranchInfo parses the history branch for branch information
+func (h *HistoryStore) ParseHistoryBranchInfo(
+	ctx context.Context,
+	request *p.ParseHistoryBranchInfoRequest,
+) (*p.ParseHistoryBranchInfoResponse, error) {
+
+	branchInfo, err := p.ParseHistoryBranchToken(request.BranchToken)
+	if err != nil {
+		return nil, err
+	}
+	return &p.ParseHistoryBranchInfoResponse{
+		BranchInfo: branchInfo,
+	}, nil
+}
+
+// UpdateHistoryBranchInfo updates the history branch with branch information
+func (h *HistoryStore) UpdateHistoryBranchInfo(
+	ctx context.Context,
+	request *p.UpdateHistoryBranchInfoRequest,
+) (*p.UpdateHistoryBranchInfoResponse, error) {
+
+	branchToken, err := p.UpdateHistoryBranchToken(request.BranchToken, request.BranchInfo)
+	if err != nil {
+		return nil, err
+	}
+	return &p.UpdateHistoryBranchInfoResponse{
+		BranchToken: branchToken,
+	}, nil
+}
+
 // NewHistoryBranch initializes a new history branch
 func (h *HistoryStore) NewHistoryBranch(
 	ctx context.Context,

common/persistence/client/fault_injection.go

@@ -618,6 +618,26 @@ func (e *FaultInjectionExecutionStore) DeleteHistoryNodes(
 	return e.baseExecutionStore.DeleteHistoryNodes(ctx, request)
 }
 
+func (e *FaultInjectionExecutionStore) ParseHistoryBranchInfo(
+	ctx context.Context,
+	request *persistence.ParseHistoryBranchInfoRequest,
+) (*persistence.ParseHistoryBranchInfoResponse, error) {
+	if err := e.ErrorGenerator.Generate(); err != nil {
+		return nil, err
+	}
+	return e.baseExecutionStore.ParseHistoryBranchInfo(ctx, request)
+}
+
+func (e *FaultInjectionExecutionStore) UpdateHistoryBranchInfo(
+	ctx context.Context,
+	request *persistence.UpdateHistoryBranchInfoRequest,
+) (*persistence.UpdateHistoryBranchInfoResponse, error) {
+	if err := e.ErrorGenerator.Generate(); err != nil {
+		return nil, err
+	}
+	return e.baseExecutionStore.UpdateHistoryBranchInfo(ctx, request)
+}
+
 func (e *FaultInjectionExecutionStore) NewHistoryBranch(
 	ctx context.Context,
 	request *persistence.NewHistoryBranchRequest,

common/persistence/dataInterfaces.go

@@ -740,7 +740,30 @@ type (
 		NodeID int64
 	}
 
+	ParseHistoryBranchInfoRequest struct {
+		// The branch token to parse the branch info from
+		BranchToken []byte
+	}
+
+	ParseHistoryBranchInfoResponse struct {
+		// The branch info parsed from the branch token
+		BranchInfo *persistencespb.HistoryBranch
+	}
+
+	UpdateHistoryBranchInfoRequest struct {
+		// The original branch token
+		BranchToken []byte
+		// The branch info to update with
+		BranchInfo *persistencespb.HistoryBranch
+	}
+
+	UpdateHistoryBranchInfoResponse struct {
+		// The newly updated branch token
+		BranchToken []byte
+	}
+
 	NewHistoryBranchRequest struct {
+		// The tree ID for the new branch token
 		TreeID string
 	}
 
@@ -1060,6 +1083,10 @@ type (
 		AppendHistoryNodes(ctx context.Context, request *AppendHistoryNodesRequest) (*AppendHistoryNodesResponse, error)
 		// AppendRawHistoryNodes add a node of raw histories to history node table
 		AppendRawHistoryNodes(ctx context.Context, request *AppendRawHistoryNodesRequest) (*AppendHistoryNodesResponse, error)
+		// ParseHistoryBranchInfo parses the history branch for branch information
+		ParseHistoryBranchInfo(ctx context.Context, request *ParseHistoryBranchInfoRequest) (*ParseHistoryBranchInfoResponse, error)
+		// UpdateHistoryBranchInfo updates the history branch with branch information
+		UpdateHistoryBranchInfo(ctx context.Context, request *UpdateHistoryBranchInfoRequest) (*UpdateHistoryBranchInfoResponse, error)
 		// NewHistoryBranch initializes a new history branch
 		NewHistoryBranch(ctx context.Context, request *NewHistoryBranchRequest) (*NewHistoryBranchResponse, error)
 		// ReadHistoryBranch returns history node data for a branch
@@ -1196,6 +1223,28 @@ func UnixMilliseconds(t time.Time) int64 {
 	return unixNano / int64(time.Millisecond)
 }
 
+func ParseHistoryBranchToken(branchToken []byte) (*persistencespb.HistoryBranch, error) {
+	// TODO: instead of always using the implementation from the serialization package, this should be injected
+	return serialization.HistoryBranchFromBlob(branchToken, enumspb.ENCODING_TYPE_PROTO3.String())
+}
+
+func UpdateHistoryBranchToken(branchToken []byte, branchInfo *persistencespb.HistoryBranch) ([]byte, error) {
+	bi, err := ParseHistoryBranchToken(branchToken)
+	if err != nil {
+		return nil, err
+	}
+	bi.TreeId = branchInfo.TreeId
+	bi.BranchId = branchInfo.BranchId
+	bi.Ancestors = branchInfo.Ancestors
+
+	// TODO: instead of always using the implementation from the serialization package, this should be injected
+	blob, err := serialization.HistoryBranchToBlob(bi)
+	if err != nil {
+		return nil, err
+	}
+	return blob.Data, nil
+}
+
 // NewHistoryBranchToken return a new branch token
 func NewHistoryBranchToken(treeID string) ([]byte, error) {
 	branchID := primitives.NewUUID().String()
@@ -1204,6 +1253,7 @@ func NewHistoryBranchToken(treeID string) ([]byte, error) {
 		BranchId:  branchID,
 		Ancestors: []*persistencespb.HistoryBranchRange{},
 	}
+	// TODO: instead of always using the implementation from the serialization package, this should be injected
 	datablob, err := serialization.HistoryBranchToBlob(bi)
 	if err != nil {
 		return nil, err
@@ -1219,6 +1269,7 @@ func NewHistoryBranchTokenByBranchID(treeID, branchID string) ([]byte, error) {
 		BranchId:  branchID,
 		Ancestors: []*persistencespb.HistoryBranchRange{},
 	}
+	// TODO: instead of always using the implementation from the serialization package, this should be injected
 	datablob, err := serialization.HistoryBranchToBlob(bi)
 	if err != nil {
 		return nil, err

common/persistence/execution_manager.go

@@ -87,7 +87,7 @@ func (m *executionManagerImpl) CreateWorkflowExecution(
 ) (*CreateWorkflowExecutionResponse, error) {
 
 	newSnapshot := request.NewWorkflowSnapshot
-	newWorkflowNewEvents, newHistoryDiff, err := m.serializeWorkflowEventBatches(request.ShardID, request.NewWorkflowEvents)
+	newWorkflowNewEvents, newHistoryDiff, err := m.serializeWorkflowEventBatches(ctx, request.ShardID, request.NewWorkflowEvents)
 	if err != nil {
 		return nil, err
 	}
@@ -140,13 +140,13 @@ func (m *executionManagerImpl) UpdateWorkflowExecution(
 	updateMutation := request.UpdateWorkflowMutation
 	newSnapshot := request.NewWorkflowSnapshot
 
-	updateWorkflowNewEvents, updateWorkflowHistoryDiff, err := m.serializeWorkflowEventBatches(request.ShardID, request.UpdateWorkflowEvents)
+	updateWorkflowNewEvents, updateWorkflowHistoryDiff, err := m.serializeWorkflowEventBatches(ctx, request.ShardID, request.UpdateWorkflowEvents)
 	if err != nil {
 		return nil, err
 	}
 	updateMutation.ExecutionInfo.ExecutionStats.HistorySize += int64(updateWorkflowHistoryDiff.SizeDiff)
 
-	newWorkflowNewEvents, newWorkflowHistoryDiff, err := m.serializeWorkflowEventBatches(request.ShardID, request.NewWorkflowEvents)
+	newWorkflowNewEvents, newWorkflowHistoryDiff, err := m.serializeWorkflowEventBatches(ctx, request.ShardID, request.NewWorkflowEvents)
 	if err != nil {
 		return nil, err
 	}
@@ -230,21 +230,21 @@ func (m *executionManagerImpl) ConflictResolveWorkflowExecution(
 	newSnapshot := request.NewWorkflowSnapshot
 	currentMutation := request.CurrentWorkflowMutation
 
-	resetWorkflowEventsNewEvents, resetWorkflowHistoryDiff, err := m.serializeWorkflowEventBatches(request.ShardID, request.ResetWorkflowEvents)
+	resetWorkflowEventsNewEvents, resetWorkflowHistoryDiff, err := m.serializeWorkflowEventBatches(ctx, request.ShardID, request.ResetWorkflowEvents)
 	if err != nil {
 		return nil, err
 	}
 	resetSnapshot.ExecutionInfo.ExecutionStats.HistorySize += int64(resetWorkflowHistoryDiff.SizeDiff)
 
-	newWorkflowEventsNewEvents, newWorkflowHistoryDiff, err := m.serializeWorkflowEventBatches(request.ShardID, request.NewWorkflowEvents)
+	newWorkflowEventsNewEvents, newWorkflowHistoryDiff, err := m.serializeWorkflowEventBatches(ctx, request.ShardID, request.NewWorkflowEvents)
 	if err != nil {
 		return nil, err
 	}
 	if newSnapshot != nil {
 		newSnapshot.ExecutionInfo.ExecutionStats.HistorySize += int64(newWorkflowHistoryDiff.SizeDiff)
 	}
 
-	currentWorkflowEventsNewEvents, currentWorkflowHistoryDiff, err := m.serializeWorkflowEventBatches(request.ShardID, request.CurrentWorkflowEvents)
+	currentWorkflowEventsNewEvents, currentWorkflowHistoryDiff, err := m.serializeWorkflowEventBatches(ctx, request.ShardID, request.CurrentWorkflowEvents)
 	if err != nil {
 		return nil, err
 	}
@@ -388,6 +388,7 @@ func (m *executionManagerImpl) SetWorkflowExecution(
 }
 
 func (m *executionManagerImpl) serializeWorkflowEventBatches(
+	ctx context.Context,
 	shardID int32,
 	eventBatches []*WorkflowEvents,
 ) ([]*InternalAppendHistoryNodesRequest, *HistoryStatistics, error) {
@@ -398,7 +399,7 @@ func (m *executionManagerImpl) serializeWorkflowEventBatches(
 
 	workflowNewEvents := make([]*InternalAppendHistoryNodesRequest, 0, len(eventBatches))
 	for _, workflowEvents := range eventBatches {
-		newEvents, err := m.serializeWorkflowEvents(shardID, workflowEvents)
+		newEvents, err := m.serializeWorkflowEvents(ctx, shardID, workflowEvents)
 		if err != nil {
 			return nil, nil, err
 		}
@@ -432,6 +433,7 @@ func (m *executionManagerImpl) DeserializeBufferedEvents( // unexport
 }
 
 func (m *executionManagerImpl) serializeWorkflowEvents(
+	ctx context.Context,
 	shardID int32,
 	workflowEvents *WorkflowEvents,
 ) (*InternalAppendHistoryNodesRequest, error) {
@@ -452,7 +454,7 @@ func (m *executionManagerImpl) serializeWorkflowEvents(
 		request.Info = BuildHistoryGarbageCleanupInfo(workflowEvents.NamespaceID, workflowEvents.WorkflowID, workflowEvents.RunID)
 	}
 
-	return m.serializeAppendHistoryNodesRequest(request)
+	return m.serializeAppendHistoryNodesRequest(ctx, request)
 }
 
 func (m *executionManagerImpl) SerializeWorkflowMutation( // unexport