feat: add versioning for request messages

lgn-messages/src/types/mod.rs

@@ -16,8 +16,8 @@ use v1::query::tasks::ProofInputKind;
 use v1::query::tasks::QueryStep;
 use v1::query::tasks::RevelationInput;
 
-use crate::routing::RoutingKey;
 use crate::KeyedPayload;
+use crate::routing::RoutingKey;
 
 pub mod v1;
 
@@ -35,6 +35,148 @@ pub enum ReplyType {
     V1Groth16(WorkerReply),
 }
 
+#[derive(Debug, Clone, Deserialize, Serialize)]
+#[serde(tag = "version")]
+pub enum RequestVersioned {
+    /// Version 1 of the request.
+    #[serde(rename = "1")]
+    V1(V1),
+
+    /// Variant for backwards compatibility.
+    ///
+    /// TODO: Remove this once phased out.
+    #[serde(untagged)]
+    Legacy(MessageEnvelope),
+}
+
+impl RequestVersioned {
+    /// Returns a reference to the requested mp2 library version.
+    pub fn mp2_version(&self) -> &str {
+        match self {
+            RequestVersioned::V1(v1) => v1.mp2_version(),
+            RequestVersioned::Legacy(message_envelope) => message_envelope.mp2_version(),
+        }
+    }
+
+    /// Returns a unique identifier for this task.
+    ///
+    /// Note: For legacy messages, the unique identifier is created by adjoining the task and query
+    /// ids.
+    pub fn id(&self) -> String {
+        match self {
+            RequestVersioned::V1(v1) => v1.id(),
+            RequestVersioned::Legacy(message_envelope) => message_envelope.id(),
+        }
+    }
+
+    /// Returns a reference to this task's id.
+    pub fn task_id(&self) -> &str {
+        match self {
+            RequestVersioned::V1(v1) => v1.task_id(),
+            RequestVersioned::Legacy(message_envelope) => message_envelope.task_id(),
+        }
+    }
+
+    /// Returns [ProverType] which supports proving this [TaskType].
+    ///
+    /// This is used to dispatch the message to the correct underlying prover.
+    pub fn to_prover_type(&self) -> ProverType {
+        match self {
+            RequestVersioned::V1(v1) => v1.to_prover_type(),
+            RequestVersioned::Legacy(message_envelope) => message_envelope.to_prover_type(),
+        }
+    }
+
+    /// Returns the task's type name.
+    ///
+    /// This is used to classify the tasks for metrics.
+    pub fn to_task_type(&self) -> &str {
+        match self {
+            RequestVersioned::V1(v1) => v1.to_task_type(),
+            RequestVersioned::Legacy(message_envelope) => message_envelope.to_task_type(),
+        }
+    }
+
+    /// Returns the inner [TaskType].
+    pub fn inner(&self) -> &TaskType {
+        match self {
+            RequestVersioned::V1(v1) => v1.inner(),
+            RequestVersioned::Legacy(message_envelope) => message_envelope.inner(),
+        }
+    }
+
+    /// Consumes self and returns the inner [TaskType].
+    pub fn into_inner(self) -> TaskType {
+        match self {
+            RequestVersioned::V1(v1) => v1.into_inner(),
+            RequestVersioned::Legacy(message_envelope) => message_envelope.into_inner(),
+        }
+    }
+
+    /// Returns the `query_id` part of the message.
+    ///
+    /// Note: This is only available for legacy messages.
+    pub fn query_id(&self) -> &str {
+        match self {
+            RequestVersioned::V1(_v1) => "",
+            RequestVersioned::Legacy(message_envelope) => message_envelope.query_id(),
+        }
+    }
+}
+
+#[derive(Debug, Clone, Deserialize, Serialize)]
+pub struct V1 {
+    /// Task id is unique for each task and helps to map replies to tasks.
+    pub task_id: String,
+
+    /// Details of the task to be executed.
+    pub inner: TaskType,
+
+    /// The proving system version.
+    pub mp2_version: String,
+}
+
+impl V1 {
+    /// Returns a reference to the requested mp2 library version.
+    pub fn mp2_version(&self) -> &str {
+        &self.mp2_version
+    }
+
+    /// Returns a copy of this task's id.
+    pub fn id(&self) -> String {
+        self.task_id.to_owned()
+    }
+
+    /// Returns a reference to this task's id.
+    pub fn task_id(&self) -> &str {
+        &self.task_id
+    }
+
+    /// Returns [ProverType] which supports proving this [TaskType].
+    ///
+    /// This is used to dispatch the message to the correct underlying prover.
+    pub fn to_prover_type(&self) -> ProverType {
+        self.inner.to_prover_type()
+    }
+
+    /// Returns the task's type name.
+    ///
+    /// This is used to classify the tasks for metrics.
+    pub fn to_task_type(&self) -> &str {
+        self.inner.to_task_type()
+    }
+
+    /// Returns the inner [TaskType].
+    pub fn inner(&self) -> &TaskType {
+        &self.inner
+    }
+
+    /// Consumes self and returns the inner [TaskType].
+    fn into_inner(self) -> TaskType {
+        self.inner
+    }
+}
+
 #[derive(Clone, Deserialize, Serialize)]
 pub struct MessageEnvelope {
     /// Query id is unique for each query and shared between all its tasks
@@ -113,6 +255,7 @@ impl MessageEnvelope {
         format!("{}-{}", self.query_id, self.task_id)
     }
 
+    /// Returns the inner [TaskType].
     pub fn inner(&self) -> &TaskType {
         &self.inner
     }
@@ -128,6 +271,11 @@ impl MessageEnvelope {
         self.inner.to_prover_type()
     }
 
+    /// Returns a reference to the requested mp2 library version.
+    pub fn mp2_version(&self) -> &str {
+        &self.version
+    }
+
     /// Returns the task's type name.
     ///
     /// This is used to classify the tasks for metrics.
@@ -210,6 +358,12 @@ impl MessageEnvelope {
             },
             TaskType::V1Groth16(..) => "groth16",
         }
