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
16use std::collections::hash_map::Entry;
17use std::collections::HashMap;
18use std::hash::Hash;
19use std::sync::{Arc, Mutex};
20
21use ylong_http::request::uri::{Authority, Scheme};
22
23pub(crate) struct Pool<K, V> {
24    pool: Arc<Mutex<HashMap<K, V>>>,
25}
26
27impl<K, V> Pool<K, V> {
28    pub(crate) fn new() -> Self {
29        Self {
30            pool: Arc::new(Mutex::new(HashMap::new())),
31        }
32    }
33}
34
35impl<K: Eq + Hash, V: Clone> Pool<K, 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)]
49pub(crate) struct PoolKey(Scheme, Authority);
50
51impl PoolKey {
52    pub(crate) fn new(scheme: Scheme, authority: Authority) -> Self {
53        Self(scheme, authority)
54    }
55}
56
57#[cfg(test)]
58mod 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]
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