nusadb 0.1.0

//! Async API on the [Tokio](https://tokio.rs) runtime (feature `async`). [`AsyncConnection`] mirrors
//! the sync [`Connection`](crate::Connection) — handshake (trust / SCRAM-SHA-256), simple and
//! extended (parameterized / prepared) queries, transactions — and [`AsyncPool`] pools them. With the
//! `async-tls` feature, connections can run over implicit TLS 1.3 (tokio-rustls).
//!
//! The framing codec, SCRAM math, value decoding, and result assembly are shared with the sync path;
//! only the I/O is async.

use std::collections::VecDeque;
use std::sync::Arc;
use std::time::Duration;

use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _, BufReader};
use tokio::net::TcpStream;

use crate::connection::{
    bind_frame, copy_count, decode_error, decode_notification, describe_portal, execute_portal,
    quote_identifier, quote_literal, CancelHandle, Config, Notification, Prepared, QueryResult,
    ResultAccumulator,
};
use crate::error::{Error, Result};
use crate::protocol::{
    auth, backend, frontend, Frame, Reader, Writer, HEADER_LEN, MAX_FRAME_LEN, PROTOCOL_MAGIC,
    PROTOCOL_MAJOR, PROTOCOL_MINOR, SCRAM_MECHANISM,
};
use crate::scram::Scram;
use crate::value::ToSql;

/// The async byte transport: a Tokio TCP socket or (feature `async-tls`) a tokio-rustls TLS stream.
enum AsyncTransport {
    Plain(TcpStream),
    #[cfg(feature = "async-tls")]
    Tls(Box<tokio_rustls::client::TlsStream<TcpStream>>),
}

impl tokio::io::AsyncRead for AsyncTransport {
    fn poll_read(
        self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
        buf: &mut tokio::io::ReadBuf<'_>,
    ) -> std::task::Poll<std::io::Result<()>> {
        match self.get_mut() {
            Self::Plain(s) => std::pin::Pin::new(s).poll_read(cx, buf),
            #[cfg(feature = "async-tls")]
            Self::Tls(s) => std::pin::Pin::new(s.as_mut()).poll_read(cx, buf),
        }
    }
}

impl tokio::io::AsyncWrite for AsyncTransport {
    fn poll_write(
        self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
        buf: &[u8],
    ) -> std::task::Poll<std::io::Result<usize>> {
        match self.get_mut() {
            Self::Plain(s) => std::pin::Pin::new(s).poll_write(cx, buf),
            #[cfg(feature = "async-tls")]
            Self::Tls(s) => std::pin::Pin::new(s.as_mut()).poll_write(cx, buf),
        }
    }

