1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
use scheduled_thread_pool::ScheduledThreadPool;
use std::fmt;
use std::marker::PhantomData;
use std::sync::Arc;
use std::time::Duration;

use crate::{
    CustomizeConnection, Error, HandleError, HandleEvent, LoggingErrorHandler, ManageConnection,
    NopConnectionCustomizer, NopEventHandler, Pool,
};

/// A builder for a connection pool.
pub struct Builder<M>
where
    M: ManageConnection,
{
    max_size: u32,
    min_idle: Option<u32>,
    test_on_check_out: bool,
    max_lifetime: Option<Duration>,
    idle_timeout: Option<Duration>,
    connection_timeout: Duration,
    error_handler: Box<dyn HandleError<M::Error>>,
    connection_customizer: Box<dyn CustomizeConnection<M::Connection, M::Error>>,
    event_handler: Box<dyn HandleEvent>,
    thread_pool: Option<Arc<ScheduledThreadPool>>,
    reaper_rate: Duration,
    _p: PhantomData<M>,
}

impl<M> fmt::Debug for Builder<M>
where
    M: ManageConnection,
{
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt.debug_struct("Builder")
            .field("max_size", &self.max_size)
            .field("min_idle", &self.min_idle)
            .field("test_on_check_out", &self.test_on_check_out)
            .field("max_lifetime", &self.max_lifetime)
            .field("idle_timeout", &self.idle_timeout)
            .field("connection_timeout", &self.connection_timeout)
            .field("error_handler", &self.error_handler)
            .field("event_handler", &self.event_handler)
            .field("connection_customizer", &self.connection_customizer)
            .finish()
    }
}

impl<M> Default for Builder<M>
where
    M: ManageConnection,
{
    fn default() -> Builder<M> {
        Builder {
            max_size: 10,
            min_idle: None,
            test_on_check_out: true,
            idle_timeout: Some(Duration::from_secs(10 * 60)),
            max_lifetime: Some(Duration::from_secs(30 * 60)),
            connection_timeout: Duration::from_secs(30),
            error_handler: Box::new(LoggingErrorHandler),
            event_handler: Box::new(NopEventHandler),
            connection_customizer: Box::new(NopConnectionCustomizer),
            thread_pool: None,
            reaper_rate: Duration::from_secs(30),
            _p: PhantomData,
        }
    }
}

