Skip to content

Commit

Permalink
Move SingleTargetLBP to common/utils.rs
Browse files Browse the repository at this point in the history
Moves `SingleTargetLBP` used by `integration/tablets.rs` to `common/utils.rs`
Code is identical except added `pub(crate)` to the struct and its
`target` field.
Moving it to allow usage in `integration/caching_session.rs` from
subsequent commit and in other tests.
  • Loading branch information
Bouncheck committed Feb 13, 2025
1 parent a253772 commit b678d0d
Show file tree
Hide file tree
Showing 2 changed files with 29 additions and 31 deletions.
28 changes: 28 additions & 0 deletions scylla/tests/common/utils.rs
Original file line number Diff line number Diff line change
Expand Up @@ -300,3 +300,31 @@ impl PerformDDL for CachingSession {
self.execute_unpaged(query, &[]).await.map(|_| ())
}
}

#[derive(Debug)]
#[allow(dead_code)]
pub(crate) struct SingleTargetLBP {
pub(crate) target: (Arc<scylla::cluster::Node>, Option<u32>),
}

impl LoadBalancingPolicy for SingleTargetLBP {
fn pick<'a>(
&'a self,
_query: &'a RoutingInfo,
_cluster: &'a ClusterState,
) -> Option<(NodeRef<'a>, Option<u32>)> {
Some((&self.target.0, self.target.1))
}

fn fallback<'a>(
&'a self,
_query: &'a RoutingInfo,
_cluster: &'a ClusterState,
) -> FallbackPlan<'a> {
Box::new(std::iter::empty())
}

fn name(&self) -> String {
"SingleTargetLBP".to_owned()
}
}
32 changes: 1 addition & 31 deletions scylla/tests/integration/tablets.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
use std::sync::Arc;

use crate::utils::SingleTargetLBP;
use crate::utils::{
scylla_supports_tablets, setup_tracing, test_with_3_node_cluster, unique_keyspace_name,
PerformDDL,
Expand All @@ -12,10 +13,6 @@ use scylla::client::execution_profile::ExecutionProfile;
use scylla::client::session::Session;
use scylla::cluster::ClusterState;
use scylla::cluster::Node;
use scylla::cluster::NodeRef;
use scylla::policies::load_balancing::FallbackPlan;
use scylla::policies::load_balancing::LoadBalancingPolicy;
use scylla::policies::load_balancing::RoutingInfo;
use scylla::prepared_statement::PreparedStatement;
use scylla::query::Query;
use scylla::response::query_result::QueryResult;
Expand Down Expand Up @@ -156,33 +153,6 @@ fn calculate_key_per_tablet(tablets: &[Tablet], prepared: &PreparedStatement) ->
value_lists
}

#[derive(Debug)]
struct SingleTargetLBP {
target: (Arc<Node>, Option<u32>),
}

impl LoadBalancingPolicy for SingleTargetLBP {
fn pick<'a>(
&'a self,
_query: &'a RoutingInfo,
_cluster: &'a ClusterState,
) -> Option<(NodeRef<'a>, Option<u32>)> {
Some((&self.target.0, self.target.1))
}

fn fallback<'a>(
&'a self,
_query: &'a RoutingInfo,
_cluster: &'a ClusterState,
) -> FallbackPlan<'a> {
Box::new(std::iter::empty())
}

fn name(&self) -> String {
"SingleTargetLBP".to_owned()
}
}

async fn send_statement_everywhere(
session: &Session,
cluster: &ClusterState,
Expand Down

0 comments on commit b678d0d

Please sign in to comment.