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 use crate::util::progress::SpeedConfig; 24 25 pub(crate) struct Pool<K, V> { 26 pool: Arc<Mutex<HashMap<K, V>>>, 27 } 28 29 impl<K, V> Pool<K, V> { new() -> Self30 pub(crate) fn new() -> Self { 31 Self { 32 pool: Arc::new(Mutex::new(HashMap::new())), 33 } 34 } 35 } 36 37 impl<K: Eq + Hash, V: Clone> Pool<K, V> { get<F>( &self, key: K, create_fn: F, allowed_num: usize, speed_conf: SpeedConfig, ) -> V where F: FnOnce(usize, SpeedConfig) -> V,38 pub(crate) fn get<F>( 39 &self, 40 key: K, 41 create_fn: F, 42 allowed_num: usize, 43 speed_conf: SpeedConfig, 44 ) -> V 45 where 46 F: FnOnce(usize, SpeedConfig) -> V, 47 { 48 let mut inner = self.pool.lock().unwrap(); 49 match (*inner).entry(key) { 50 Entry::Occupied(conns) => conns.get().clone(), 51 Entry::Vacant(e) => e.insert(create_fn(allowed_num, speed_conf)).clone(), 52 } 53 } 54 } 55 56 #[derive(Debug, PartialEq, Eq, Hash, Clone)] 57 pub(crate) struct PoolKey(Scheme, Authority); 58 59 impl PoolKey { new(scheme: Scheme, authority: Authority) -> Self60 pub(crate) fn new(scheme: Scheme, authority: Authority) -> Self { 61 Self(scheme, authority) 62 } 63 } 64 65 #[cfg(test)] 66 mod ut_pool { 67 use ylong_http::request::uri::Uri; 68 69 use crate::pool::{Pool, PoolKey}; 70 use crate::util::progress::SpeedConfig; 71 72 /// UT test cases for `Pool::get`. 73 /// 74 /// # Brief 75 /// 1. Creates a `pool` by calling `Pool::new()`. 76 /// 2. Uses `pool::get` to get connection. 77 /// 3. Checks if the results are correct. 78 #[test] ut_pool_get()79 fn ut_pool_get() { 80 let uri = Uri::from_bytes(b"http://example1.com:80/foo?a=1").unwrap(); 81 let key = PoolKey::new( 82 uri.scheme().unwrap().clone(), 83 uri.authority().unwrap().clone(), 84 ); 85 let data = String::from("Data info"); 86 let consume_and_return_data = move |_size: usize, _conf: SpeedConfig| data; 87 let pool = Pool::new(); 88 let res = pool.get(key, consume_and_return_data, 6, SpeedConfig::none()); 89 assert_eq!(res, "Data info".to_string()); 90 } 91 } 92