Have aws_smithy_types::config_bag::Layer return the replaced value

[kikuomax]

Summary

Make the unset, store_put, and store_or_unset methods of aws_smithy_types::config_bag::Layer (Layer) return the previously stored value(s).

These methods currently return a mutable reference to the Layer and never return the previously stored value.

Background

I am currently writing a crate that facilitates X-Ray tracing in AWS Lambda. I came up with an idea to utilize interceptor to start and end a subsegment for an AWS service operation. Here is what in my mind:

impl Intercept for XrayIntercept {
    // starts a subsegment at every attempt
    fn read_before_attempt(
        &self,
        _context: &BeforeTransmitInterceptorContextRef<'_>,
        _runtime_components: &RuntimeComponents,
        cfg: &mut ConfigBag,
    ) -> Result<(), BoxError> {
        let session = self.enter_subsegment(self.service, self.operation);
        cfg.interceptor_state().store_put(session);
        Ok(())
    }

    // injects the X-Amzn-Trace-Id header
    fn modify_before_transmit(
        &self,
        context: &mut BeforeTransmitInterceptorContextMut<'_>,
        _runtime_components: &RuntimeComponents,
        cfg: &mut ConfigBag,
    ) -> Result<(), BoxError> {
        let trace_id = cfg.interceptor_state().load::<SubsegmentSession>()
            .and_then(|s| s.x_amzn_trace_id());
        if let Some(trace_id) = trace_id {
            context
                .request_mut()
                .headers_mut()
                .insert("X-Amzn-Trace-Id", trace_id);
        }
        Ok(())
    }

    // ends the subsegment updated with the request ID and response
    fn read_after_attempt(
        &self,
        context: &FinalizerInterceptorContextRef<'_>,
        _runtime_components: &RuntimeComponents,
        cfg: &mut ConfigBag,
    ) -> Result<(), BoxError> {
        let session = cfg.interceptor_state().unset::<SubsegmentSession>();
        if let Some(mut session) = session {
            if let Some(response) = context.response() {
                session.response_status(response.status().as_u16());
                if let Some(request_id) = response.request_id() {
                    session.request_id(request_id);
                }
            }
        }
        // session reports the subsegment to the X-Ray daemon when it is dropped
        Ok(())
    }
}

Unfortunately, the above idea is not feasible because the unset method of Layer won't return the replaced value; i.e., I have no chance to update the session information before it is dropped.

What I am currently doing is to make the above XrayIntercept have a session wrapped in Mutex, which I feel awkward:

struct XrayIntercept {
    // wrapped in Mutex because Intercept has to be immutable
    session: Arc<Mutex<Option<SubsegmentSession>>>,
}

impl Intercept for XrayIntercept {
    fn read_before_attempt(
        &self,
        _context: &BeforeTransmitInterceptorContextRef<'_>,
        _runtime_components: &RuntimeComponents,
        _cfg: &mut ConfigBag,
    ) -> Result<(), BoxError> {
        let session = self.enter_subsegment(self.service, self.operation);
        *self.session.lock().unwrap() = Some(session);
        Ok(())
    }

    fn modify_before_transmit(
        &self,
        context: &mut BeforeTransmitInterceptorContextMut<'_>,
        _runtime_components: &RuntimeComponents,
        _cfg: &mut ConfigBag,
    ) -> Result<(), BoxError> {
        let trace_id = self
            .session
            .lock()
            .unwrap()
            .as_ref()
            .and_then(|s| s.x_amzn_trace_id());
        if let Some(trace_id) = trace_id {
            context
                .request_mut()
                .headers_mut()
                .insert("X-Amzn-Trace-Id", trace_id);
        }
        Ok(())
    }

    fn read_after_attempt(
        &self,
        context: &FinalizerInterceptorContextRef<'_>,
        _runtime_components: &RuntimeComponents,
        _cfg: &mut ConfigBag,
    ) -> Result<(), BoxError> {
        let mut session = self.session.lock().unwrap();
        if let Some(mut session) = session.take() {
            if let Some(response) = context.response() {
                session.response_status(response.status().as_u16());
                if let Some(request_id) = response.request_id() {
                    session.request_id(request_id);
                }
            }
        }
        Ok(())
    }
}

You can find the actual code here.

Suggested changes (PoC)

I have a PoC of my suggested changes.

