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

Code before:
        let mut e = Event::default();
        e.region_id = 233;
        {
            let memory_quota = Arc::new(MemoryQuota::new(1024));
            let (mut tx, mut rx) = channel(10, memory_quota);

            let truncated = Arc::new(AtomicBool::new(false));
            let event = CdcEvent::Event(e.clone());
            let size = event.size() as usize;
            let _ = block_on(tx.send_all(vec![event], truncated));

            let memory_quota = rx.memory_quota.clone();
            let mut drain = rx.drain();
            assert_matches!(block_on(drain.next()), Some((CdcEvent::Event(_), _)));
            assert_eq!(memory_quota.in_use(), size);
        }
        {
            let memory_quota = Arc::new(MemoryQuota::new(1024));
            let (mut tx, mut rx) = channel(10, memory_quota);

            let truncated = Arc::new(AtomicBool::new(true));
            let _ = block_on(tx.send_all(vec![CdcEvent::Event(e)], truncated));

            let memory_quota = rx.memory_quota.clone();
            let mut drain = rx.drain();
            block_on_timeout(drain.next(), Duration::from_millis(100)).unwrap_err();
            assert_eq!(memory_quota.in_use(), 0);
        }
    }
    
    #[test]
    fn test_barrier() {
        let force_send = false;
        let (mut send, mut rx) = new_test_channel(10, usize::MAX, force_send);
        send(CdcEvent::Event(Default::default())).unwrap();
        let (btx1, brx1) = mpsc::channel();
        send(CdcEvent::Barrier(Some(Box::new(move |()| {
            btx1.send(()).unwrap();
        }))))
        .unwrap();
        send(CdcEvent::ResolvedTs(Default::default())).unwrap();
        let (btx2, brx2) = mpsc::channel();
        send(CdcEvent::Barrier(Some(Box::new(move |()| {
            btx2.send(()).unwrap();
        }))))
        .unwrap();

        let mut drain = rx.drain();
        brx1.recv_timeout(Duration::from_millis(100)).unwrap_err();
        brx2.recv_timeout(Duration::from_millis(100)).unwrap_err();
        assert_matches!(block_on(drain.next()), Some((CdcEvent::Event(_), _)));

Review comment (suggestion):
since the truncate event is set to None, so call the `block_on` would block forever. use `block_on_timeout` as workaround

Code after:
            let (mut tx, mut rx) = channel(ConnId::default(), 10, memory_quota);

            let truncated = Arc::new(AtomicBool::new(false));
            let event = CdcEvent::Event(e.clone());
            let size = event.size() as usize;
            let _ = block_on(tx.send_all(vec![event], truncated));

            let memory_quota = rx.memory_quota.clone();
            let mut drain = rx.drain();
            assert_matches!(block_on(drain.next()), Some((CdcEvent::Event(_), _)));
            assert_eq!(memory_quota.in_use(), size);
        }
        {
            let memory_quota = Arc::new(MemoryQuota::new(1024));
            let (mut tx, mut rx) = channel(ConnId::default(), 10, memory_quota);

            let truncated = Arc::new(AtomicBool::new(true));
            let _ = block_on(tx.send_all(vec![CdcEvent::Event(e)], truncated));

            let memory_quota = rx.memory_quota.clone();
            let mut drain = rx.drain();
            recv_timeout(&mut drain, Duration::from_millis(100)).unwrap_err();
            assert_eq!(memory_quota.in_use(), 0);
        }
    }

    #[test]
    fn test_barrier() {
        let force_send = false;
        let (mut send, mut rx) = new_test_channel(10, usize::MAX, force_send);
        send(CdcEvent::Event(Default::default())).unwrap();
        let (btx1, brx1) = mpsc::channel();
        send(CdcEvent::Barrier(Some(Box::new(move |()| {
            btx1.send(()).unwrap();
        }))))
        .unwrap();
        send(CdcEvent::ResolvedTs(Default::default())).unwrap();
        let (btx2, brx2) = mpsc::channel();
        send(CdcEvent::Barrier(Some(Box::new(move |()| {
            btx2.send(()).unwrap();
        }))))
        .unwrap();

        let mut drain = rx.drain();
        brx1.recv_timeout(Duration::from_millis(100)).unwrap_err();
        brx2.recv_timeout(Duration::from_millis(100)).unwrap_err();
        assert_matches!(block_on(drain.next()), Some((CdcEvent::Event(_), _)));
        brx1.recv_timeout(Duration::from_millis(100)).unwrap_err();
        brx2.recv_timeout(Duration::from_millis(100)).unwrap_err();
        assert_matches!(block_on(drain.next()), Some((CdcEvent::Barrier(_), _)));
        brx1.recv_timeout(Duration::from_millis(100)).unwrap();