Reduce code duplication due to monomorphization

bin/crates_code_size.bloaty

@@ -1,4 +1,12 @@
-#!/usr/bin/env bloaty -d crates -s vm -c
+#!/usr/bin/env -S bloaty -d crates -s vm -c
+#
+# Uses bloaty from https://github.com/google/bloaty. Run as
+#
+#   ./crates_code_size.bloaty target/aarch64-linux-android/release/libsignal_jni.so -- baseline.so
+#
+# where baseline.so is the same file (here, libsignal_jni.so) built at the
+# version to compare against.
+#
 # We use VM size because otherwise the debug sections are included.
 
 custom_data_source: {

rust/bridge/shared/src/net/tokio.rs

@@ -8,6 +8,7 @@ use std::future::Future;
 use std::sync::atomic::AtomicU64;
 use std::sync::{Arc, Mutex};
 
+use futures_util::future::BoxFuture;
 use futures_util::FutureExt as _;
 use libsignal_bridge_macros::bridge_fn;
 
@@ -103,6 +104,37 @@ where
         &self,
         make_future: impl FnOnce(TokioContextCancellation) -> F,
         completer: <F::Output as ResultReporter>::Receiver,
+    ) -> CancellationId {
+        // Delegate to a non-templated function with dynamic dispatch to save on
+        // compiled code size.
+        self.run_future_boxed(Box::new(move |cancellation| {
+            let future = make_future(cancellation);
+            async {
+                let reporter = future.await;
+                let report_cb: Box<dyn FnOnce() + Send> =
+                    Box::new(move || reporter.report_to(completer));
+                report_cb
+            }
+            .boxed()
+        }))
+    }
+}
+
+type ReportResultBoxed = Box<dyn FnOnce() + Send>;
+
+impl TokioAsyncContext {
+    /// Create and spawn a `Future` as a task, then spawn a blocking task to
+    /// execute the output callback.
+    ///
+    /// This intentionally uses dynamic dispatch to save on code size. Future
+    /// spawning in templated code (that gets duplicated during
+    /// monomorphization) should call this method with a callback of the
+    /// appropriate type.
+    fn run_future_boxed<'s>(
+        &'s self,
+        make_future: Box<
+            dyn 's + FnOnce(TokioContextCancellation) -> BoxFuture<'static, ReportResultBoxed>,
+        >,
     ) -> CancellationId {
         let (cancel_tx, cancel_rx) = tokio::sync::oneshot::channel();
 
@@ -128,9 +160,8 @@ where
 
         #[allow(clippy::let_underscore_future)]
         let _: tokio::task::JoinHandle<()> = self.rt.spawn(async move {
-            let completed = future.await;
-            let _: tokio::task::JoinHandle<()> =
-                handle.spawn_blocking(move || completed.report_to(completer));
+            let report_fn = future.await;
+            let _: tokio::task::JoinHandle<()> = handle.spawn_blocking(report_fn);
             // What happens if we don't get here? We leak an entry in the task map. Also, we
             // probably have bigger problems, because in practice all the `bridge_io` futures are
             // supposed to be catching panics.

rust/net/src/enclave.rs

@@ -44,7 +44,7 @@ impl<'a> AsRaftConfig<'a> for &'a RaftConfig {
 }
 
 pub trait EnclaveKind {
-    type RaftConfigType: AsRaftConfig<'static> + Clone;
+    type RaftConfigType: AsRaftConfig<'static> + Clone + Sync;
     fn url_path(enclave: &[u8]) -> PathAndQuery;
 }
 
@@ -177,7 +177,9 @@ impl<S: Send> PpssSetup<S> for Svr3Env<'_> {
 #[derive_where(Clone, Copy; Bytes)]
 pub struct MrEnclave<Bytes, E> {
     inner: Bytes,
-    enclave_kind: PhantomData<E>,
+    // Using fn instead of E directly so that `MrEnclave` implements `Send +
+    // Sync` even if `E` does not.
+    enclave_kind: PhantomData<fn(E) -> E>,
 }
 
 impl<Bytes, E: EnclaveKind> MrEnclave<Bytes, E> {
@@ -257,36 +259,58 @@ impl<E: EnclaveKind + NewHandshake, C: ConnectionManager> EnclaveEndpointConnect
     where
         C: ConnectionManager,
     {
-        let auth_decorator = auth.into();
-        let connector = ServiceConnectorWithDecorator::new(
-            WebSocketClientConnector::<_, WebSocketServiceError>::new(
-                transport_connector,
-                self.endpoint_connection.config.clone(),
-            ),
-            auth_decorator,
-        );
-        let service_initializer =
-            ServiceInitializer::new(connector, &self.endpoint_connection.manager);
-        let connection_attempt_result = service_initializer.connect().await;
-        let websocket = match connection_attempt_result {
-            ServiceState::Active(websocket, _) => Ok(websocket),
-            ServiceState::Error(e) => Err(Error::WebSocketConnect(e)),
-            ServiceState::Cooldown(_) | ServiceState::ConnectionTimedOut => {
-                Err(Error::ConnectionTimedOut)
-            }
-            ServiceState::Inactive => {
-                unreachable!("can't be returned by the initializer")
-            }
-        }?;
-        let attested = AttestedConnection::connect(websocket, |attestation_msg| {
-            E::new_handshake(&self.params, attestation_msg)
-        })
-        .await?;
-
-        Ok(attested)
+        // Delegate to a function that dynamically-dispatches. This could be
+        // inlined, but then the body would be duplicated in the generated code
+        // for each instantiation of this trait (of which there is one per
+        // unique `E: EnclaveKind`).
+        connect_attested(
+            &self.endpoint_connection,
+            auth,
+            transport_connector,
+            &move |attestation_message| E::new_handshake(&self.params, attestation_message),
+        )
+        .await
     }
 }
 
+/// Create an `AttestedConnection`.
+///
+/// Making the handshaker a concrete type (via `&dyn`) prevents this from being
+/// instantiated multiple times and duplicated in the generated code.
+async fn connect_attested<
+    C: ConnectionManager,
+    T: TransportConnector<Stream = S>,
+    S: AsyncDuplexStream,
+>(
+    endpoint_connection: &EndpointConnection<C>,
+    auth: impl HttpBasicAuth,
+    transport_connector: T,
+    do_handshake: &(dyn Sync + Fn(&[u8]) -> enclave::Result<enclave::Handshake>),
+) -> Result<AttestedConnection<S>, Error> {
+    let auth_decorator = auth.into();
+    let connector = ServiceConnectorWithDecorator::new(
+        WebSocketClientConnector::<_, WebSocketServiceError>::new(
+            transport_connector,
+            endpoint_connection.config.clone(),
+        ),
+        auth_decorator,
+    );
+    let service_initializer = ServiceInitializer::new(connector, &endpoint_connection.manager);
+    let connection_attempt_result = service_initializer.connect().await;
+    let websocket = match connection_attempt_result {
+        ServiceState::Active(websocket, _) => Ok(websocket),
+        ServiceState::Error(e) => Err(Error::WebSocketConnect(e)),
+        ServiceState::Cooldown(_) | ServiceState::ConnectionTimedOut => {
+            Err(Error::ConnectionTimedOut)
+        }
+        ServiceState::Inactive => {
+            unreachable!("can't be returned by the initializer")
+        }
+    }?;
+    let attested = AttestedConnection::connect(websocket, do_handshake).await?;
+    Ok(attested)
+}
+
 impl<E: EnclaveKind> EnclaveEndpointConnection<E, SingleRouteThrottlingConnectionManager> {
     pub fn new(endpoint: &EnclaveEndpoint<'static, E>, connect_timeout: Duration) -> Self {
         Self {