Prepare Query API Plumbing #1398

cberkhoff · 2024-11-02T00:54:49Z

A few changes required to make Prepare Query API work for the sharded scenario:

HandlerRef was generalized to support both <ShardIndex> and <HelperIdentity>
app::Setup returns also a HandlerRef<ShardIndex> to give to the ShardHttpTransport.
Added ShardTransports trait. This one provides an additional function to retrieve my shard configuration (Sharded). I need shard transports to know how many shards are in total. I will use this functionality in the Prepare Query API (query processor) so that the leader shard can send messages to all other shards. Note there's no Context at this stage, these are the preliminary steps of the query.
Created a simple net::errors::ShardError. This one only wraps the general net Error with shard information. This is needed by the processor to be able to distinguish errors that happen in the MPC transport vs the ones in the shard transport. Also provides more context information for the error itself.
Related to 3; So that TestWorlds can create In-Memory Shard Transport I had to expand ShardContext = Option<ShardIndex> to also contain the total count of shards. For simplicity I ended up using Option<ShardIndex>. The name ShardContext was overly generic so I removed it.

codecov · 2024-11-02T01:25:30Z

Codecov Report

Attention: Patch coverage is 63.15789% with 49 lines in your changes missing coverage. Please review.

Project coverage is 93.40%. Comparing base (bbf2deb) to head (c730a02).
Report is 2 commits behind head on main.

Files with missing lines	Patch %	Lines
ipa-core/src/helpers/transport/mod.rs	18.75%	13 Missing ⚠️
ipa-core/src/net/transport.rs	18.75%	13 Missing ⚠️
ipa-core/src/app.rs	55.55%	8 Missing ⚠️
ipa-core/src/sharding.rs	33.33%	6 Missing ⚠️
ipa-core/src/bin/helper.rs	0.00%	3 Missing ⚠️
ipa-core/src/helpers/gateway/transport.rs	0.00%	3 Missing ⚠️
...-core/src/helpers/transport/in_memory/transport.rs	76.92%	3 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1398      +/-   ##
==========================================
- Coverage   93.48%   93.40%   -0.09%     
==========================================
  Files         223      223              
  Lines       38483    38461      -22     
==========================================
- Hits        35976    35924      -52     
- Misses       2507     2537      +30

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

akoshelev · 2024-11-04T16:35:22Z

ipa-core/src/bin/helper.rs

@@ -185,11 +185,14 @@ async fn server(args: ServerArgs, logging_handle: LoggingHandle) -> Result<(), B
    let shard_network_config = NetworkConfig::new_shards(vec![], shard_clients_config);
    let (shard_transport, _shard_server) = ShardHttpTransport::new(
        IpaRuntime::from_tokio_runtime(&http_runtime),
-        ShardIndex::FIRST,
+        Sharded {


something like Sharded::with_shards(1) could be more concise to read

This initialization to be the single shard is temporary. I created a new function to help with the conciseness.

akoshelev · 2024-11-04T16:37:37Z

ipa-core/src/config.rs

+    /// In the unlikely event a usize cannot be turned into a u32
+    #[must_use]
+    pub fn shard_count(&self) -> ShardIndex {
+        ShardIndex::from(u32::try_from(self.peers.len()).unwrap())


Suggested change

ShardIndex::from(u32::try_from(self.peers.len()).unwrap())

ShardIndex::try_from(self.peers.len()).unwrap()

akoshelev · 2024-11-04T16:38:07Z

ipa-core/src/config.rs

@@ -131,6 +131,13 @@ impl NetworkConfig<Shard> {
            identities,
        }
    }
+
+    /// # Panics
+    /// In the unlikely event a usize cannot be turned into a u32


nit: I wouldn't say it is unlikely for usize, but it is surely unexpected to have 4B shards

akoshelev · 2024-11-04T16:40:53Z

ipa-core/src/sharding.rs

 pub trait ShardBinding: Debug + Send + Sync + Clone + 'static {
-    fn context(&self) -> ShardContext;
+    fn shard_config(&self) -> Option<Sharded>;


this deserves documenting the purpose of it like we had before for ShardContext

akoshelev · 2024-11-04T16:44:27Z

ipa-core/src/helpers/transport/mod.rs

@@ -333,6 +333,10 @@ pub trait Transport: Clone + Send + Sync + 'static {
    }
 }

