Function stream_yields

datafusion/core/tests/execution/coop.rs:774–820

(
    mut stream: SendableRecordBatchStream,
)

}

async fn stream_yields(
    mut stream: SendableRecordBatchStream,
) -> Result<(), Box<dyn Error>> {
    // Create an independent executor pool
    let child_runtime = Runtime::new()?;

    // Spawn a task that tries to poll the stream
    // The task returns Ready when the stream yielded with either Ready or Pending
    let join_handle = child_runtime.spawn(std::future::poll_fn(move |cx| {
        match stream.poll_next_unpin(cx) {
            Poll::Ready(Some(Ok(_))) => Poll::Ready(Poll::Ready(Ok(()))),
            Poll::Ready(Some(Err(e))) => Poll::Ready(Poll::Ready(Err(e))),
            Poll::Ready(None) => Poll::Ready(Poll::Ready(Ok(()))),
            Poll::Pending => Poll::Ready(Poll::Pending),
        }
    }));

    let abort_handle = join_handle.abort_handle();

    // Now select on the join handle of the task running in the child executor with a timeout
    let yielded = select! {
        result = join_handle => {
            match result {
                Ok(Poll::Pending) => Yielded::ReadyOrPending,
                Ok(Poll::Ready(Ok(_))) => Yielded::ReadyOrPending,
                Ok(Poll::Ready(Err(e))) => Yielded::Err(e),
                Err(_) => Yielded::Err(exec_datafusion_err!("join error")),
            }
        },
        _ = tokio::time::sleep(Duration::from_secs(10)) => {
            Yielded::Timeout
        }
    };

    // Try to abort the poll task and shutdown the child runtime
    abort_handle.abort();
    Handle::current().spawn_blocking(move || {
        child_runtime.shutdown_timeout(Duration::from_secs(5));
    });

    // Finally, check if poll_next yielded
    assert!(
        matches!(yielded, Yielded::ReadyOrPending),
        "Result is not Ready or Pending: {yielded:?}"
    );
    Ok(())
}

async fn query_yields(
    plan: Arc<dyn ExecutionPlan>,