Add corruption protection to ms (#4560)

Author: wxing1292

common/metrics/metric_defs.go

@@ -1497,6 +1497,7 @@ var (
 	WorkflowRunTimeoutOverrideCount                = NewCounterDef("workflow_run_timeout_overrides")
 	ReplicationTaskCleanupCount                    = NewCounterDef("replication_task_cleanup_count")
 	ReplicationTaskCleanupFailure                  = NewCounterDef("replication_task_cleanup_failed")
+	MutableStateDirty                              = NewCounterDef("mutable_state_dirty")
 	MutableStateChecksumMismatch                   = NewCounterDef("mutable_state_checksum_mismatch")
 	MutableStateChecksumInvalidated                = NewCounterDef("mutable_state_checksum_invalidated")
 	ClusterMetadataLockLatency                     = NewTimerDef("cluster_metadata_lock_latency")

service/history/workflow/history_builder.go

@@ -154,6 +154,13 @@ func NewImmutableHistoryBuilder(
 	}
 }
 
+func (b *HistoryBuilder) IsDirty() bool {
+	return len(b.memEventsBatches) > 0 ||
+		len(b.memLatestBatch) > 0 ||
+		len(b.memBufferBatch) > 0 ||
+		len(b.scheduledIDToStartedID) > 0
+}
+
 // NOTE:
 // originalRunID is the runID when the WorkflowExecutionStarted event is written
 // firstRunID is the very first runID along the chain of ContinueAsNew and Reset

service/history/workflow/mutable_state_impl.go

@@ -4378,6 +4378,17 @@ func (ms *MutableStateImpl) UpdateWorkflowStateStatus(
 func (ms *MutableStateImpl) StartTransaction(
 	namespaceEntry *namespace.Namespace,
 ) (bool, error) {
+	if ms.hBuilder.IsDirty() || len(ms.InsertTasks) > 0 {
+		ms.logger.Error("MutableState encountered dirty transaction",
+			tag.WorkflowNamespaceID(ms.executionInfo.NamespaceId),
+			tag.WorkflowID(ms.executionInfo.WorkflowId),
+			tag.WorkflowRunID(ms.executionState.RunId),
+			tag.Value(ms.hBuilder),
+		)
+		ms.metricsHandler.Counter(metrics.MutableStateChecksumInvalidated.GetMetricName()).Record(1)
+		return false, serviceerror.NewUnavailable("MutableState encountered dirty transaction")
+	}
+
 	namespaceEntry, err := ms.startTransactionHandleNamespaceMigration(namespaceEntry)
 	if err != nil {
 		return false, err