1 // Copyright (c) 2023 Huawei Device Co., Ltd.
2 // Licensed under the Apache License, Version 2.0 (the "License");
3 // you may not use this file except in compliance with the License.
4 // You may obtain a copy of the License at
5 //
6 //     http://www.apache.org/licenses/LICENSE-2.0
7 //
8 // Unless required by applicable law or agreed to in writing, software
9 // distributed under the License is distributed on an "AS IS" BASIS,
10 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11 // See the License for the specific language governing permissions and
12 // limitations under the License.
13 
14 //! Connection pool implementation.
15 
16 use std::collections::hash_map::Entry;
17 use std::collections::HashMap;
18 use std::hash::Hash;
19 use std::sync::{Arc, Mutex};
20 
21 use ylong_http::request::uri::{Authority, Scheme};
22 
23 pub(crate) struct Pool<K, V> {
24     pool: Arc<Mutex<HashMap<K, V>>>,
25 }
26 
27 impl<K, V> Pool<K, V> {
new() -> Self28     pub(crate) fn new() -> Self {
29         Self {
30             pool: Arc::new(Mutex::new(HashMap::new())),
31         }
32     }
33 }
34 
35 impl<K: Eq + Hash, V: Clone> Pool<K, V> {
get<F>(&self, key: K, create_fn: F) -> V where F: FnOnce() -> V,36     pub(crate) fn get<F>(&self, key: K, create_fn: F) -> V
37     where
38         F: FnOnce() -> V,
39     {
40         let mut inner = self.pool.lock().unwrap();
41         match (*inner).entry(key) {
42             Entry::Occupied(conns) => conns.get().clone(),
43             Entry::Vacant(e) => e.insert(create_fn()).clone(),
44         }
45     }
46 }
47 
48 #[derive(Debug, PartialEq, Eq, Hash, Clone)]
49 pub(crate) struct PoolKey(Scheme, Authority);
50 
51 impl PoolKey {
new(scheme: Scheme, authority: Authority) -> Self52     pub(crate) fn new(scheme: Scheme, authority: Authority) -> Self {
53         Self(scheme, authority)
54     }
55 }
56 
57 #[cfg(test)]
58 mod ut_pool {
59     use ylong_http::request::uri::Uri;
60 
61     use crate::pool::{Pool, PoolKey};
62 
63     /// UT test cases for `Pool::get`.
64     ///
65     /// # Brief
66     /// 1. Creates a `pool` by calling `Pool::new()`.
67     /// 2. Uses `pool::get` to get connection.
68     /// 3. Checks if the results are correct.
69     #[test]
ut_pool_get()70     fn ut_pool_get() {
71         let uri = Uri::from_bytes(b"http://example1.com:80/foo?a=1").unwrap();
72         let key = PoolKey::new(
73             uri.scheme().unwrap().clone(),
74             uri.authority().unwrap().clone(),
75         );
76         let data = String::from("Data info");
77         let consume_and_return_data = move || data;
78         let pool = Pool::new();
79         let res = pool.get(key, consume_and_return_data);
80         assert_eq!(res, "Data info".to_string());
81     }
82 }
83