Fix internal flags not always being written when they should be (#532)

* Update protos subtree to 84a8e55fa275185934f242616e3e0f0f84788971

Author: Sushisource

client/src/raw.rs

@@ -740,19 +740,19 @@ proxier! {
         }
     );
     (
-        update_worker_build_id_ordering,
-        UpdateWorkerBuildIdOrderingRequest,
-        UpdateWorkerBuildIdOrderingResponse,
+        update_worker_build_id_compatibility,
+        UpdateWorkerBuildIdCompatibilityRequest,
+        UpdateWorkerBuildIdCompatibilityResponse,
         |r| {
             let mut labels = AttachMetricLabels::namespace(r.get_ref().namespace.clone());
             labels.task_q_str(r.get_ref().task_queue.clone());
             r.extensions_mut().insert(labels);
         }
     );
     (
-        get_worker_build_id_ordering,
-        GetWorkerBuildIdOrderingRequest,
-        GetWorkerBuildIdOrderingResponse,
+        get_worker_build_id_compatibility,
+        GetWorkerBuildIdCompatibilityRequest,
+        GetWorkerBuildIdCompatibilityResponse,
         |r| {
             let mut labels = AttachMetricLabels::namespace(r.get_ref().namespace.clone());
             labels.task_q_str(r.get_ref().task_queue.clone());
@@ -768,6 +768,15 @@ proxier! {
             r.extensions_mut().insert(labels);
         }
     );
+    (
+        poll_workflow_execution_update,
+        PollWorkflowExecutionUpdateRequest,
+        PollWorkflowExecutionUpdateResponse,
+        |r| {
+            let labels = AttachMetricLabels::namespace(r.get_ref().namespace.clone());
+            r.extensions_mut().insert(labels);
+        }
+    );
     (
         start_batch_operation,
         StartBatchOperationRequest,

core/Cargo.toml

@@ -28,6 +28,7 @@ dashmap = "5.0"
 derive_builder = "0.12"
 derive_more = "0.99"
 enum_dispatch = "0.3"
+enum-iterator = "1.4"
 flate2 = "1.0"
 futures = "0.3"
 futures-util = "0.3"

core/src/core_tests/workflow_tasks.rs

@@ -1,5 +1,7 @@
 use crate::{
-    advance_fut, job_assert,
+    advance_fut,
+    internal_flags::CoreInternalFlags,
+    job_assert,
     replay::TestHistoryBuilder,
     test_help::{
         build_fake_worker, build_mock_pollers, build_multihist_mock_sg, canned_histories,
@@ -14,9 +16,9 @@ use crate::{
 use futures::{stream, FutureExt};
 use rstest::{fixture, rstest};
 use std::{
-    collections::{HashMap, VecDeque},
+    collections::{HashMap, HashSet, VecDeque},
     sync::{
-        atomic::{AtomicBool, AtomicU64, Ordering},
+        atomic::{AtomicU64, Ordering},
         Arc,
     },
     time::Duration,
@@ -54,9 +56,7 @@ use temporal_sdk_core_protos::{
     },
     DEFAULT_ACTIVITY_TYPE, DEFAULT_WORKFLOW_TYPE,
 };
-use temporal_sdk_core_test_utils::{
-    fanout_tasks, schedule_activity_cmd, start_timer_cmd, WorkerTestHelpers,
-};
+use temporal_sdk_core_test_utils::{fanout_tasks, start_timer_cmd, WorkerTestHelpers};
 use tokio::{
     join,
     sync::{Barrier, Semaphore},
@@ -2380,54 +2380,36 @@ async fn lang_internal_flags() {
     core.shutdown().await;
 }
 
-// Verify we send flags to server when they're used
+// Verify we send all core internal flags on the first non-replay WFT
 #[tokio::test]
 async fn core_internal_flags() {
     let mut t = TestHistoryBuilder::default();
     t.add_by_type(EventType::WorkflowExecutionStarted);
-    t.add_full_wf_task();
-    let act_scheduled_event_id = t.add_activity_task_scheduled("act-id");
-    let act_started_event_id = t.add_activity_task_started(act_scheduled_event_id);
-    t.add_activity_task_completed(
-        act_scheduled_event_id,
-        act_started_event_id,
-        Default::default(),
-    );
-    t.add_full_wf_task();
-    t.add_workflow_execution_completed();
+    t.add_workflow_task_scheduled_and_started();
 
     let mut mh = MockPollCfg::from_resp_batches(
         "fake_wf_id",
         t,
-        [ResponseType::ToTaskNum(1), ResponseType::ToTaskNum(2)],
+        [ResponseType::ToTaskNum(1)],
         mock_workflow_client(),
     );
-    let first_poll = AtomicBool::new(true);
     mh.completion_asserts = Some(Box::new(move |c| {
-        if !first_poll.load(Ordering::Acquire) {
-            assert_matches!(c.sdk_metadata.core_used_flags.as_slice(), &[1]);
-        }
-        first_poll.store(false, Ordering::Release);
+        assert_eq!(
+            c.sdk_metadata
+                .core_used_flags
+                .iter()
+                .copied()
+                .collect::<HashSet<_>>(),
+            CoreInternalFlags::all_except_too_high()
+                .into_iter()
+                .map(|f| f as u32)
+                .collect()
+        );
     }));
     let mut mock = build_mock_pollers(mh);
     mock.worker_cfg(|wc| wc.max_cached_workflows = 1);
     let core = mock_worker(mock);
 
-    let act = core.poll_workflow_activation().await.unwrap();
-    core.complete_workflow_activation(WorkflowActivationCompletion::from_cmd(
-        act.run_id,
-        schedule_activity_cmd(
-            1,
-            "whatever",
-            "act-id",
-            ActivityCancellationType::TryCancel,
-            Duration::from_secs(60),
-            Duration::from_secs(60),
-        ),
-    ))
-    .await
-    .unwrap();
-
     let act = core.poll_workflow_activation().await.unwrap();
     core.complete_execution(&act.run_id).await;
     core.shutdown().await;

core/src/internal_flags.rs

@@ -1,7 +1,11 @@
 //! Utilities for and tracking of internal versions which alter history in incompatible ways
 //! so that we can use older code paths for workflows executed on older core versions.
 
-use std::collections::{BTreeSet, HashSet};
+use itertools::Either;
+use std::{
+    collections::{BTreeSet, HashSet},
+    iter,
+};
 use temporal_sdk_core_protos::temporal::api::{
     history::v1::WorkflowTaskCompletedEventAttributes, sdk::v1::WorkflowTaskCompletedMetadata,
     workflowservice::v1::get_system_info_response,
@@ -15,7 +19,7 @@ use temporal_sdk_core_protos::temporal::api::{
 /// that removing older variants does not create any change in existing values. Removed flag
 /// variants must be reserved forever (a-la protobuf), and should be called out in a comment.
 #[repr(u32)]
-#[derive(PartialEq, Eq, PartialOrd, Ord, Hash, Copy, Clone, Debug)]
+#[derive(PartialEq, Eq, PartialOrd, Ord, Hash, Copy, Clone, Debug, enum_iterator::Sequence)]
 pub(crate) enum CoreInternalFlags {
     /// In this flag additional checks were added to a number of state machines to ensure that
     /// the ID and type of activities, local activities, and child workflows match during replay.
@@ -28,83 +32,125 @@ pub(crate) enum CoreInternalFlags {
 }
 
 #[derive(Debug, Clone, PartialEq, Eq)]
-pub(crate) struct InternalFlags {
-    enabled: bool,
-    core: BTreeSet<CoreInternalFlags>,
-    lang: BTreeSet<u32>,
-    core_since_last_complete: HashSet<CoreInternalFlags>,
-    lang_since_last_complete: HashSet<u32>,
+pub(crate) enum InternalFlags {
+    Enabled {
+        core: BTreeSet<CoreInternalFlags>,
+        lang: BTreeSet<u32>,
+        core_since_last_complete: HashSet<CoreInternalFlags>,
+        lang_since_last_complete: HashSet<u32>,
+    },
+    Disabled,
 }
 
 impl InternalFlags {
     pub fn new(server_capabilities: &get_system_info_response::Capabilities) -> Self {
-        Self {
-            enabled: server_capabilities.sdk_metadata,
-            core: Default::default(),
-            lang: Default::default(),
-            core_since_last_complete: Default::default(),
-            lang_since_last_complete: Default::default(),
+        match server_capabilities.sdk_metadata {
+            true => Self::Enabled {
+                core: Default::default(),
+                lang: Default::default(),
+                core_since_last_complete: Default::default(),
+                lang_since_last_complete: Default::default(),
+            },
+            false => Self::Disabled,
         }
     }
 
     pub fn add_from_complete(&mut self, e: &WorkflowTaskCompletedEventAttributes) {
-        if !self.enabled {
-            return;
-        }
-
-        if let Some(metadata) = e.sdk_metadata.as_ref() {
-            self.core.extend(
-                metadata
-                    .core_used_flags
-                    .iter()
-                    .map(|u| CoreInternalFlags::from_u32(*u)),
-            );
-            self.lang.extend(metadata.lang_used_flags.iter());
+        if let Self::Enabled { core, lang, .. } = self {
+            if let Some(metadata) = e.sdk_metadata.as_ref() {
+                core.extend(
+                    metadata
+                        .core_used_flags
+                        .iter()
+                        .map(|u| CoreInternalFlags::from_u32(*u)),
+                );
+                lang.extend(metadata.lang_used_flags.iter());
+            }
         }
     }
 
     pub fn add_lang_used(&mut self, flags: impl IntoIterator<Item = u32>) {
-        if !self.enabled {
-            return;
+        if let Self::Enabled {
+            lang_since_last_complete,
+            ..
+        } = self
+        {
+            lang_since_last_complete.extend(flags.into_iter());
         }
-
-        self.lang_since_last_complete.extend(flags.into_iter());
     }
 
     /// Returns true if this flag may currently be used. If `should_record` is true, always returns
     /// true and records the flag as being used, for taking later via
     /// [Self::gather_for_wft_complete].
     pub fn try_use(&mut self, core_patch: CoreInternalFlags, should_record: bool) -> bool {
-        if !self.enabled {
+        match self {
+            Self::Enabled {
+                core,
+                core_since_last_complete,
+                ..
+            } => {
+                if should_record {
+                    core_since_last_complete.insert(core_patch);
+                    true
+                } else {
+                    core.contains(&core_patch)
+                }
+            }
             // If the server does not support the metadata field, we must assume we can never use
             // any internal flags since they can't be recorded for future use
-            return false;
+            Self::Disabled => false,
         }
+    }
 
-        if should_record {
-            self.core_since_last_complete.insert(core_patch);
-            true
-        } else {
-            self.core.contains(&core_patch)
+    /// Writes all known core flags to the set which should be recorded in the current WFT if not
+    /// already known. Must only be called if not replaying.
+    pub fn write_all_known(&mut self) {
+        if let Self::Enabled {
+            core_since_last_complete,
+            ..
+        } = self
+        {
+            core_since_last_complete.extend(CoreInternalFlags::all_except_too_high());
         }
     }
 
     /// Wipes the recorded flags used during the current WFT and returns a partially filled
     /// sdk metadata message that can be combined with any existing data before sending the WFT
     /// complete
     pub fn gather_for_wft_complete(&mut self) -> WorkflowTaskCompletedMetadata {
-        WorkflowTaskCompletedMetadata {
-            core_used_flags: self
-                .core_since_last_complete
-                .drain()
-                .map(|p| p as u32)
-                .collect(),
-            lang_used_flags: self.lang_since_last_complete.drain().collect(),
+        match self {
+            Self::Enabled {
+                core_since_last_complete,
+                lang_since_last_complete,
+                core,
+                lang,
+            } => {
+                let core_newly_used: Vec<_> = core_since_last_complete
+                    .iter()
+                    .filter(|f| !core.contains(f))
+                    .map(|p| *p as u32)
+                    .collect();
+                let lang_newly_used: Vec<_> = lang_since_last_complete
+                    .iter()
+                    .filter(|f| !lang.contains(f))
+                    .copied()
+                    .collect();
+                core.extend(core_since_last_complete.iter());
+                lang.extend(lang_since_last_complete.iter());
+                WorkflowTaskCompletedMetadata {
+                    core_used_flags: core_newly_used,
+                    lang_used_flags: lang_newly_used,
+                }
+            }
+            Self::Disabled => WorkflowTaskCompletedMetadata::default(),
         }
     }
 
-    pub fn all_lang(&self) -> &BTreeSet<u32> {
-        &self.lang
+    pub fn all_lang(&self) -> impl Iterator<Item = u32> + '_ {
+        match self {
+            Self::Enabled { lang, .. } => Either::Left(lang.iter().copied()),
+            Self::Disabled => Either::Right(iter::empty()),
+        }
     }
 }
 
@@ -116,6 +162,11 @@ impl CoreInternalFlags {
             _ => Self::TooHigh,
         }
     }
+
+    pub fn all_except_too_high() -> impl Iterator<Item = CoreInternalFlags> {
+        enum_iterator::all::<CoreInternalFlags>()
+            .filter(|f| !matches!(f, CoreInternalFlags::TooHigh))
+    }
 }
 
 #[cfg(test)]
@@ -138,4 +189,39 @@ mod tests {
         assert_matches!(gathered.core_used_flags.as_slice(), &[]);
         assert_matches!(gathered.lang_used_flags.as_slice(), &[]);
     }
+
+    #[test]
+    fn all_have_u32_from_impl() {
+        let all_known = CoreInternalFlags::all_except_too_high();
+        for flag in all_known {
+            let as_u32 = flag as u32;
+            assert_eq!(CoreInternalFlags::from_u32(as_u32), flag);
+        }
+    }
+
+    #[test]
+    fn only_writes_new_flags() {
+        let mut f = InternalFlags::new(&Capabilities {
+            sdk_metadata: true,
+            ..Default::default()
+        });
+        f.add_lang_used([1]);
+        f.try_use(CoreInternalFlags::IdAndTypeDeterminismChecks, true);
+        let gathered = f.gather_for_wft_complete();
+        assert_matches!(gathered.core_used_flags.as_slice(), &[1]);
+        assert_matches!(gathered.lang_used_flags.as_slice(), &[1]);
+
+        f.add_from_complete(&WorkflowTaskCompletedEventAttributes {
+            sdk_metadata: Some(WorkflowTaskCompletedMetadata {
+                core_used_flags: vec![2],
+                lang_used_flags: vec![2],
+            }),
+            ..Default::default()
+        });
+        f.add_lang_used([2]);
+        f.try_use(CoreInternalFlags::UpsertSearchAttributeOnPatch, true);
+        let gathered = f.gather_for_wft_complete();
+        assert_matches!(gathered.core_used_flags.as_slice(), &[]);
+        assert_matches!(gathered.lang_used_flags.as_slice(), &[]);
+    }
 }