hickory_net/client/
memoize_client_handle.rs

1// Copyright 2015-2016 Benjamin Fry <benjaminfry@me.com>
2//
3// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
4// https://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
5// https://opensource.org/licenses/MIT>, at your option. This file may not be
6// copied, modified, or distributed except according to those terms.
7use core::pin::Pin;
8use std::collections::HashMap;
9use std::sync::Arc;
10
11use futures_util::future::FutureExt;
12use futures_util::lock::Mutex;
13use futures_util::stream::Stream;
14
15use super::{
16    ClientHandle,
17    rc_stream::{RcStream, rc_stream},
18};
19use crate::{
20    NetError,
21    proto::op::{DnsRequest, DnsResponse, Query},
22    xfer::DnsHandle,
23};
24
25// TODO: move to proto
26/// A ClientHandle for memoized (cached) responses to queries.
27///
28/// This wraps a ClientHandle, changing the implementation `send()` to store the response against
29///  the Message.Query that was sent. This should reduce network traffic especially during things
30///  like DNSSEC validation. *Warning* this will currently cache for the life of the Client.
31#[derive(Clone)]
32#[must_use = "queries can only be sent through a ClientHandle"]
33pub struct MemoizeClientHandle<H: ClientHandle> {
34    client: H,
35    active_queries: Arc<Mutex<HashMap<Query, RcStream<<H as DnsHandle>::Response>>>>,
36}
37
38impl<H> MemoizeClientHandle<H>
39where
40    H: ClientHandle,
41{
42    /// Returns a new handle wrapping the specified client
43    pub fn new(client: H) -> Self {
44        Self {
45            client,
46            active_queries: Arc::new(Mutex::new(HashMap::new())),
47        }
48    }
49
50    async fn inner_send(
51        request: DnsRequest,
52        active_queries: Arc<Mutex<HashMap<Query, RcStream<<H as DnsHandle>::Response>>>>,
53        client: H,
54    ) -> impl Stream<Item = Result<DnsResponse, NetError>> {
55        // TODO: what if we want to support multiple queries (non-standard)?
56        let query = request.queries.first().expect("no query!").clone();
57
58        // lock all the currently running queries
59        let mut active_queries = active_queries.lock().await;
60
61        // TODO: we need to consider TTL on the records here at some point
62        // If the query is running, grab that existing one...
63        if let Some(rc_stream) = active_queries.get(&query) {
64            return rc_stream.clone();
65        };
66
67        // Otherwise issue a new query and store in the map
68        active_queries
69            .entry(query)
70            .or_insert_with(|| rc_stream(client.send(request)))
71            .clone()
72    }
73}
74
75impl<H: ClientHandle> DnsHandle for MemoizeClientHandle<H> {
76    type Response = Pin<Box<dyn Stream<Item = Result<DnsResponse, NetError>> + Send>>;
77    type Runtime = H::Runtime;
78
79    fn send(&self, request: DnsRequest) -> Self::Response {
80        Box::pin(
81            Self::inner_send(
82                request,
83                Arc::clone(&self.active_queries),
84                self.client.clone(),
85            )
86            .flatten_stream(),
87        )
88    }
89}
90
91#[cfg(test)]
92mod test {
93    #![allow(clippy::dbg_macro, clippy::print_stdout)]
94
95    use core::pin::Pin;
96    use std::sync::Arc;
97
98    use futures_util::lock::Mutex;
99    use futures_util::stream;
100
101    use super::*;
102    use crate::{
103        proto::{
104            op::{DnsRequest, DnsResponse, Message, MessageType, OpCode, Query},
105            rr::RecordType,
106        },
107        runtime::TokioRuntimeProvider,
108        xfer::{DnsHandle, FirstAnswer},
109    };
110    use test_support::subscribe;
111
112    #[derive(Clone)]
113    struct TestClient {
114        i: Arc<Mutex<u16>>,
115    }
116
117    impl DnsHandle for TestClient {
118        type Response = Pin<Box<dyn Stream<Item = Result<DnsResponse, NetError>> + Send>>;
119        type Runtime = TokioRuntimeProvider;
120
121        fn send(&self, request: DnsRequest) -> Self::Response {
122            let i = Arc::clone(&self.i);
123            Box::pin(stream::once(async move {
124                let mut i = i.lock().await;
125                let message = Message::new(*i, MessageType::Query, OpCode::Query).into_response();
126                std::println!(
127                    "sending {}: {}",
128                    *i,
129                    request.queries.first().expect("no query!").clone()
130                );
131
132                *i += 1;
133
134                Ok(DnsResponse::from_message(message).unwrap())
135            }))
136        }
137    }
138
139    #[test]
140    fn test_memoized() {
141        use futures_executor::block_on;
142
143        subscribe();
144
145        let client = MemoizeClientHandle::new(TestClient {
146            i: Arc::new(Mutex::new(0)),
147        });
148
149        let mut test1 = Message::query();
150        test1.add_query(Query::new().set_query_type(RecordType::A).clone());
151
152        let mut test2 = Message::query();
153        test2.add_query(Query::new().set_query_type(RecordType::AAAA).clone());
154
155        let result = block_on(client.send(DnsRequest::from(test1.clone())).first_answer()).unwrap();
156        assert_eq!(result.id, 0);
157
158        let result = block_on(client.send(DnsRequest::from(test2.clone())).first_answer()).unwrap();
159        assert_eq!(result.id, 1);
160
161        // should get the same result for each...
162        let result = block_on(client.send(DnsRequest::from(test1)).first_answer()).unwrap();
163        assert_eq!(result.id, 0);
164
165        let result = block_on(client.send(DnsRequest::from(test2)).first_answer()).unwrap();
166        assert_eq!(result.id, 1);
167    }
168}
hickory_net/client/memoize_client_handle.rs

hickory_net/client/
memoize_client_handle.rs