+pub trait ShardedTransport: Transport {


I think we should extend Transport trait instead with these

fn peer_count(&self) -> u32; /// Broadcasts a message to all peers, excluding this instance, collecting all failures and /// successes. This method waits for all responses and returns only when all peers responded. async fn broadcast<Q, S, R>( &self, route: R, ) -> Result<(), BroadcastError<Self::Identity, Self::Error>> where Option<QueryId>: From<Q>, Option<Gate>: From<S>, Q: QueryIdBinding, S: StepBinding, R: RouteParams<RouteId, Q, S> + Clone;

It is easy to implement, works for both MPC and Shard and avoids transport specialization which this code invites us to do

I understand you would still allow Transport users to check the number of shards via peer_count?

that's right. Depending on prepare query orchestration logic, you may just use broadcast

akoshelev · 2024-11-04T16:46:46Z

ipa-core/src/helpers/transport/handler.rs

@@ -173,18 +173,13 @@ pub enum Error {
 /// Trait for custom-handling different request types made against MPC helper parties.
 /// There is a limitation for RPITIT that traits can't be made object-safe, hence the use of async_trait
 #[async_trait]
-pub trait RequestHandler: Send + Sync {
-    type Identity: TransportIdentity;
+pub trait RequestHandler<I: TransportIdentity>: Send + Sync {


It was intentional to have Identity as an associated type - that ensures for each peer type we have exactly one handler. Unless we have to change it, I would prefer to have it as AT

Can you please clarify "ensures for each peer type we have exactly one handler"? The following code is perfectly fine to add:

struct AnotherHandler; #[async_trait] impl RequestHandler for AnotherHandler { type Identity = HelperIdentity; async fn handle( &self, req: Addr<HelperIdentity>, data: BodyStream, ) -> Result<HelperResponse, ApiError> { Ok(HelperResponse::ok()) } }

In other words, an associated type CAN have multiple handlers.

An associated type prevents having a single implementor support multiple associated types. I don't think that's a desirable property here. More concretely, It's simpler to have Inner implement both Shard and Mpc Handlers. The alternative is to have 2 structs which are going to be interconnected anyway because they want to use the same Query Processor and transports. I would advocate on keeping things simple and avoid proliferation of wrappers. That's what I'm proposing on this PR.

Let me know what you think, but I will need some more clarification on what you're overall trying to do to find the right abstraction.

The concrete case where it's needed / desirable to implement both kinds of handlers is the HelperApp (it wasn't needed previously, but it is needed now, to handle prepare-query messages from the leader shard).

A lot of the uses of RequestHandler are via make_owned_handler, where there's an indirect limitation that the handler can only handle one identity type, because the handler wraps Fn(Addr<I>).

I'm a bit surprised, but I can't even find other concrete types implementing RequestHandler at all, so I can't assess risk of incorrectly implementing RequestHandler for multiple identity types. Are there places I'm missing?

@andyleiserson We wanted to keep request handlers in one place, so not having too many implementations of RequestHandlers was the goal of #992.

Given that we need a single type to implement both, generic is the way to go

andyleiserson · 2024-11-04T18:08:30Z

ipa-core/src/helpers/transport/in_memory/transport.rs

@@ -268,7 +274,7 @@ impl Debug for InMemoryStream {
 }

 pub struct Setup<I> {
-    identity: I,
+    pub identity: I,


Suggested change

pub identity: I,

identity: I,

Seems like pub is not needed here?

andyleiserson · 2024-11-04T18:10:19Z

ipa-core/src/test_fixture/world.rs

@@ -1097,7 +1098,9 @@ mod tests {
                config.role_assignment(),
                |ctx: &MaliciousHelperContext, data: &mut Vec<u8>| {
                    assert!(ctx.shard.is_some());


Suggested change

assert!(ctx.shard.is_some());

No need to have the assert and the unwrap.

akoshelev · 2024-11-05T17:04:11Z

ipa-core/src/helpers/transport/mod.rs

+    /// Broadcasts a message to all peers, excluding this instance, collecting all failures and
+    /// successes. This method waits for all responses and returns only when all peers responded.
+    /// The routes and data will be cloned.
+    async fn broadcast<Q, S, R, D>(&self, route: R, data: D) -> Result<(), Self::ShardError>


can we make it part of the Transport trait? My goal here was to avoid having ShardedTransport trait altogether. Specializing shard transport based on the existing transport interface seems like a wrong way to go to me.

While it's technically possible, I don't see advantage in having one Transport. In fact, I think making both have the same API can lead to misuse. I rather have the compiler help the developer identify the transport it's using and its capabilities.

As an example; The query processor needs to know if it's dealing with a sharded or mpc transport for the Prepare API. Say you're the leader of helper 2 and you need to let your shards know about the incoming query. If both transports are interchangeable, then you could by mistake (a bug) call the broadcast on the mpc but not on the shard transport. With 2 different traits the compiler can protect you.

Imo, the distinction is already made in Transport::Identity, so it is not possible to swap them. Currently both are implemented using the same backend and making them conform to the same API is a better strategy as long as we plan to use the same implementation for both.

cberkhoff added 2 commits November 1, 2024 17:28

Transport plumbing

8a0a7d4

Fixing re-branching errors

d3348cf

cberkhoff changed the title ~~Si plumb3~~ Prepare Query API Plumbing Nov 2, 2024

cberkhoff requested review from andyleiserson and akoshelev November 2, 2024 00:55

akoshelev requested changes Nov 4, 2024

View reviewed changes

Addressing smaller comments

b4f591d

andyleiserson reviewed Nov 4, 2024

View reviewed changes

cberkhoff added 5 commits November 4, 2024 15:43

Addressing Andy's comments

16b992a

Clippy

3938b17

ShardTransport Broadcast

ac55d35

Merge branch 'main' into si-plumb3

fd98f20

Fixing merge conflicts

3c6e7e3

cberkhoff requested a review from akoshelev November 5, 2024 02:08

akoshelev reviewed Nov 5, 2024

View reviewed changes

cberkhoff added 2 commits November 5, 2024 19:04

No ShardTransport

fecbc6b

Prunning un-needed stuff

c730a02

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Prepare Query API Plumbing #1398

Prepare Query API Plumbing #1398

cberkhoff commented Nov 2, 2024

codecov bot commented Nov 2, 2024 •

edited

Loading

akoshelev Nov 4, 2024

cberkhoff Nov 4, 2024

akoshelev Nov 4, 2024

akoshelev Nov 4, 2024

akoshelev Nov 4, 2024

akoshelev Nov 4, 2024

cberkhoff Nov 4, 2024

akoshelev Nov 5, 2024

akoshelev Nov 4, 2024

cberkhoff Nov 4, 2024

andyleiserson Nov 4, 2024

akoshelev Nov 4, 2024

andyleiserson Nov 4, 2024

andyleiserson Nov 4, 2024

akoshelev Nov 5, 2024

cberkhoff Nov 5, 2024 •

edited

Loading

akoshelev Nov 5, 2024

	ShardIndex::from(u32::try_from(self.peers.len()).unwrap())
	ShardIndex::try_from(self.peers.len()).unwrap()

Prepare Query API Plumbing #1398

Are you sure you want to change the base?

Prepare Query API Plumbing #1398

Conversation

cberkhoff commented Nov 2, 2024

codecov bot commented Nov 2, 2024 • edited Loading

Codecov Report

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

cberkhoff Nov 5, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

codecov bot commented Nov 2, 2024 •

edited

Loading

cberkhoff Nov 5, 2024 •

edited

Loading