https://github.com/codemonger-io/smithy-rs/blob/daa7e3fda222d6efa155022d908d666307f9d4e9/rust-runtime/aws-smithy-types/src/config_bag.rs#L416-L443

    /// Remove `T` from this bag
    pub fn unset<T: Send + Sync + Debug + 'static>(&mut self) -> UnsetResult<'_> {
        let value = self.put_directly::<StoreReplace<T>>(Value::ExplicitlyUnset(type_name::<T>()));
        UnsetResult::new(self, value)
    }

    /// Stores `item` of type `T` into the config bag, overriding a previous value of the same type
    pub fn store_put<T>(&mut self, item: T) -> UnsetResult<'_>
    where
        T: Storable<Storer = StoreReplace<T>>,
    {
        let old = self.put_directly::<StoreReplace<T>>(Value::Set(item));
        UnsetResult::new(self, old)
    }

    /// Stores `item` of type `T` into the config bag, overriding a previous value of the same type,
    /// or unsets it by passing a `None`
    pub fn store_or_unset<T>(&mut self, item: Option<T>) -> UnsetResult<'_>
    where
        T: Storable<Storer = StoreReplace<T>>,
    {
        let item = match item {
            Some(item) => Value::Set(item),
            None => Value::ExplicitlyUnset(type_name::<T>()),
        };
        let old = self.put_directly::<StoreReplace<T>>(item);
        UnsetResult::new(self, old)
    }

unset, store_put, and store_or_unset no longer return &mut Layer but UnsetResult which can be derefed as &mut Layer for backward compatibility.

https://github.com/codemonger-io/smithy-rs/blob/daa7e3fda222d6efa155022d908d666307f9d4e9/rust-runtime/aws-smithy-types/src/config_bag.rs#L518-L551

/// Unset result.
#[derive(Debug)]
pub struct UnsetResult<'a> {
    layer: &'a mut Layer,
    value: Option<UnsetResultValue>,
}

impl<'a> UnsetResult<'a> {
    fn new(layer: &'a mut Layer, value: Option<TypeErasedBox>) -> Self {
        Self {
            layer,
            value: value.map(UnsetResultValue),
        }
    }

    /// Decomposes the result into the layer and the value.
    pub fn into_parts(self) -> (&'a mut Layer, Option<UnsetResultValue>) {
        (self.layer, self.value)
    }
}

impl<'a> std::ops::Deref for UnsetResult<'a> {
    type Target = Layer;

    fn deref(&self) -> &Self::Target {
        self.layer
    }
}

impl<'a> std::ops::DerefMut for UnsetResult<'a> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        self.layer
    }
}

UnsetResultValue is introduced to extract the underlying value that is wrapped in the private enum Value when it is stored in the Layer.

https://github.com/codemonger-io/smithy-rs/blob/daa7e3fda222d6efa155022d908d666307f9d4e9/rust-runtime/aws-smithy-types/src/config_bag.rs#L553-L568

/// Value part of [`UnsetResult`].
#[derive(Debug)]
pub struct UnsetResultValue(TypeErasedBox);

impl UnsetResultValue {
    /// Downcasts the underlying value to a given type `T`.
    ///
    /// This method extracts the mutable reference to the underlying value that
    /// is wrapped in `Value::Set`.
    pub fn downcast_mut<T: fmt::Debug + Send + Sync + 'static>(&mut self) -> Option<&mut T> {
        match self.0.downcast_mut::<Value<T>>() {
            Some(Value::Set(t)) => Some(t),
            _ => None,
        }
    }
}

The read_after_attempt method in my initial intended code would become something similar to the following if we incorporated the above suggestion into:

    // ends the subsegment updated with the request ID and response
    fn read_after_attempt(
        &self,
        context: &FinalizerInterceptorContextRef<'_>,
        _runtime_components: &RuntimeComponents,
        cfg: &mut ConfigBag,
    ) -> Result<(), BoxError> {
        let result = cfg.interceptor_state().unset::<SubsegmentSession>();
        let (_, mut session) = result.into_parts();
        if let Some(session) = session.downcast_mut::<SubsegmentSession>() {
            if let Some(response) = context.response() {
                session.response_status(response.status().as_u16());
                if let Some(request_id) = response.request_id() {
                    session.request_id(request_id);
                }
            }
        }
        // session reports the subsegment to the X-Ray daemon when it is dropped
        Ok(())
    }