Repository: launchbadge/sqlx
File: sqlx-core/src/postgres/listener.rs (Rust)

Code before:
            .await?;

        self.channels.push(channel.to_owned());

        Ok(())
    }

    /// Starts listening for notifications on all channels.
    pub async fn listen_all(
        &mut self,
        channels: impl IntoIterator<Item = &str>,
    ) -> Result<(), Error> {
        let beg = self.channels.len();
        self.channels.extend(channels.into_iter().map(|s| s.into()));

        let query = build_listen_all_query(&self.channels[beg..]);
        self.connection().await?.execute(&*query).await?;

        Ok(())
    }

    /// Stops listening for notifications on a channel.
    /// The channel name is quoted here to ensure case sensitivity.
    pub async fn unlisten(&mut self, channel: &str) -> Result<(), Error> {
        self.connection()
            .await?
            .execute(&*format!(r#"UNLISTEN "{}""#, ident(channel)))
            .await?;

        if let Some(pos) = self.channels.iter().position(|s| s == channel) {
            self.channels.remove(pos);
        }

        Ok(())
    }

    /// Stops listening for notifications on all channels.
    pub async fn unlisten_all(&mut self) -> Result<(), Error> {
        self.connection().await?.execute("UNLISTEN *").await?;

        self.channels.clear();

        Ok(())
    }

    #[inline]
    async fn connect_if_needed(&mut self) -> Result<(), Error> {
        if self.connection.is_none() {
            let mut connection = self.pool.acquire().await?;
            connection.stream.notifications = self.buffer_tx.take();


Review comment (suggestion):
This and `unlisten_all()` don't technically need to reconnect right away; if the connection is dead, they just need to delete the channels from the local set.

In fact, this is rather wasteful as `.connection()` will internally reconnect and then execute a `LISTEN` query for every channel and then this will immediately execute `UNLISTEN {channel}`.

Code after:
            .await?;

        self.channels.push(channel.to_owned());

        Ok(())
    }

    /// Starts listening for notifications on all channels.
    pub async fn listen_all(
        &mut self,
        channels: impl IntoIterator<Item = &str>,
    ) -> Result<(), Error> {
        let beg = self.channels.len();
        self.channels.extend(channels.into_iter().map(|s| s.into()));

        let query = build_listen_all_query(&self.channels[beg..]);
        self.connection().await?.execute(&*query).await?;

        Ok(())
    }

    /// Stops listening for notifications on a channel.
    /// The channel name is quoted here to ensure case sensitivity.
    pub async fn unlisten(&mut self, channel: &str) -> Result<(), Error> {
        // use RAW connection and do NOT re-connect automatically, since this is not required for
        // UNLISTEN (we've disconnected anyways)
        if let Some(connection) = self.connection.as_mut() {
            connection
                .execute(&*format!(r#"UNLISTEN "{}""#, ident(channel)))
                .await?;
        }

        if let Some(pos) = self.channels.iter().position(|s| s == channel) {
            self.channels.remove(pos);
        }

        Ok(())
    }

    /// Stops listening for notifications on all channels.
    pub async fn unlisten_all(&mut self) -> Result<(), Error> {
        // use RAW connection and do NOT re-connect automatically, since this is not required for
        // UNLISTEN (we've disconnected anyways)
        if let Some(connection) = self.connection.as_mut() {
            connection.execute("UNLISTEN *").await?;
        }

        self.channels.clear();

        Ok(())
    }