+        self.inner.to_task_type()
+    }
+
+    /// Consumes self and returns the inner [TaskType].
+    fn into_inner(self) -> TaskType {
+        self.inner
     }
 }
 
@@ -391,4 +545,87 @@ impl TaskType {
             TaskType::V1Groth16(_) => ProverType::V1Groth16,
         }
     }
+
+    /// Returns the task's type name.
+    ///
+    /// This is used to classify the tasks for metrics.
+    fn to_task_type(&self) -> &str {
+        match self {
+            TaskType::V1Preprocessing(worker_task) => {
+                match &worker_task.task_type {
+                    v1::preprocessing::WorkerTaskType::Extraction(extraction_type) => {
+                        match extraction_type {
+                            ExtractionType::MptExtraction(mpt) => {
+                                match mpt.mpt_type {
+                                    MptType::MappingLeaf(..) => "mapping_leaf",
+                                    MptType::MappingBranch(..) => "mapping_branch",
+                                    MptType::VariableLeaf(..) => "multi_var_leaf",
+                                    MptType::VariableBranch(..) => "multi_var_branch",
+                                }
+                            },
+                            ExtractionType::LengthExtraction(..) => "length",
+                            ExtractionType::ContractExtraction(..) => "contract",
+                            ExtractionType::BlockExtraction(..) => "block",
+                            ExtractionType::FinalExtraction(final_extraction) => {
+                                match &**final_extraction {
+                                    FinalExtraction::Single(single_table_extraction) => {
+                                        match single_table_extraction.extraction_type {
+                                            FinalExtractionType::Simple(..) => "final_extraction",
+                                            FinalExtractionType::Lengthed => {
+                                                "final_extraction_lengthed"
+                                            },
+                                        }
+                                    },
+                                    FinalExtraction::Merge(..) => "final_extraction_merge",
+                                }
+                            },
+                        }
+                    },
+                    v1::preprocessing::WorkerTaskType::Database(database_type) => {
+                        match database_type {
+                            DatabaseType::Cell(db_cell_type) => {
+                                match db_cell_type {
+                                    DbCellType::Leaf(..) => "cell_leaf",
+                                    DbCellType::Partial(..) => "cell_partial",
+                                    DbCellType::Full(..) => "cell_full",
+                                }
+                            },
+                            DatabaseType::Row(db_row_type) => {
+                                match db_row_type {
+                                    DbRowType::Leaf(..) => "row_leaf",
+                                    DbRowType::Partial(..) => "row_partial",
+                                    DbRowType::Full(..) => "row_full",
+                                }
+                            },
+                            DatabaseType::Index(..) => "index",
+                            DatabaseType::IVC(..) => "ivc",
+                        }
+                    },
+                }
+            },
+            TaskType::V1Query(worker_task) => {
+                match &worker_task.task_type {
+                    v1::query::WorkerTaskType::Query(query_input) => {
+                        match &query_input.query_step {
+                            QueryStep::Tabular(..) => "tabular",
+                            QueryStep::Aggregation(aggregation_input) => {
+                                match aggregation_input.input_kind {
+                                    ProofInputKind::RowsChunk(..) => "rows_chunk",
+                                    ProofInputKind::ChunkAggregation(..) => "chunk_aggregation",
+                                    ProofInputKind::NonExistence(..) => "non_existence",
+                                }
+                            },
+                            QueryStep::Revelation(revelation_input) => {
+                                match revelation_input {
+                                    RevelationInput::Aggregated { .. } => "revelation_aggregated",
+                                    RevelationInput::Tabular { .. } => "revelation_tabular",
+                                }
+                            },
+                        }
+                    },
+                }
+            },
+            TaskType::V1Groth16(..) => "groth16",
+        }
+    }
 }

