Cleanup node setup in test and add version control to stacks client

Signed-off-by: Jacinta Ferrant <jacinta@trustmachines.co>

stacks-signer/src/config.rs

@@ -154,9 +154,7 @@ impl TryFrom<RawConfigFile> for Config {
             .endpoint
             .clone()
             .to_socket_addrs()
-            .map_err(|_| {
-                ConfigError::BadField("endpoint".to_string(), raw_data.endpoint.clone())
-            })?
+            .map_err(|_| ConfigError::BadField("endpoint".to_string(), raw_data.endpoint.clone()))?
             .next()
             .ok_or(ConfigError::BadField(
                 "endpoint".to_string(),

stacks-signer/src/runloop.rs

@@ -5,8 +5,8 @@ use frost_signer::{
 };
 use libsigner::{SignerRunLoop, StackerDBChunksEvent};
 use p256k1::ecdsa;
-use slog::{slog_debug, slog_info, slog_warn};
-use stacks_common::{debug, info, warn};
+use slog::{slog_info, slog_warn};
+use stacks_common::{info, warn};
 use std::time::Duration;
 use wsts::Point;
 

stacks-signer/src/stacks_client.rs

@@ -1,9 +1,10 @@
 use bincode::Error as BincodeError;
 use frost_signer::{net::Message, signing_round::MessageTypes};
+use hashbrown::HashMap;
 use libsigner::{RPCError, SignerSession, StackerDBSession};
 use libstackerdb::{Error as StackerDBError, StackerDBChunkAckData, StackerDBChunkData};
-use slog::slog_debug;
-use stacks_common::{debug, types::chainstate::StacksPrivateKey};
+use slog::{slog_debug, slog_warn};
+use stacks_common::{debug, types::chainstate::StacksPrivateKey, warn};
 
 use crate::config::Config;
 
@@ -21,6 +22,9 @@ pub enum ClientError {
     /// Failed to write to stacker-db due to RPC error
     #[error("Failed to write to stacker-db instance: {0}")]
     PutChunkFailed(#[from] RPCError),
+    /// Stacker-db instance rejected the chunk
+    #[error("Stacker-db rejected the chunk. Reason: {0}")]
+    PutChunkRejected(String),
 }
 
 /// TODO: Add stacks node communication to this
@@ -30,6 +34,8 @@ pub struct StacksClient {
     stackerdb_session: StackerDBSession,
     /// The private key used in all stacks node communications
     stacks_private_key: StacksPrivateKey,
+    /// A map of a slot ID to last chunk version
+    slot_versions: HashMap<u32, u32>,
 }
 
 impl From<&Config> for StacksClient {
@@ -40,6 +46,7 @@ impl From<&Config> for StacksClient {
                 config.stackerdb_contract_id.clone(),
             ),
             stacks_private_key: config.stacks_private_key,
+            slot_versions: HashMap::new(),
         }
     }
 }
@@ -52,12 +59,29 @@ impl StacksClient {
         message: Message,
     ) -> Result<StackerDBChunkAckData, ClientError> {
         let message_bytes = bincode::serialize(&message)?;
+        let slot_id = slot_id(id, &message.msg);
 
-        let mut chunk = StackerDBChunkData::new(slot_id(id, &message.msg), 1, message_bytes);
-        chunk.sign(&self.stacks_private_key)?;
-        let chunk_ack = self.stackerdb_session.put_chunk(chunk)?;
-        debug!("{:?}", chunk_ack.clone());
-        Ok(chunk_ack)
+        loop {
+            let slot_version = *self.slot_versions.entry(slot_id).or_insert(0) + 1;
+            let mut chunk = StackerDBChunkData::new(slot_id, slot_version, message_bytes.clone());
+            chunk.sign(&self.stacks_private_key)?;
+            debug!("Sending a chunk to stackerdb!\n{:?}", chunk.clone());
+            let chunk_ack = self.stackerdb_session.put_chunk(chunk)?;
+            self.slot_versions.insert(slot_id, slot_version);
+            if chunk_ack.accepted {
+                debug!("Chunk accepted by stackerdb! ACK: {:?}", chunk_ack);
+                return Ok(chunk_ack);
+            }
+            if let Some(reason) = chunk_ack.reason {
+                // TODO: fix this jankiness. Update stackerdb to use an error code mapping instead of just a string
+                if reason == "Data for this slot and version already exist" {
+                    warn!("Failed to send message to stackerdb due to wrong version number {}. Incrementing and retrying...", slot_version);
+                } else {
+                    warn!("Failed to send message to stackerdb: {}", reason);
+                    return Err(ClientError::PutChunkRejected(reason));
+                }
+            }
+        }
     }
 
     /// Retrieve the total number of slots allocated to a stacker-db writer

testnet/stacks-node/src/tests/signer.rs

@@ -3,7 +3,7 @@ use std::time::Duration;
 use std::{env, thread};
 
 use crate::{
-    config::{EventKeyType, EventObserverConfig, InitialBalance},
+    config::{Config as NeonConfig, EventKeyType, EventObserverConfig, InitialBalance},
     neon,
     tests::{
         bitcoin_regtest::BitcoinCoreController,
@@ -25,6 +25,15 @@ use stacks_signer::runloop::RunLoopCommand;
 
 const SLOTS_PER_USER: u32 = 16;
 
+// Helper struct for holding the btc and stx neon nodes
+#[allow(dead_code)]
+struct RunningNodes {
+    pub btc_regtest_controller: BitcoinRegtestController,
+    pub btcd_controller: BitcoinCoreController,
+    pub join_handle: thread::JoinHandle<()>,
+    pub conf: NeonConfig,
+}
+
 fn build_contract(num_signers: u32, signer_stacks_private_keys: &[StacksPrivateKey]) -> String {
     let mut stackerdb_contract = String::new(); // "
     stackerdb_contract += "        ;; stacker DB\n";
@@ -141,18 +150,11 @@ fn spawn_running_signer(
     signer.spawn(endpoint).unwrap()
 }
 
-#[test]
-fn test_stackerdb_dkg() {
-    let num_signers: u32 = 5;
-    let signer_stacks_private_keys = (0..num_signers)
-        .map(|_| StacksPrivateKey::new())
-        .collect::<Vec<StacksPrivateKey>>();
-
-    // Setup the neon node
-    if env::var("BITCOIND_TEST") != Ok("1".into()) {
-        return;
-    }
-
+fn setup_stx_btc_node(
+    num_signers: u32,
+    signer_stacks_private_keys: &[StacksPrivateKey],
+    stackerdb_contract: &str,
+) -> RunningNodes {
     let (mut conf, _) = neon_integration_test_conf();
     conf.events_observers.push(EventObserverConfig {
         endpoint: format!("localhost:{}", test_observer::EVENT_OBSERVER_PORT),
@@ -190,7 +192,7 @@ fn test_stackerdb_dkg() {
     let mut run_loop = neon::RunLoop::new(conf.clone());
     let blocks_processed = run_loop.get_blocks_processed_arc();
 
-    thread::spawn(move || run_loop.start(None, 0));
+    let join_handle = thread::spawn(move || run_loop.start(None, 0));
 
     // Give the run loop some time to start up!
     eprintln!("Wait for runloop...");
@@ -209,14 +211,13 @@ fn test_stackerdb_dkg() {
     next_block_and_wait(&mut btc_regtest_controller, &blocks_processed);
 
     let http_origin = format!("http://{}", &conf.node.rpc_bind);
-    let stackerdb_contract = build_contract(num_signers, &signer_stacks_private_keys);
     eprintln!("Send contract-publish...");
     let tx = make_contract_publish(
         &signer_stacks_private_keys[0],
         0,
         10_000,
         "hello-world",
-        &stackerdb_contract,
+        stackerdb_contract,
     );
     submit_tx(&http_origin, &tx);
 
@@ -225,12 +226,41 @@ fn test_stackerdb_dkg() {
     next_block_and_wait(&mut btc_regtest_controller, &blocks_processed);
     next_block_and_wait(&mut btc_regtest_controller, &blocks_processed);
 
+    RunningNodes {
+        btcd_controller,
+        btc_regtest_controller,
+        join_handle,
+        conf,
+    }
+}
+
+#[test]
+fn test_stackerdb_dkg() {
+    let num_signers: u32 = 5;
+    let signer_stacks_private_keys = (0..num_signers)
+        .map(|_| StacksPrivateKey::new())
+        .collect::<Vec<StacksPrivateKey>>();
+
+    // Setup the neon node
+    if env::var("BITCOIND_TEST") != Ok("1".into()) {
+        return;
+    }
+
+    // Build the stackerdb contract
+    let stackerdb_contract = build_contract(num_signers, &signer_stacks_private_keys);
+    // Setup the nodes and deploy the contract to it
+    let node = setup_stx_btc_node(
+        num_signers,
+        &signer_stacks_private_keys,
+        &stackerdb_contract,
+    );
+
     // Setup the signer and coordinator configurations
-    let contract_id = conf.node.stacker_dbs[0].clone();
+    let contract_id = node.conf.node.stacker_dbs[0].clone();
     let signer_configs = build_signer_config_tomls(
         num_signers,
         &signer_stacks_private_keys,
-        &conf.node.rpc_bind,
+        &node.conf.node.rpc_bind,
         contract_id.to_string(),
     );