Repository: tikv/tikv
File: components/cdc/src/channel.rs (Rust)

Code before:
            }
            sink.flush().await?;
            total_event_bytes.inc_by(event_bytes as u64);
            total_resolved_ts_bytes.inc_by(resolved_ts_bytes as u64);
        }
        Ok(())
    }
}

impl Drop for Drain {
    fn drop(&mut self) {
        self.bounded_receiver.close();
        self.unbounded_receiver.close();
        let start = Instant::now();
        let mut total_bytes = 0;
        let mut drain = Box::pin(async move {
            let memory_quota = self.memory_quota.clone();
            let mut drain = self.drain();
            while let Some((_, bytes)) = drain.next().await {
                total_bytes += bytes;
            }
            memory_quota.free(total_bytes);
        });
        block_on(&mut drain);
        let takes = start.saturating_elapsed();
        info!("drop Drain finished, free memory";"takes" => ?takes,
            "freed_bytes" => total_bytes, "inuse_bytes" => self.memory_quota.in_use());
    }
}

#[allow(clippy::result_unit_err)]
pub fn recv_timeout<S, I>(s: &mut S, dur: std::time::Duration) -> Result<Option<I>, ()>
where
    S: Stream<Item = I> + Unpin,
{
    block_on_timeout(s.next(), dur)
}

#[cfg(test)]
mod tests {
    use std::{
        assert_matches::assert_matches,
        sync::{mpsc, Arc},
        time::Duration,
    };

    use futures::executor::block_on;
    use kvproto::cdcpb::{
        ChangeDataEvent, Event, EventEntries, EventRow, Event_oneof_event, ResolvedTs,
    };


Review comment (suggestion):
The log doesn't show which connection the `Drain` belongs to. It's better to add this information.

Code after:
                // Buffer messages and flush them at once.
                let write_flags = WriteFlags::default().buffer_hint(i + 1 != resps_len);
                sink.feed((e, write_flags)).await?;
            }
            sink.flush().await?;
            total_event_bytes.inc_by(event_bytes as u64);
            total_resolved_ts_bytes.inc_by(resolved_ts_bytes as u64);
        }
        Ok(())
    }
}

impl Drop for Drain {
    fn drop(&mut self) {
        self.bounded_receiver.close();
        self.unbounded_receiver.close();
        let start = Instant::now();
        let mut total_bytes = 0;
        let mut drain = Box::pin(async move {
            let conn_id = self.conn_id;
            let memory_quota = self.memory_quota.clone();
            let mut drain = self.drain();
            while let Some((_, bytes)) = drain.next().await {
                total_bytes += bytes;
            }
            memory_quota.free(total_bytes);
            info!("drop Drain finished, free memory"; "conn_id" => ?conn_id,
                "freed_bytes" => total_bytes, "inuse_bytes" => memory_quota.in_use());
        });
        block_on(&mut drain);
        let takes = start.saturating_elapsed();
        if takes >= Duration::from_millis(200) {
            warn!("drop Drain too slow"; "takes" => ?takes);
        }
    }
}

#[allow(clippy::result_unit_err)]
pub fn recv_timeout<S, I>(s: &mut S, dur: std::time::Duration) -> Result<Option<I>, ()>
where
    S: Stream<Item = I> + Unpin,
{
    block_on_timeout(s.next(), dur)
}

#[cfg(test)]
mod tests {
    use std::{
        assert_matches::assert_matches,
        sync::{mpsc, Arc},
        time::Duration,