lgn-provers/src/provers/mod.rs

@@ -1,5 +1,5 @@
-use lgn_messages::types::MessageEnvelope;
 use lgn_messages::types::MessageReplyEnvelope;
+use lgn_messages::types::RequestVersioned;
 
 pub mod v1;
 
@@ -10,6 +10,6 @@ pub mod v1;
 pub trait LgnProver {
     fn run(
         &self,
-        envelope: MessageEnvelope,
+        envelope: RequestVersioned,
     ) -> anyhow::Result<MessageReplyEnvelope>;
 }

lgn-provers/src/provers/v1/groth16/dummy_prover.rs

@@ -1,9 +1,9 @@
 use anyhow::bail;
 use lgn_messages::types::v1::groth16::keys::ProofKey;
-use lgn_messages::types::MessageEnvelope;
 use lgn_messages::types::MessageReplyEnvelope;
 use lgn_messages::types::ProofCategory;
 use lgn_messages::types::ReplyType;
+use lgn_messages::types::RequestVersioned;
 use lgn_messages::types::TaskType;
 use lgn_messages::types::WorkerReply;
 
@@ -18,17 +18,18 @@ pub struct Groth16DummyProver;
 impl LgnProver for Groth16DummyProver {
     fn run(
         &self,
-        envelope: MessageEnvelope,
+        envelope: RequestVersioned,
     ) -> anyhow::Result<MessageReplyEnvelope> {
-        let query_id = envelope.query_id.clone();
-        let task_id = envelope.task_id.clone();
+        let query_id = envelope.query_id();
+        let task_id = envelope.task_id();
         if let TaskType::V1Groth16(task) = envelope.inner() {
             let key = ProofKey(query_id.to_string()).to_string();
             let proof = dummy_proof(PROOF_SIZE);
             let reply =
                 WorkerReply::new(task.chain_id, Some((key, proof)), ProofCategory::Querying);
             let reply_type = ReplyType::V1Groth16(reply);
-            let reply_envelope = MessageReplyEnvelope::new(query_id, task_id, reply_type);
+            let reply_envelope =
+                MessageReplyEnvelope::new(query_id.to_owned(), task_id.to_owned(), reply_type);
             Ok(reply_envelope)
         } else {
             bail!("Unexpected task: {:?}", envelope);