Merge pull request #675 from primitivefinance/feat/event-logger-metadata

feat(arbiter-core): event logger metadata

arbiter-core/data/output.json

@@ -0,0 +1 @@
+{"events":{"arbx":{"ApprovalFilter":[{"owner":"0x2efdc9eecfee3a776209fcb8e9a83a6b221d74f5","spender":"0x2efdc9eecfee3a776209fcb8e9a83a6b221d74f5","amount":"0x1"},{"owner":"0x2efdc9eecfee3a776209fcb8e9a83a6b221d74f5","spender":"0x2efdc9eecfee3a776209fcb8e9a83a6b221d74f5","amount":"0x1"},{"owner":"0x2efdc9eecfee3a776209fcb8e9a83a6b221d74f5","spender":"0x2efdc9eecfee3a776209fcb8e9a83a6b221d74f5","amount":"0x1"},{"owner":"0x2efdc9eecfee3a776209fcb8e9a83a6b221d74f5","spender":"0x2efdc9eecfee3a776209fcb8e9a83a6b221d74f5","amount":"0x1"},{"owner":"0x2efdc9eecfee3a776209fcb8e9a83a6b221d74f5","spender":"0x2efdc9eecfee3a776209fcb8e9a83a6b221d74f5","amount":"0x1"}]},"arby":{"ApprovalFilter":[{"owner":"0x2efdc9eecfee3a776209fcb8e9a83a6b221d74f5","spender":"0x2efdc9eecfee3a776209fcb8e9a83a6b221d74f5","amount":"0x1"},{"owner":"0x2efdc9eecfee3a776209fcb8e9a83a6b221d74f5","spender":"0x2efdc9eecfee3a776209fcb8e9a83a6b221d74f5","amount":"0x1"},{"owner":"0x2efdc9eecfee3a776209fcb8e9a83a6b221d74f5","spender":"0x2efdc9eecfee3a776209fcb8e9a83a6b221d74f5","amount":"0x1"},{"owner":"0x2efdc9eecfee3a776209fcb8e9a83a6b221d74f5","spender":"0x2efdc9eecfee3a776209fcb8e9a83a6b221d74f5","amount":"0x1"},{"owner":"0x2efdc9eecfee3a776209fcb8e9a83a6b221d74f5","spender":"0x2efdc9eecfee3a776209fcb8e9a83a6b221d74f5","amount":"0x1"}]},"lex":{"PriceChangeFilter":[{"price":"0xde0b6b3a7640000"},{"price":"0xde0b6b3a7640000"},{"price":"0xde0b6b3a7640000"},{"price":"0xde0b6b3a7640000"},{"price":"0xde0b6b3a7640000"}]}},"metadata":{"name":"test"}}

arbiter-core/src/data_collection.rs

@@ -53,10 +53,11 @@ type FilterDecoder =
 /// * `E` - Type that implements the `EthLogDecode`, `Debug`, `Serialize`
 ///   traits, and has a static lifetime.
 pub struct EventLogger {
-    directory: Option<String>,
-    file_name: Option<String>,
     decoder: FilterDecoder,
     receiver: Option<crossbeam_channel::Receiver<Broadcast>>,
+    directory: Option<String>,
+    file_name: Option<String>,
+    metadata: Option<Value>,
 }
 
 impl EventLogger {
@@ -72,6 +73,7 @@ impl EventLogger {
             file_name: None,
             decoder: BTreeMap::new(),
             receiver: None,
+            metadata: None,
         }
     }
 
@@ -141,6 +143,22 @@ impl EventLogger {
         self
     }
 
+    /// Sets the metadata for the `EventLogger`.
+    ///
+    /// # Arguments
+    ///
+    /// * `metadata` - The metadata to be stored with the event logs which must
+    ///   implement the `Serialize` trait.
+    ///
+    /// # Returns
+    ///
+    /// The `EventLogger` instance with the specified metadata.
+    pub fn metadata(mut self, metadata: impl Serialize) -> Result<Self, serde_json::Error> {
+        let metadata = serde_json::to_value(metadata)?;
+        self.metadata = Some(metadata);
+        Ok(self)
+    }
+
     /// Executes the `EventLogger`.
     ///
     /// This function starts the event logging process. It first deletes the
@@ -165,8 +183,9 @@ impl EventLogger {
         let receiver = self.receiver.unwrap();
         let dir = self.directory.unwrap_or("./data".into());
         let file_name = self.file_name.unwrap_or("output".into());
+        let metadata = self.metadata.clone();
         std::thread::spawn(move || {
-            let mut logs: BTreeMap<String, BTreeMap<String, Vec<Value>>> = BTreeMap::new();
+            let mut events: BTreeMap<String, BTreeMap<String, Vec<Value>>> = BTreeMap::new();
             while let Ok(broadcast) = receiver.recv() {
                 match broadcast {
                     Broadcast::StopSignal => {
@@ -176,7 +195,18 @@ impl EventLogger {
                         let file_path = output_dir.join(format!("{}.json", file_name));
                         let file = std::fs::File::create(file_path).unwrap();
                         let writer = BufWriter::new(file);
-                        serde_json::to_writer(writer, &logs).expect("Unable to write data");
+
+                        // Create a struct to hold both logs and metadata
+                        #[derive(Serialize)]
+                        struct OutputData<T> {
+                            events: BTreeMap<String, BTreeMap<String, Vec<Value>>>,
+                            metadata: Option<T>,
+                        }
+
+                        let data = OutputData { events, metadata };
+
+                        // Write the data to the file
+                        serde_json::to_writer(writer, &data).expect("Unable to write data");
                         break;
                     }
                     Broadcast::Event(event) => {
@@ -191,11 +221,11 @@ impl EventLogger {
                                     .unwrap();
                                     let event_as_object = event_as_value.as_object().unwrap();
 
-                                    let contract = logs.get(contract_name);
+                                    let contract = events.get(contract_name);
                                     if contract.is_none() {
-                                        logs.insert(contract_name.clone(), BTreeMap::new());
+                                        events.insert(contract_name.clone(), BTreeMap::new());
                                     }
-                                    let contract = logs.get_mut(contract_name).unwrap();
+                                    let contract = events.get_mut(contract_name).unwrap();
 
                                     let event_name =
                                         event_as_object.clone().keys().collect::<Vec<&String>>()[0]

arbiter-core/src/tests/data_output.rs

@@ -1,16 +1,29 @@
+use serde::Serialize;
+
 use super::*;
 use crate::data_collection::EventLogger;
 
+#[derive(Serialize, Clone)]
+struct MockMetadata {
+    pub name: String,
+}
+
 #[tokio::test]
 async fn data_capture() {
     let (env, client) = startup_user_controlled().unwrap();
     let (arbx, arby, lex) = deploy_liquid_exchange(client.clone()).await.unwrap();
     println!("Deployed contracts");
 
+    let metadata = MockMetadata {
+        name: "test".to_string(),
+    };
+
     let listener = EventLogger::builder()
         .add(arbx.events(), "arbx")
         .add(arby.events(), "arby")
-        .add(lex.events(), "lex");
+        .add(lex.events(), "lex")
+        .metadata(metadata)
+        .unwrap();
 
     listener.run().unwrap();
 
@@ -37,5 +50,5 @@ async fn data_capture() {
 
     env.stop().unwrap();
     let path = std::env::current_dir().unwrap().join("data");
-    std::fs::remove_dir_all(path).unwrap();
+    // std::fs::remove_dir_all(path).unwrap();
 }