From 76bcf19fa96917be65c2c72ca465f397a66ceb95 Mon Sep 17 00:00:00 2001 From: Jonas Bostoen Date: Thu, 9 Jan 2025 11:53:14 +0100 Subject: [PATCH] refactor(db): some changes, add InMemoryDb --- src/db/memory.rs | 61 ++++++++++++++++++++++++++++++++++++++++++++++++ src/db/mod.rs | 11 +++++---- src/db/noop.rs | 36 ---------------------------- src/db/sql.rs | 27 ++++++++++++++------- 4 files changed, 87 insertions(+), 48 deletions(-) create mode 100644 src/db/memory.rs delete mode 100644 src/db/noop.rs diff --git a/src/db/memory.rs b/src/db/memory.rs new file mode 100644 index 0000000..b17d0f2 --- /dev/null +++ b/src/db/memory.rs @@ -0,0 +1,61 @@ +use std::sync::{Arc, RwLock}; + +use alloy::primitives::Address; +use tracing::info; + +use super::{BlsPublicKey, DbResult, Operator, Registration, RegistryDb}; + +#[derive(Debug, Clone)] +pub(crate) struct InMemoryDb { + validator_registrations: Arc>>, + operator_registrations: Arc>>, +} + +#[async_trait::async_trait] +impl RegistryDb for InMemoryDb { + async fn register_validators(&self, registration: Registration) -> DbResult<()> { + info!( + keys_count = registration.validator_pubkeys.len(), + sig_count = registration.signatures.len(), + digest = ?registration.digest(), + "NoOpDb: register_validators" + ); + + let mut registrations = self.validator_registrations.write().unwrap(); + registrations.push(registration); + + Ok(()) + } + + async fn register_operator(&self, operator: Operator) -> DbResult<()> { + info!(signer = %operator.signer, "NoOpDb: register_operator"); + + let mut operators = self.operator_registrations.write().unwrap(); + operators.push(operator); + + Ok(()) + } + + async fn get_operator(&self, signer: Address) -> DbResult> { + let operators = self.operator_registrations.read().unwrap(); + let operator = operators.iter().find(|op| op.signer == signer); + + match operator { + Some(op) => Ok(Some(op.clone())), + None => Ok(None), + } + } + + async fn get_validator_registration( + &self, + pubkey: BlsPublicKey, + ) -> DbResult> { + let registrations = self.validator_registrations.read().unwrap(); + let registration = registrations.iter().find(|reg| reg.validator_pubkeys.contains(&pubkey)); + + match registration { + Some(reg) => Ok(Some(reg.clone())), + None => Ok(None), + } + } +} diff --git a/src/db/mod.rs b/src/db/mod.rs index 1e7bd9f..64e822a 100644 --- a/src/db/mod.rs +++ b/src/db/mod.rs @@ -12,8 +12,8 @@ use crate::primitives::{ mod types; /// No-op database implementation. -mod noop; -pub(crate) use noop::NoOpDb; +mod memory; +pub(crate) use memory::InMemoryDb; /// SQL database backend implementation. mod sql; @@ -50,8 +50,11 @@ pub(crate) trait RegistryDb: Clone { async fn register_operator(&self, operator: Operator) -> DbResult<()>; /// Get an operator from the database. - async fn get_operator(&self, signer: Address) -> DbResult; + async fn get_operator(&self, signer: Address) -> DbResult>; /// Get a validator registration from the database. - async fn get_validator_registration(&self, pubkey: BlsPublicKey) -> DbResult; + async fn get_validator_registration( + &self, + pubkey: BlsPublicKey, + ) -> DbResult>; } diff --git a/src/db/noop.rs b/src/db/noop.rs deleted file mode 100644 index c89374b..0000000 --- a/src/db/noop.rs +++ /dev/null @@ -1,36 +0,0 @@ -use alloy::primitives::Address; - -use super::{BlsPublicKey, DbResult, Operator, Registration, RegistryDb}; - -#[derive(Debug, Clone)] -pub(crate) struct NoOpDb; - -#[async_trait::async_trait] -impl RegistryDb for NoOpDb { - async fn register_validators(&self, _registration: Registration) -> DbResult<()> { - Ok(()) - } - - async fn register_operator(&self, _operator: Operator) -> DbResult<()> { - Ok(()) - } - - async fn get_operator(&self, signer: Address) -> DbResult { - Ok(Operator { - signer, - rpc_endpoint: "https://grugbrain.dev".parse()?, - collateral_tokens: vec![], - collateral_amounts: vec![], - }) - } - - async fn get_validator_registration(&self, pubkey: BlsPublicKey) -> DbResult { - Ok(Registration { - validator_pubkeys: vec![pubkey], - operator: Address::default(), - gas_limit: 0, - expiry: 0, - signatures: vec![], - }) - } -} diff --git a/src/db/sql.rs b/src/db/sql.rs index 2dc8900..0100650 100644 --- a/src/db/sql.rs +++ b/src/db/sql.rs @@ -97,8 +97,8 @@ impl RegistryDb for SQLDb { Ok(()) } - async fn get_operator(&self, signer: Address) -> DbResult { - let row: OperatorRow = sqlx::query_as( + async fn get_operator(&self, signer: Address) -> DbResult> { + let row: Option = sqlx::query_as( " SELECT signer, rpc, protocol, source, collateral_tokens, collateral_amounts, last_update FROM operators @@ -106,14 +106,21 @@ impl RegistryDb for SQLDb { ", ) .bind(signer.to_vec()) - .fetch_one(&self.conn) + .fetch_optional(&self.conn) .await?; - row.try_into() + let Some(row) = row else { + return Ok(None); + }; + + Ok(Some(row.try_into()?)) } - async fn get_validator_registration(&self, pubkey: BlsPublicKey) -> DbResult { - let row: ValidatorRegistrationRow = sqlx::query_as( + async fn get_validator_registration( + &self, + pubkey: BlsPublicKey, + ) -> DbResult> { + let row: Option = sqlx::query_as( " SELECT pubkey, signature, expiry, gas_limit, operator, priority, source, last_update FROM validator_registrations @@ -121,9 +128,13 @@ impl RegistryDb for SQLDb { ", ) .bind(pubkey.serialize().to_vec()) - .fetch_one(&self.conn) + .fetch_optional(&self.conn) .await?; - row.try_into() + let Some(row) = row else { + return Ok(None); + }; + + Ok(Some(row.try_into()?)) } }