tower_batch_control/
layer.rs1use std::{fmt, marker::PhantomData};
4
5use tower::layer::Layer;
6use tower::Service;
7
8use crate::RequestWeight;
9
10use super::{service::Batch, BatchControl};
11
12pub struct BatchLayer<Request> {
19 max_items_weight_in_batch: usize,
20 max_batches: Option<usize>,
21 max_latency: std::time::Duration,
22
23 _handles_requests: PhantomData<fn(Request)>,
27}
28
29impl<Request> BatchLayer<Request> {
30 pub fn new(
35 max_items_weight_in_batch: usize,
36 max_batches: impl Into<Option<usize>>,
37 max_latency: std::time::Duration,
38 ) -> Self {
39 BatchLayer {
40 max_items_weight_in_batch,
41 max_batches: max_batches.into(),
42 max_latency,
43 _handles_requests: PhantomData,
44 }
45 }
46}
47
48impl<S, Request: RequestWeight> Layer<S> for BatchLayer<Request>
49where
50 S: Service<BatchControl<Request>> + Send + 'static,
51 S::Future: Send,
52 S::Response: Send,
53 S::Error: Into<crate::BoxError> + Send + Sync,
54 Request: Send + 'static,
55{
56 type Service = Batch<S, Request>;
57
58 fn layer(&self, service: S) -> Self::Service {
59 Batch::new(
60 service,
61 self.max_items_weight_in_batch,
62 self.max_batches,
63 self.max_latency,
64 )
65 }
66}
67
68impl<Request> fmt::Debug for BatchLayer<Request> {
69 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
70 f.debug_struct("BufferLayer")
71 .field("max_items_weight_in_batch", &self.max_items_weight_in_batch)
72 .field("max_batches", &self.max_batches)
73 .field("max_latency", &self.max_latency)
74 .finish()
75 }
76}