Log duration of major IPA steps

src/cli/verbosity.rs

@@ -3,7 +3,9 @@ use clap::Parser;
 use metrics_tracing_context::MetricsLayer;
 use std::io::stderr;
 use tracing::{info, metadata::LevelFilter, Level};
-use tracing_subscriber::{fmt, layer::SubscriberExt, util::SubscriberInitExt};
+use tracing_subscriber::{
+    fmt, fmt::format::FmtSpan, layer::SubscriberExt, util::SubscriberInitExt,
+};
 
 #[derive(Debug, Parser)]
 pub struct Verbosity {
@@ -25,7 +27,9 @@ impl Verbosity {
     #[must_use]
     pub fn setup_logging(&self) -> LoggingHandle {
         let filter_layer = self.level_filter();
-        let fmt_layer = fmt::layer().without_time().with_writer(stderr);
+        let fmt_layer = fmt::layer()
+            .with_span_events(FmtSpan::NEW | FmtSpan::CLOSE)
+            .with_writer(stderr);
 
         tracing_subscriber::registry()
             .with(self.level_filter())

src/protocol/attribution/accumulate_credit.rs

@@ -84,6 +84,7 @@ where
 /// # Errors
 ///
 /// Fails if the multiplication fails.
+#[tracing::instrument(name = "accumulate_credit", skip_all)]
 pub async fn accumulate_credit<F, C, T>(
     ctx: C,
     input: &[MCAccumulateCreditInputRow<F, T>],

src/protocol/attribution/aggregate_credit.rs

@@ -44,6 +44,9 @@ const SIMPLE_AGGREGATION_BREAK_EVEN_POINT: u32 = 32;
 ///
 /// # Errors
 /// propagates errors from multiplications
+#[tracing::instrument(name = "aggregate_credit", skip_all)]
+// instrumenting this function makes the return type look bad to Clippy
+#[allow(clippy::type_complexity)]
 pub async fn aggregate_credit<C, V, F, BK, I, S>(
     validator: V,
     sh_ctx: C,

src/protocol/attribution/apply_attribution_window.rs

@@ -23,6 +23,7 @@ use std::{
 ///
 /// # Errors
 /// Fails if sub-protocols fails.
+#[tracing::instrument(name = "apply_window", skip_all)]
 pub async fn apply_attribution_window<C, S, F>(
     ctx: C,
     input: &[MCApplyAttributionWindowInputRow<F, S>],

src/protocol/attribution/credit_capping.rs

@@ -26,6 +26,7 @@ use std::iter::{repeat, zip};
 /// ## Errors
 /// Fails if the multiplication protocol fails, or if the `cap` is larger than
 /// 1/2 of the prime number.
+#[tracing::instrument(name = "user_capping", skip_all)]
 pub async fn credit_capping<F, C, S>(
     ctx: C,
     input: &[MCCreditCappingInputRow<F, S>],

src/protocol/attribution/mod.rs

@@ -44,8 +44,9 @@ use std::iter::{empty, once, zip};
 ///
 /// # Errors
 /// propagates errors from multiplications
+#[tracing::instrument(name = "attribute", skip_all)]
 pub async fn secure_attribution<C, S, SB, F, BK>(
-    sh_ctx: C,
+    ctx: C,
     validator: C::Validator<F>,
     binary_validator: C::Validator<Gf2>,
     sorted_match_keys: Vec<Vec<SB>>,
@@ -74,7 +75,7 @@ where
     let helper_bits_gf2 = compute_helper_bits_gf2(m_binary_ctx, &sorted_match_keys).await?;
     let validated_helper_bits_gf2 = binary_validator.validate(helper_bits_gf2).await?;
     let semi_honest_fp_helper_bits =
-        mod_conv_helper_bits(sh_ctx.clone(), &validated_helper_bits_gf2).await?;
+        mod_conv_helper_bits(ctx.clone(), &validated_helper_bits_gf2).await?;
     let helper_bits = once(S::ZERO)
         .chain(m_ctx.upgrade(semi_honest_fp_helper_bits).await?)
         .collect::<Vec<_>>();
@@ -125,7 +126,7 @@ where
 
     let (validator, output) = aggregate_credit(
         validator,
-        sh_ctx,
+        ctx,
         user_capped_credits.into_iter(),
         config.max_breakdown_key,
         config.num_multi_bits,

src/protocol/context/validator.rs

@@ -235,6 +235,7 @@ impl<'a, F: ExtendableField> Validator<MaliciousContext<'a>, F> for Malicious<'a
     ///
     /// ## Panics
     /// Will panic if the mutex is poisoned
+    #[tracing::instrument(name = "validate", skip_all, fields(step = %self.validate_ctx.step()))]
     async fn validate<D: DowngradeMalicious>(self, values: D) -> Result<D::Target, Error> {
         // send our `u_i+1` value to the helper on the right
         let (u_share, w_share) = self.propagate_u_and_w().await?;

src/protocol/modulus_conversion/convert_shares.rs

@@ -139,6 +139,8 @@ where
 /// Propagates errors from convert shares
 /// # Panics
 /// Propagates panics from convert shares
+#[tracing::instrument(name = "modulus_conversion", skip_all, fields(bits = %num_bits, parallel = %num_multi_bits, step = %ctx.step()))]
+// #[tracing::instrument("IpaQuery::execute", skip_all, fields(query_id = %self.0))]
 pub async fn convert_all_bits<F, C, S>(
     ctx: &C,
     locally_converted_bits: &[Vec<BitConversionTriple<S>>],

src/protocol/sort/apply_sort/mod.rs

@@ -17,6 +17,7 @@ use crate::{
 
 /// # Errors
 /// Propagates errors from shuffle/reshare
+#[tracing::instrument(name = "apply_sort", skip_all, fields(step = %ctx.step()))]
 pub async fn apply_sort_permutation<C, I>(
     ctx: C,
     input: Vec<I>,

src/protocol/sort/generate_permutation.rs

@@ -98,8 +98,9 @@ where
 /// If unable to convert sort keys length to u32
 /// # Errors
 /// If unable to convert sort keys length to u32
+#[tracing::instrument(name = "sort_permutation", skip_all)]
 pub async fn generate_permutation_and_reveal_shuffled<C, S, F>(
-    sh_ctx: C,
+    ctx: C,
     sort_keys: impl Iterator<Item = &Vec<Vec<Replicated<F>>>>,
 ) -> Result<RevealedAndRandomPermutations, Error>
 where
@@ -110,7 +111,7 @@ where
     ShuffledPermutationWrapper<S, C::UpgradedContext<F>>: DowngradeMalicious<Target = Vec<u32>>,
 {
     let (validator, sort_permutation) =
-        generate_permutation_opt(sh_ctx.narrow(&SortKeys), sort_keys).await?;
+        generate_permutation_opt(ctx.narrow(&SortKeys), sort_keys).await?;
 
     let m_ctx = validator.context();
     shuffle_and_reveal_permutation::<C, _, _>(

src/query/executor.rs

@@ -77,7 +77,7 @@ pub fn execute(
     input: ByteArrStream,
 ) -> JoinHandle<QueryResult> {
     match (config.query_type, config.field_type) {
-        #[cfg(any(test, feature = "cli", feature = "test-fixture"))]
+        #[cfg(any(test, feature = "weak-field"))]
         (QueryType::TestMultiply, FieldType::Fp31) => {
             do_query(config, gateway, input, |prss, gateway, input| {
                 Box::pin(execute_test_multiply::<crate::ff::Fp31>(

src/query/runner/ipa.rs

@@ -52,6 +52,7 @@ where
         DowngradeMalicious<Target = MCAggregateCreditOutputRow<F, AdditiveShare<F>, BreakdownKey>>,
     AdditiveShare<F>: Serializable,
 {
+    #[tracing::instrument("ipa_query", skip_all)]
     pub async fn execute<'a>(
         self,
         ctx: C,