Added optional wrapping of message types in common combined message type (#82)

* Added keys to slot and transaction messages

Signed-off-by: Anthony Milbourne <18662115+amilbourne@users.noreply.github.com>

* Added optional wrapping of message types in common combined message type

Signed-off-by: Anthony Milbourne <18662115+amilbourne@users.noreply.github.com>

* Added documentation for message keys, address filtering, and message wrapping

Signed-off-by: Anthony Milbourne <18662115+amilbourne@users.noreply.github.com>

---------

Signed-off-by: Anthony Milbourne <18662115+amilbourne@users.noreply.github.com>

Author: amilbourne

README.md

@@ -43,6 +43,7 @@ Config is specified via the plugin's JSON config file.
   "slot_status_topic": "solana.testnet.slot_status",
   "transaction_topic": "solana.testnet.transactions",
   "publish_all_accounts": false,
+  "wrap_messages": false,
   "program_ignores": [
     "Sysvar1111111111111111111111111111111111111",
     "Vote111111111111111111111111111111111111111"
@@ -59,7 +60,32 @@ Config is specified via the plugin's JSON config file.
 - `update_account_topic`: Topic name of account updates. Omit to disable.
 - `slot_status_topic`: Topic name of slot status update. Omit to disable.
 - `publish_all_accounts`: Publish all accounts on startup. Omit to disable.
-- `program_ignores`: Solana program IDs for which to ignore updates for owned accounts.
+- `wrap_messages`: Wrap all messages in a unified wrapper object. Omit to disable (see Message Wrapping below).
+- `program_ignores`: Account addresses to ignore (see Filtering below).
+
+### Message Keys
+
+The message types are keyed as follows:
+- **Account update:** account address (public key)
+- **Slot status:** slot number
+- **Transaction notification:** transaction signature
+
+### Filtering
+
+If `program_ignores` are specified, then these addresses will be filtered out of the account updates
+and transaction notifications.  More specifically, account update messages for these accounts will not be emitted,
+and transaction notifications for any transaction involving these accounts will not be emitted.
+
+### Message Wrapping
+
+In some cases it may be desirable to send multiple types of messages to the same topic,
+for instance to preserve relative order.  In this case it is helpful if all messages conform to a single schema.
+Setting `wrap_messages` to true will wrap all three message types in a uniform wrapper object so that they
+conform to a single schema.
+
+Note that if `wrap_messages` is true, in order to avoid key collision, the message keys are prefixed with a single byte,
+which is dependent on the type of the message being wrapped.  Account update message keys are prefixed with
+65 (A), slot status keys with 83 (S), and transaction keys with 84 (T).
 
 ## Buffering
 

build.rs

@@ -1,6 +1,8 @@
 use std::io::Result;
 
 fn main() -> Result<()> {
-    prost_build::compile_protos(&["proto/event.proto"], &["proto/"])?;
+    let mut config = prost_build::Config::new();
+    config.boxed(".blockdaemon.solana.accountsdb_plugin_kafka.types.MessageWrapper");
+    config.compile_protos(&["proto/event.proto"], &["proto/"])?;
     Ok(())
 }

proto/event.proto

@@ -172,3 +172,11 @@ message TransactionEvent {
   TransactionStatusMeta transaction_status_meta = 4;
   uint64 slot = 5;
 }
+
+message MessageWrapper {
+  oneof event_message {
+    UpdateAccountEvent account = 1;
+    SlotStatusEvent slot = 2;
+    TransactionEvent transaction = 3;
+  }
+}

src/config.rs

@@ -49,6 +49,9 @@ pub struct Config {
     /// Publish all accounts on startup.
     #[serde(default)]
     pub publish_all_accounts: bool,
+    /// Wrap all event message in a single message type.
+    #[serde(default)]
+    pub wrap_messages: bool,
 }
 
 impl Default for Config {
@@ -61,6 +64,7 @@ impl Default for Config {
             transaction_topic: "".to_owned(),
             program_ignores: Vec::new(),
             publish_all_accounts: false,
+            wrap_messages: false,
         }
     }
 }

src/publisher.rs

@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use crate::message_wrapper::EventMessage;
+use crate::message_wrapper::EventMessage::{Account, Slot, Transaction};
 use {
     crate::*,
     prost::Message,
@@ -29,6 +31,8 @@ pub struct Publisher {
     update_account_topic: String,
     slot_status_topic: String,
     transaction_topic: String,
+
+    wrap_messages: bool,
 }
 
 impl Publisher {
@@ -39,26 +43,53 @@ impl Publisher {
             update_account_topic: config.update_account_topic.clone(),
             slot_status_topic: config.slot_status_topic.clone(),
             transaction_topic: config.transaction_topic.clone(),
+            wrap_messages: config.wrap_messages,
         }
     }
 
     pub fn update_account(&self, ev: UpdateAccountEvent) -> Result<(), KafkaError> {
-        let buf = ev.encode_to_vec();
+        let temp_key;
+        let (key, buf) = if self.wrap_messages {
+            temp_key = self.copy_and_prepend(ev.pubkey.as_slice(), 65u8);
+            (&temp_key, Self::encode_with_wrapper(Account(Box::new(ev))))
+        } else {
+            (&ev.pubkey, ev.encode_to_vec())
+        };
         let record = BaseRecord::<Vec<u8>, _>::to(&self.update_account_topic)
-            .key(&ev.pubkey)
+            .key(key)
             .payload(&buf);
         self.producer.send(record).map(|_| ()).map_err(|(e, _)| e)
     }
 
     pub fn update_slot_status(&self, ev: SlotStatusEvent) -> Result<(), KafkaError> {
-        let buf = ev.encode_to_vec();
-        let record = BaseRecord::<(), _>::to(&self.slot_status_topic).payload(&buf);
+        let temp_key;
+        let (key, buf) = if self.wrap_messages {
+            temp_key = self.copy_and_prepend(&ev.slot.to_le_bytes(), 83u8);
+            (&temp_key, Self::encode_with_wrapper(Slot(Box::new(ev))))
+        } else {
+            temp_key = ev.slot.to_le_bytes().to_vec();
+            (&temp_key, ev.encode_to_vec())
+        };
+        let record = BaseRecord::<Vec<u8>, _>::to(&self.slot_status_topic)
+            .key(key)
+            .payload(&buf);
         self.producer.send(record).map(|_| ()).map_err(|(e, _)| e)
     }
 
     pub fn update_transaction(&self, ev: TransactionEvent) -> Result<(), KafkaError> {
-        let buf = ev.encode_to_vec();
-        let record = BaseRecord::<(), _>::to(&self.transaction_topic).payload(&buf);
+        let temp_key;
+        let (key, buf) = if self.wrap_messages {
+            temp_key = self.copy_and_prepend(ev.signature.as_slice(), 84u8);
+            (
+                &temp_key,
+                Self::encode_with_wrapper(Transaction(Box::new(ev))),
+            )
+        } else {
+            (&ev.signature, ev.encode_to_vec())
+        };
+        let record = BaseRecord::<Vec<u8>, _>::to(&self.transaction_topic)
+            .key(key)
+            .payload(&buf);
         self.producer.send(record).map(|_| ()).map_err(|(e, _)| e)
     }
 
@@ -73,6 +104,20 @@ impl Publisher {
     pub fn wants_transaction(&self) -> bool {
         !self.transaction_topic.is_empty()
     }
+
+    fn encode_with_wrapper(message: EventMessage) -> Vec<u8> {
+        MessageWrapper {
+            event_message: Some(message),
+        }
+        .encode_to_vec()
+    }
+
+    fn copy_and_prepend(&self, data: &[u8], prefix: u8) -> Vec<u8> {
+        let mut temp_key = Vec::with_capacity(data.len() + 1);
+        temp_key.push(prefix);
+        temp_key.extend_from_slice(data);
+        temp_key
+    }
 }
 
 impl Drop for Publisher {