    fn poll_flush(
        self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<std::io::Result<()>> {
        match self.get_mut() {
            Self::Plain(s) => std::pin::Pin::new(s).poll_flush(cx),
            #[cfg(feature = "async-tls")]
            Self::Tls(s) => std::pin::Pin::new(s.as_mut()).poll_flush(cx),
        }
    }

    fn poll_shutdown(
        self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<std::io::Result<()>> {
        match self.get_mut() {
            Self::Plain(s) => std::pin::Pin::new(s).poll_shutdown(cx),
            #[cfg(feature = "async-tls")]
            Self::Tls(s) => std::pin::Pin::new(s.as_mut()).poll_shutdown(cx),
        }
    }
}

async fn read_frame<R: tokio::io::AsyncRead + Unpin>(r: &mut R) -> Result<Frame> {
    let mut header = [0u8; HEADER_LEN];
    r.read_exact(&mut header).await?;
    let kind = header[0];
    let total = u32::from_be_bytes([header[1], header[2], header[3], header[4]]);
    if total < HEADER_LEN as u32 {
        return Err(Error::Protocol(format!("malformed frame: len {total} < 5")));
    }
    if total > MAX_FRAME_LEN {
        return Err(Error::Protocol(format!(
            "frame too large: {total} > {MAX_FRAME_LEN}"
        )));
    }
    let mut payload = vec![0u8; (total - HEADER_LEN as u32) as usize];
    r.read_exact(&mut payload).await?;
    Ok(Frame { kind, payload })
}

async fn write_frames<W: tokio::io::AsyncWrite + Unpin>(w: &mut W, frames: &[u8]) -> Result<()> {
    w.write_all(frames).await?;
    w.flush().await?;
    Ok(())
}

/// An async connection to a NusaDB server.
pub struct AsyncConnection {
    stream: BufReader<AsyncTransport>,
    addr: (String, u16),
    backend_pid: u32,
    backend_secret: u32,
    stmt_counter: u64,
    notifications: VecDeque<Notification>,
}

impl AsyncConnection {
    /// Open a connection from a `nusadb://…` (or `nusadbs://…` for TLS) URL.
    pub async fn connect(url: &str) -> Result<Self> {
        Self::connect_config(&Config::from_url(url)?).await
    }

    /// Open a connection from an explicit [`Config`].
    pub async fn connect_config(config: &Config) -> Result<Self> {
        let connect = TcpStream::connect((config.host.as_str(), config.port));
        let tcp = match config.connect_timeout {
            Some(timeout) => tokio::time::timeout(timeout, connect).await.map_err(|_| {
                Error::Io(std::io::Error::new(
                    std::io::ErrorKind::TimedOut,
                    "connect timed out",
                ))
            })??,
            None => connect.await?,
        };
        tcp.set_nodelay(true).ok();
        let transport = build_async_transport(tcp, config).await?;
        let mut conn = Self {
            stream: BufReader::new(transport),
            addr: (config.host.clone(), config.port),
            backend_pid: 0,
            backend_secret: 0,
            stmt_counter: 0,
            notifications: VecDeque::new(),
        };
        conn.handshake(config).await?;
        Ok(conn)
    }

    async fn send(&mut self, frames: &[u8]) -> Result<()> {
        write_frames(self.stream.get_mut(), frames).await
    }

    /// Read one frame, transparently buffering asynchronous `NotificationResponse` frames (see the
    /// sync [`Connection::notifications`](crate::Connection::notifications) for the rationale).
    async fn recv(&mut self) -> Result<Frame> {
        loop {
            let frame = read_frame(&mut self.stream).await?;
            if frame.kind == backend::NOTIFICATION_RESPONSE {
                self.notifications.push_back(decode_notification(&frame)?);
                continue;
            }
            return Ok(frame);
        }
    }

    async fn handshake(&mut self, config: &Config) -> Result<()> {
        let startup = Writer::new()
            .u32(PROTOCOL_MAGIC)
            .u16(PROTOCOL_MAJOR)
            .u16(PROTOCOL_MINOR)
            .str(&config.user)
            .str(&config.database)
            .frame(frontend::STARTUP);
        self.send(&startup).await?;

        // Authentication: a single top-level message — either an immediate AuthOk (trust) or an
        // AuthSASL that begins SCRAM. `scram` drives the rest of the SASL exchange and consumes its
        // trailing AuthOk, so exactly one auth frame is read here before the post-auth stream.
        let frame = self.recv().await?;
        match frame.kind {
            backend::AUTH => {
                let mut r = Reader::new(&frame.payload);
                match r.u32()? {
                    auth::OK => {}
                    auth::SASL => self.scram(config).await?,
                    other => {
                        return Err(Error::Protocol(format!("unexpected auth sub-code {other}")))
                    }
                }
            }
            backend::ERROR => return Err(decode_error(&frame.payload)),
            other => {
                return Err(Error::Protocol(format!(
                    "unexpected message {:?} during auth",
                    other as char
                )))
            }
        }

        loop {
            let frame = self.recv().await?;
            match frame.kind {
                backend::BACKEND_KEY => {
                    let mut r = Reader::new(&frame.payload);
                    self.backend_pid = r.u32()?;
                    self.backend_secret = r.u32()?;
                }
                // Runtime parameters announced at startup carry nothing the driver needs.
                backend::PARAMETER_STATUS => {}
                backend::READY => return Ok(()),
                backend::ERROR => return Err(decode_error(&frame.payload)),
                other => {
                    return Err(Error::Protocol(format!(
                        "unexpected message {:?} after auth",
                        other as char
                    )))
                }
            }
        }
    }

    async fn scram(&mut self, config: &Config) -> Result<()> {
        let password = config
            .password
            .as_deref()
            .ok_or_else(|| Error::Auth("server requires a password (SCRAM)".to_owned()))?;
        let (mut state, client_first) = Scram::start(&config.user, password)?;
        let initial = Writer::new()
            .str(SCRAM_MECHANISM)
            .u32(client_first.len() as u32)
            .raw(client_first.as_bytes())
            .frame(frontend::SASL_INITIAL);
        self.send(&initial).await?;

        let frame = self.recv().await?;
        let server_first = expect_auth_sub(frame, auth::SASL_CONTINUE)?;
        let client_final = state.client_final(&server_first)?;
        let response = Writer::new()
            .raw(client_final.as_bytes())
            .frame(frontend::SASL_RESPONSE);
        self.send(&response).await?;

        let frame = self.recv().await?;
        let server_final = expect_auth_sub(frame, auth::SASL_FINAL)?;
        if !state.verify(&server_final) {
            return Err(Error::Auth(
                "server signature verification failed".to_owned(),
            ));
        }
        let frame = self.recv().await?;
        let _ = expect_auth_sub(frame, auth::OK)?;
        Ok(())
    }

    async fn collect(&mut self) -> Result<QueryResult> {
        let mut acc = ResultAccumulator::new();
        loop {
            let frame = self.recv().await?;
            if let Some(result) = acc.push(&frame)? {
                return Ok(result);
            }
        }
    }

    /// Run a simple (text) query with no parameters.
    pub async fn query(&mut self, sql: &str) -> Result<QueryResult> {
        let frame = Writer::new().str(sql).frame(frontend::QUERY);
        self.send(&frame).await?;
        self.collect().await
    }

    /// Run a query with positional `$1..$n` parameters (extended protocol, one-shot unnamed statement).
    pub async fn query_params(&mut self, sql: &str, params: &[&dyn ToSql]) -> Result<QueryResult> {
        let encoded: Vec<Option<Vec<u8>>> = params.iter().map(|p| p.to_sql_text()).collect();
        let mut batch = Writer::new().str("").str(sql).u16(0).frame(frontend::PARSE);
        batch.extend_from_slice(&bind_frame("", "", &encoded));
        batch.extend_from_slice(&describe_portal(""));
        batch.extend_from_slice(&execute_portal("", 0));
        batch.extend_from_slice(&Writer::new().frame(frontend::SYNC));
        self.send(&batch).await?;
        self.collect().await
    }

    /// Prepare a statement for repeated execution. Reuse it with [`Self::query_prepared`].
    pub async fn prepare(&mut self, sql: &str) -> Result<Prepared> {
        self.stmt_counter += 1;
        let name = format!("nusadb_stmt_{}", self.stmt_counter);
        let mut batch = Writer::new()
            .str(&name)
            .str(sql)
            .u16(0)
            .frame(frontend::PARSE);
        batch.extend_from_slice(&Writer::new().frame(frontend::SYNC));
        self.send(&batch).await?;
        self.collect().await?;
        Ok(Prepared::new(name))
    }

    /// Execute a previously [`Self::prepare`]d statement with parameters.
    pub async fn query_prepared(
        &mut self,
        stmt: &Prepared,
        params: &[&dyn ToSql],
    ) -> Result<QueryResult> {
        let encoded: Vec<Option<Vec<u8>>> = params.iter().map(|p| p.to_sql_text()).collect();
        let mut batch = bind_frame("", stmt.name(), &encoded);
        batch.extend_from_slice(&describe_portal(""));
        batch.extend_from_slice(&execute_portal("", 0));
        batch.extend_from_slice(&Writer::new().frame(frontend::SYNC));
        self.send(&batch).await?;
        self.collect().await
    }

    /// Run a statement and return its affected/returned row count.
    pub async fn execute(&mut self, sql: &str) -> Result<u64> {
        Ok(self.query(sql).await?.affected())
    }

    /// Open a transaction (`BEGIN`).
    pub async fn begin(&mut self) -> Result<()> {
        self.query("BEGIN").await.map(|_| ())
    }

    /// Commit the open transaction (`COMMIT`).
    pub async fn commit(&mut self) -> Result<()> {
        self.query("COMMIT").await.map(|_| ())
    }

    /// Roll back the open transaction (`ROLLBACK`).
    pub async fn rollback(&mut self) -> Result<()> {
        self.query("ROLLBACK").await.map(|_| ())
    }

    /// Establish a savepoint in the open transaction (`SAVEPOINT name`) — the async mirror of
    /// [`Connection::savepoint`](crate::Connection::savepoint).
    pub async fn savepoint(&mut self, name: &str) -> Result<()> {
        self.query(&format!("SAVEPOINT {}", quote_identifier(name)))
            .await
            .map(|_| ())
    }

    /// Roll back to a savepoint (`ROLLBACK TO SAVEPOINT name`); the transaction stays open.
    pub async fn rollback_to_savepoint(&mut self, name: &str) -> Result<()> {
        self.query(&format!("ROLLBACK TO SAVEPOINT {}", quote_identifier(name)))
            .await
            .map(|_| ())
    }

    /// Release (forget) a savepoint (`RELEASE SAVEPOINT name`), keeping its work.
    pub async fn release_savepoint(&mut self, name: &str) -> Result<()> {
        self.query(&format!("RELEASE SAVEPOINT {}", quote_identifier(name)))
            .await
            .map(|_| ())
    }

    /// Subscribe to asynchronous notifications on `channel` (`LISTEN channel`) — the async mirror of
    /// [`Connection::listen`](crate::Connection::listen).
    pub async fn listen(&mut self, channel: &str) -> Result<()> {
        self.query(&format!("LISTEN {}", quote_identifier(channel)))
            .await
            .map(|_| ())
    }

    /// Stop listening on `channel` (`UNLISTEN channel`).
    pub async fn unlisten(&mut self, channel: &str) -> Result<()> {
        self.query(&format!("UNLISTEN {}", quote_identifier(channel)))
            .await
            .map(|_| ())
    }

    /// Stop listening on every channel (`UNLISTEN *`).
    pub async fn unlisten_all(&mut self) -> Result<()> {
        self.query("UNLISTEN *").await.map(|_| ())
    }

    /// Send a notification on `channel` with an optional `payload` (`NOTIFY channel[, 'payload']`).
    pub async fn notify(&mut self, channel: &str, payload: Option<&str>) -> Result<()> {
        let sql = match payload {
            Some(p) => format!("NOTIFY {}, {}", quote_identifier(channel), quote_literal(p)),
            None => format!("NOTIFY {}", quote_identifier(channel)),
        };
        self.query(&sql).await.map(|_| ())
    }

    /// Drain the notifications already received (buffered while reading other responses). Does not
    /// touch the socket.
    pub fn notifications(&mut self) -> Vec<Notification> {
        self.notifications.drain(..).collect()
    }

    /// Wait up to `timeout` for the next notification (a buffered one is returned immediately;
    /// `None` on timeout). `timeout` of `None` waits indefinitely. Only meaningful after
    /// [`AsyncConnection::listen`].
    pub async fn poll_notification(
        &mut self,
        timeout: Option<Duration>,
    ) -> Result<Option<Notification>> {
        if let Some(note) = self.notifications.pop_front() {
            return Ok(Some(note));
        }
        let read = read_frame(&mut self.stream);
        let frame = match timeout {
            Some(dur) => match tokio::time::timeout(dur, read).await {
                Ok(res) => res?,
                Err(_) => return Ok(None),
            },
            None => read.await?,
        };
        if frame.kind == backend::NOTIFICATION_RESPONSE {
            Ok(Some(decode_notification(&frame)?))
        } else {
            Err(Error::Protocol(format!(
                "unexpected message {:?} while polling for notifications",
                frame.kind as char
            )))
        }
    }

    /// Bulk-load a table via `COPY <table> FROM STDIN` (§12.1) — the async mirror of
    /// [`Connection::copy_in`](crate::Connection::copy_in). Streams `source`'s text-format bytes
    /// (tab-delimited, `\N` for NULL) as `CopyData` frames and returns the number of rows loaded.
    ///
    /// On a read error from `source` the driver sends `CopyFail`, drains to `ReadyForQuery`, and
    /// returns the I/O error, leaving the connection usable. A refused COPY surfaces as
    /// [`Error::Server`].
    pub async fn copy_in<R: tokio::io::AsyncRead + Unpin>(
        &mut self,
        sql: &str,
        source: &mut R,
    ) -> Result<u64> {
        self.send(&Writer::new().str(sql).frame(frontend::QUERY))
            .await?;
        self.await_copy_start(backend::COPY_IN).await?;
        let mut buf = vec![0u8; 64 * 1024];
        loop {
            match source.read(&mut buf).await {
                Ok(0) => break,
                Ok(n) => {
                    self.send(&Writer::new().raw(&buf[..n]).frame(frontend::COPY_DATA))
                        .await?;
                }
                Err(e) => {
                    let _ = self
                        .send(
                            &Writer::new()
                                .str("nusadb: client read error during COPY")
                                .frame(frontend::COPY_FAIL),
                        )
                        .await;
                    self.drain_to_ready().await;
                    return Err(Error::Io(e));
                }
            }
        }
        self.send(&Writer::new().frame(frontend::COPY_DONE)).await?;
        self.finish_copy().await
    }

    /// Bulk-export a table via `COPY <table> TO STDOUT` (§12.2) — the async mirror of
    /// [`Connection::copy_out`](crate::Connection::copy_out). Writes every `CopyData` chunk into
    /// `sink` (raw text-format bytes) and returns the number of rows exported.
    pub async fn copy_out<W: tokio::io::AsyncWrite + Unpin>(
        &mut self,
        sql: &str,
        sink: &mut W,
    ) -> Result<u64> {
        self.send(&Writer::new().str(sql).frame(frontend::QUERY))
            .await?;
        self.await_copy_start(backend::COPY_OUT).await?;
        loop {
            let frame = self.recv().await?;
            match frame.kind {
                backend::COPY_DATA => sink.write_all(&frame.payload).await?,
                backend::COPY_DONE => break,
                backend::ERROR => {
                    let err = decode_error(&frame.payload);
                    self.drain_to_ready().await;
                    return Err(err);
                }
                other => {
                    return Err(Error::Protocol(format!(
                        "unexpected message {:?} during COPY TO STDOUT",
                        other as char
                    )))
                }
            }
        }
        self.finish_copy().await
    }

    /// Read frames until the expected `CopyInResponse`/`CopyOutResponse` arrives, surfacing a refused
    /// COPY (`Error` then `ReadyForQuery`) with the connection left ready.
    async fn await_copy_start(&mut self, expected: u8) -> Result<()> {
        let mut err: Option<Error> = None;
        loop {
            let frame = self.recv().await?;
            match frame.kind {
                k if k == expected => return Ok(()),
                backend::ERROR => err = Some(decode_error(&frame.payload)),
                backend::READY => {
                    return Err(err.unwrap_or_else(|| {
                        Error::Protocol("COPY: server did not start the copy".to_owned())
                    }))
                }
                other => {
                    return Err(Error::Protocol(format!(
                        "unexpected message {:?} starting COPY",
                        other as char
                    )))
                }
            }
        }
    }

    /// Drain the tail of a COPY exchange to `ReadyForQuery`, returning the `COPY <n>` row count.
    async fn finish_copy(&mut self) -> Result<u64> {
        let mut tag = String::new();
        let mut err: Option<Error> = None;
        loop {
            let frame = self.recv().await?;
            match frame.kind {
                backend::COMMAND_COMPLETE => tag = Reader::new(&frame.payload).str()?,
                backend::COPY_DONE => {}
                backend::ERROR => err = Some(decode_error(&frame.payload)),
                backend::READY => {
                    if let Some(e) = err {
                        return Err(e);
                    }
                    return Ok(copy_count(&tag));
                }
                other => {
                    return Err(Error::Protocol(format!(
                        "unexpected message {:?} finishing COPY",
                        other as char
                    )))
                }
            }
        }
    }

    /// Best-effort resync: read and discard frames until `ReadyForQuery`.
    async fn drain_to_ready(&mut self) {
        while let Ok(frame) = self.recv().await {
            if frame.kind == backend::READY {
                break;
            }
        }
    }

    /// This connection's `(pid, secret)` for out-of-band cancellation (§13).
    #[must_use]
    pub fn cancel_handle(&self) -> CancelHandle {
        CancelHandle::new(
            self.addr.0.clone(),
            self.addr.1,
            self.backend_pid,
            self.backend_secret,
        )
    }

    /// Politely close the connection (`Terminate`).
    pub async fn close(mut self) {
        let _ = self.send(&Writer::new().frame(frontend::TERMINATE)).await;
    }
}

fn expect_auth_sub(frame: Frame, expect: u32) -> Result<String> {
    if frame.kind == backend::ERROR {
        return Err(decode_error(&frame.payload));
    }
    if frame.kind != backend::AUTH {
        return Err(Error::Protocol(format!(
            "expected auth message, got {:?}",
            frame.kind as char
        )));
    }
    let mut r = Reader::new(&frame.payload);
    let sub = r.u32()?;
    if sub != expect {
        return Err(Error::Protocol(format!(
            "expected auth sub-code {expect}, got {sub}"
        )));
    }
    let rest = r.rest()?;
    String::from_utf8(rest.to_vec())
        .map_err(|_| Error::Protocol("invalid UTF-8 in SASL message".to_owned()))
}

async fn build_async_transport(tcp: TcpStream, config: &Config) -> Result<AsyncTransport> {
    let Some(tls) = &config.tls else {
        return Ok(AsyncTransport::Plain(tcp));
    };
    #[cfg(feature = "async-tls")]
    {
        let client_config = crate::tls::client_config(tls)?;
        let server_name = crate::tls::server_name(&config.host)?;
        let connector = tokio_rustls::TlsConnector::from(client_config);
        let stream = connector
            .connect(server_name, tcp)
            .await
            .map_err(|e| Error::Protocol(format!("TLS handshake failed: {e}")))?;
        Ok(AsyncTransport::Tls(Box::new(stream)))
    }
    #[cfg(not(feature = "async-tls"))]
    {
        let _ = tls;
        Err(Error::Config(
            "TLS requested but the `async-tls` feature is not enabled".to_owned(),
        ))
    }
}

/// An async connection pool: [`AsyncPool::get`] awaits a free connection (up to `max_size`), opening
/// one on demand, and returns it to the pool on drop.
#[derive(Clone)]
pub struct AsyncPool {
    config: Arc<Config>,
    sem: Arc<tokio::sync::Semaphore>,
    idle: Arc<std::sync::Mutex<Vec<AsyncConnection>>>,
}

impl AsyncPool {
    /// Create a pool of up to `max_size` connections (created lazily). `max_size` is clamped to >= 1.
    pub fn new(config: Config, max_size: usize) -> Result<Self> {
        let max_size = max_size.max(1);
        Ok(Self {
            config: Arc::new(config),
            sem: Arc::new(tokio::sync::Semaphore::new(max_size)),
            idle: Arc::new(std::sync::Mutex::new(Vec::with_capacity(max_size))),
        })
    }

    /// Check out a connection, awaiting a free slot when the pool is saturated.
    pub async fn get(&self) -> Result<AsyncPooled> {
        let permit = Arc::clone(&self.sem)
            .acquire_owned()
            .await
            .map_err(|_| Error::Pool("pool closed".to_owned()))?;
        let pooled = self.idle.lock().ok().and_then(|mut idle| idle.pop());
        let conn = match pooled {
            Some(c) => c,
            None => AsyncConnection::connect_config(&self.config).await?,
        };
        Ok(AsyncPooled {
            conn: Some(conn),
            idle: Arc::clone(&self.idle),
            _permit: permit,
        })
    }
}

/// A connection checked out of an [`AsyncPool`]; returns to the pool on drop. Derefs to
/// [`AsyncConnection`], so call its async methods directly.
pub struct AsyncPooled {
    conn: Option<AsyncConnection>,
    idle: Arc<std::sync::Mutex<Vec<AsyncConnection>>>,
    _permit: tokio::sync::OwnedSemaphorePermit,
}

impl std::ops::Deref for AsyncPooled {
    type Target = AsyncConnection;
    fn deref(&self) -> &AsyncConnection {
        self.conn.as_ref().expect("connection present until drop")
    }
}

impl std::ops::DerefMut for AsyncPooled {
    fn deref_mut(&mut self) -> &mut AsyncConnection {
        self.conn.as_mut().expect("connection present until drop")
    }
}

impl Drop for AsyncPooled {
    fn drop(&mut self) {
        if let (Some(conn), Ok(mut idle)) = (self.conn.take(), self.idle.lock()) {
            idle.push(conn);
        }
    }
}