@@ -126,3 +126,188 @@ impl<T> RerunHeadersExtractorExt for tonic::Request<T> {
126126 Ok ( Some ( entry_name) )
127127 }
128128}
129+
130+ /// Creates a new [`tower::Layer`] middleware that always makes sure to propagate Rerun headers
131+ /// back and forth across requests and responses.
132+ pub fn new_rerun_headers_propagation_layer ( ) -> PropagateHeadersLayer {
133+ PropagateHeadersLayer :: new (
134+ [
135+ http:: HeaderName :: from_static ( RERUN_HTTP_HEADER_ENTRY_ID ) ,
136+ http:: HeaderName :: from_static ( "x-request-id" ) ,
137+ ]
138+ . into_iter ( )
139+ . collect ( ) ,
140+ )
141+ }
142+
143+ // ---
144+
145+ // NOTE: This if a fork of <https://docs.rs/tower-http/0.6.6/tower_http/propagate_header/struct.PropagateHeader.html>.
146+ //
147+ // It exists to prevent never-ending chains of generics when propagating multiple headers, e.g.:
148+ // ```
149+ // pub type RedapClientInner =
150+ // re_perf_telemetry::external::tower_http::propagate_header::PropagateHeader<
151+ // re_perf_telemetry::external::tower_http::propagate_header::PropagateHeader<
152+ // re_perf_telemetry::external::tower_http::propagate_header::PropagateHeader<
153+ // re_perf_telemetry::external::tower_http::propagate_header::PropagateHeader<
154+ // re_perf_telemetry::external::tower_http::trace::Trace<
155+ // tonic::service::interceptor::InterceptedService<
156+ // tonic::service::interceptor::InterceptedService<
157+ // tonic::transport::Channel,
158+ // re_auth::client::AuthDecorator,
159+ // >,
160+ // re_perf_telemetry::TracingInjectorInterceptor,
161+ // >,
162+ // re_perf_telemetry::external::tower_http::classify::SharedClassifier<
163+ // re_perf_telemetry::external::tower_http::classify::GrpcErrorsAsFailures,
164+ // >,
165+ // re_perf_telemetry::GrpcMakeSpan,
166+ // >,
167+ // >,
168+ // >,
169+ // >,
170+ // >;
171+ // ```
172+ // which instead becomes this:
173+ // ```
174+ // pub type RedapClientInner =
175+ // PropagateHeaders<
176+ // re_perf_telemetry::external::tower_http::trace::Trace<
177+ // tonic::service::interceptor::InterceptedService<
178+ // tonic::service::interceptor::InterceptedService<
179+ // tonic::transport::Channel,
180+ // re_auth::client::AuthDecorator,
181+ // >,
182+ // re_perf_telemetry::TracingInjectorInterceptor,
183+ // >,
184+ // re_perf_telemetry::external::tower_http::classify::SharedClassifier<
185+ // re_perf_telemetry::external::tower_http::classify::GrpcErrorsAsFailures,
186+ // >,
187+ // re_perf_telemetry::GrpcMakeSpan,
188+ // >,
189+ // >;
190+ // ```
191+
192+ use std:: collections:: HashSet ;
193+ use std:: future:: Future ;
194+ use std:: {
195+ pin:: Pin ,
196+ task:: { Context , Poll , ready} ,
197+ } ;
198+
199+ use http:: { HeaderValue , Request , Response , header:: HeaderName } ;
200+ use pin_project_lite:: pin_project;
201+ use tower:: Service ;
202+ use tower:: layer:: Layer ;
203+
204+ /// Layer that applies [`PropagateHeaders`] which propagates multiple headers at once from requests to responses.
205+ ///
206+ /// If the headers are present on the request they'll be applied to the response as well. This could
207+ /// for example be used to propagate headers such as `x-rerun-entry-id`, `x-rerun-client-version`, etc.
208+ #[ derive( Clone , Debug ) ]
209+ pub struct PropagateHeadersLayer {
210+ headers : HashSet < HeaderName > ,
211+ }
212+
213+ impl PropagateHeadersLayer {
214+ /// Create a new [`PropagateHeadersLayer`].
215+ pub fn new ( headers : HashSet < HeaderName > ) -> Self {
216+ Self { headers }
217+ }
218+ }
219+
220+ impl < S > Layer < S > for PropagateHeadersLayer {
221+ type Service = PropagateHeaders < S > ;
222+
223+ fn layer ( & self , inner : S ) -> Self :: Service {
224+ PropagateHeaders {
225+ inner,
226+ headers : self . headers . clone ( ) ,
227+ }
228+ }
229+ }
230+
231+ /// Middleware that propagates multiple headers at once from requests to responses.
232+ ///
233+ /// If the headers are present on the request they'll be applied to the response as well. This could
234+ /// for example be used to propagate headers such as `x-rerun-entry-id`, `x-rerun-client-version`, etc.
235+ #[ derive( Clone , Debug ) ]
236+ pub struct PropagateHeaders < S > {
237+ inner : S ,
238+ headers : HashSet < HeaderName > ,
239+ }
240+
241+ impl < S > PropagateHeaders < S > {
242+ /// Create a new [`PropagateHeaders`] that propagates the given header.
243+ pub fn new ( inner : S , headers : HashSet < HeaderName > ) -> Self {
244+ Self { inner, headers }
245+ }
246+
247+ /// Returns a new [`Layer`] that wraps services with a `PropagateHeaders` middleware.
248+ ///
249+ /// [`Layer`]: tower::layer::Layer
250+ pub fn layer ( headers : HashSet < HeaderName > ) -> PropagateHeadersLayer {
251+ PropagateHeadersLayer :: new ( headers)
252+ }
253+ }
254+
255+ impl < ReqBody , ResBody , S > Service < Request < ReqBody > > for PropagateHeaders < S >
256+ where
257+ S : Service < Request < ReqBody > , Response = Response < ResBody > > ,
258+ {
259+ type Response = S :: Response ;
260+ type Error = S :: Error ;
261+ type Future = ResponseFuture < S :: Future > ;
262+
263+ #[ inline]
264+ fn poll_ready ( & mut self , cx : & mut Context < ' _ > ) -> Poll < Result < ( ) , Self :: Error > > {
265+ self . inner . poll_ready ( cx)
266+ }
267+
268+ fn call ( & mut self , req : Request < ReqBody > ) -> Self :: Future {
269+ let headers_and_values = self
270+ . headers
271+ . iter ( )
272+ . filter_map ( |name| {
273+ req. headers ( )
274+ . get ( name)
275+ . cloned ( )
276+ . map ( |value| ( name. clone ( ) , value) )
277+ } )
278+ . collect ( ) ;
279+
280+ ResponseFuture {
281+ future : self . inner . call ( req) ,
282+ headers_and_values,
283+ }
284+ }
285+ }
286+
287+ pin_project ! {
288+ /// Response future for [`PropagateHeaders`].
289+ [ derive( Debug ) ]
290+ pub struct ResponseFuture <F > {
291+ #[ pin]
292+ future: F ,
293+ headers_and_values: Vec <( HeaderName , HeaderValue ) >,
294+ }
295+ }
296+
297+ impl < F , ResBody , E > Future for ResponseFuture < F >
298+ where
299+ F : Future < Output = Result < Response < ResBody > , E > > ,
300+ {
301+ type Output = F :: Output ;
302+
303+ fn poll ( self : Pin < & mut Self > , cx : & mut Context < ' _ > ) -> Poll < Self :: Output > {
304+ let this = self . project ( ) ;
305+ let mut res = ready ! ( this. future. poll( cx) ?) ;
306+
307+ for ( header, value) in std:: mem:: take ( this. headers_and_values ) {
308+ res. headers_mut ( ) . insert ( header, value) ;
309+ }
310+
311+ Poll :: Ready ( Ok ( res) )
312+ }
313+ }
0 commit comments