impl<M> Builder<M>
where
    M: ManageConnection,
{
    /// Constructs a new `Builder`.
    ///
    /// Parameters are initialized with their default values.
    pub fn new() -> Builder<M> {
        Builder::default()
    }

    /// Sets the maximum number of connections managed by the pool.
    ///
    /// Defaults to 10.
    ///
    /// # Panics
    ///
    /// Panics if `max_size` is 0.
    pub fn max_size(mut self, max_size: u32) -> Builder<M> {
        assert!(max_size > 0, "max_size must be positive");
        self.max_size = max_size;
        self
    }

    /// Sets the minimum idle connection count maintained by the pool.
    ///
    /// If set, the pool will try to maintain at least this many idle
    /// connections at all times, while respecting the value of `max_size`.
    ///
    /// Defaults to `None` (equivalent to the value of `max_size`).
    pub fn min_idle(mut self, min_idle: Option<u32>) -> Builder<M> {
        self.min_idle = min_idle;
        self
    }

    /// Sets the thread pool used for asynchronous operations such as connection
    /// creation.
    ///
    /// Defaults to a new pool with 3 threads.
    pub fn thread_pool(mut self, thread_pool: Arc<ScheduledThreadPool>) -> Builder<M> {
        self.thread_pool = Some(thread_pool);
        self
    }

    /// If true, the health of a connection will be verified via a call to
    /// `ConnectionManager::is_valid` before it is checked out of the pool.
    ///
    /// Defaults to true.
    pub fn test_on_check_out(mut self, test_on_check_out: bool) -> Builder<M> {
        self.test_on_check_out = test_on_check_out;
        self
    }

    /// Sets the maximum lifetime of connections in the pool.
    ///
    /// If set, connections will be closed after existing for at most 30 seconds
    /// beyond this duration.
    ///
    /// If a connection reaches its maximum lifetime while checked out it will
    /// be closed when it is returned to the pool.
    ///
    /// Defaults to 30 minutes.
    ///
    /// # Panics
    ///
    /// Panics if `max_lifetime` is the zero `Duration`.
    pub fn max_lifetime(mut self, max_lifetime: Option<Duration>) -> Builder<M> {
        assert_ne!(max_lifetime, Some(Duration::from_secs(0)), "max_lifetime must be positive");
        self.max_lifetime = max_lifetime;
        self
    }

    /// Sets the idle timeout used by the pool.
    ///
    /// If set, connections will be closed after sitting idle for at most 30
    /// seconds beyond this duration.
    ///
    /// Defaults to 10 minutes.
    ///
    /// # Panics
    ///
    /// Panics if `idle_timeout` is the zero `Duration`.
    pub fn idle_timeout(mut self, idle_timeout: Option<Duration>) -> Builder<M> {
        assert_ne!(idle_timeout, Some(Duration::from_secs(0)), "idle_timeout must be positive");
        self.idle_timeout = idle_timeout;
        self
    }

    /// Sets the connection timeout used by the pool.
    ///
    /// Calls to `Pool::get` will wait this long for a connection to become
    /// available before returning an error.
    ///
    /// Defaults to 30 seconds.
    ///
    /// # Panics
    ///
    /// Panics if `connection_timeout` is the zero duration
    pub fn connection_timeout(mut self, connection_timeout: Duration) -> Builder<M> {
        assert!(
            connection_timeout > Duration::from_secs(0),
            "connection_timeout must be positive"
        );
        self.connection_timeout = connection_timeout;
        self
    }

    /// Sets the handler for errors reported in the pool.
    ///
    /// Defaults to the `LoggingErrorHandler`.
    pub fn error_handler(mut self, error_handler: Box<dyn HandleError<M::Error>>) -> Builder<M> {
        self.error_handler = error_handler;
        self
    }

    /// Sets the handler for events reported by the pool.
    ///
    /// Defaults to the `NopEventHandler`.
    pub fn event_handler(mut self, event_handler: Box<dyn HandleEvent>) -> Builder<M> {
        self.event_handler = event_handler;
        self
    }

    /// Sets the connection customizer used by the pool.
    ///
    /// Defaults to the `NopConnectionCustomizer`.
    pub fn connection_customizer(
        mut self,
        connection_customizer: Box<dyn CustomizeConnection<M::Connection, M::Error>>,
    ) -> Builder<M> {
        self.connection_customizer = connection_customizer;
        self
    }

    // used by tests
    #[allow(dead_code)]
    pub(crate) fn reaper_rate(mut self, reaper_rate: Duration) -> Builder<M> {
        self.reaper_rate = reaper_rate;
        self
    }

    /// Consumes the builder, returning a new, initialized pool.
    ///
    /// It will block until the pool has established its configured minimum
    /// number of connections, or it times out.
    ///
    /// # Errors
    ///
    /// Returns an error if the pool is unable to open its minimum number of
    /// connections.
    ///
    /// # Panics
    ///
    /// Panics if `min_idle` is greater than `max_size`.
    pub fn build(self, manager: M) -> Result<Pool<M>, Error> {
        let pool = self.build_unchecked(manager);
        pool.wait_for_initialization()?;
        Ok(pool)
    }

    /// Consumes the builder, returning a new pool.
    ///
    /// Unlike `build`, this method does not wait for any connections to be
    /// established before returning.
    ///
    /// # Panics
    ///
    /// Panics if `min_idle` is greater than `max_size`.
    pub fn build_unchecked(self, manager: M) -> Pool<M> {
        if let Some(min_idle) = self.min_idle {
            assert!(
                self.max_size >= min_idle,
                "min_idle must be no larger than max_size"
            );
        }

        let thread_pool = match self.thread_pool {
            Some(thread_pool) => thread_pool,
            None => Arc::new(ScheduledThreadPool::with_name("r2d2-worker-{}", 3)),
        };

        let config = Config {
            max_size: self.max_size,
            min_idle: self.min_idle,
            test_on_check_out: self.test_on_check_out,
            max_lifetime: self.max_lifetime,
            idle_timeout: self.idle_timeout,
            connection_timeout: self.connection_timeout,
            error_handler: self.error_handler,
            event_handler: self.event_handler,
            connection_customizer: self.connection_customizer,
            thread_pool,
        };

        Pool::new_inner(config, manager, self.reaper_rate)
    }
}

pub struct Config<C, E> {
    pub max_size: u32,
    pub min_idle: Option<u32>,
    pub test_on_check_out: bool,
    pub max_lifetime: Option<Duration>,
    pub idle_timeout: Option<Duration>,
    pub connection_timeout: Duration,
    pub error_handler: Box<dyn HandleError<E>>,
    pub event_handler: Box<dyn HandleEvent>,
    pub connection_customizer: Box<dyn CustomizeConnection<C, E>>,
    pub thread_pool: Arc<ScheduledThreadPool>,
}

// manual to avoid bounds on C and E
impl<C, E> fmt::Debug for Config<C, E> {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt.debug_struct("Config")
            .field("max_size", &self.max_size)
            .field("min_idle", &self.min_idle)
            .field("test_on_check_out", &self.test_on_check_out)
            .field("max_lifetime", &self.max_lifetime)
            .field("idle_timeout", &self.idle_timeout)
            .field("connection_timeout", &self.connection_timeout)
            .field("error_handler", &self.error_handler)
            .field("event_handler", &self.event_handler)
            .field("connection_customizer", &self.connection_customizer)
            .finish()
    }
}