use std::collections::VecDeque;
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};
use asupersync::Cx;
use futures_core::Stream;
use oracledb_protocol::thin::{ColumnMetadata, ExecuteOptions, QueryResult, QueryValue};
use crate::request::{DeadlineExpiry, QueryDeadline};
use crate::rows::first_cursor_from_result;
use crate::{
columns_require_define, CancelHandle, Connection, Cursor, Error, Params, Query, Result,
};
type OwnedRow = Vec<Option<QueryValue>>;
type FetchOutcome = (Connection, Result<QueryResult>);
type FetchFuture = Pin<Box<dyn Future<Output = FetchOutcome>>>;
enum OwnedRowStreamState {
Buffered(VecDeque<OwnedRow>),
Fetching {
future: FetchFuture,
cancel_handle: CancelHandle,
},
Done,
Poisoned,
}
#[must_use = "an OwnedRowStream holds the connection until drained or recovered with into_connection().await"]
pub struct OwnedRowStream {
connection: Option<Connection>,
cx: Cx,
columns: Arc<[ColumnMetadata]>,
cursor: Option<Cursor>,
cursor_id: u32,
arraysize: u32,
deadline: QueryDeadline,
more_rows: bool,
previous_row: Option<OwnedRow>,
state: OwnedRowStreamState,
}
impl OwnedRowStream {
fn from_first_page(
connection: Option<Connection>,
cx: Cx,
arraysize: u32,
deadline: QueryDeadline,
result: QueryResult,
) -> Self {
let cursor_id = result.cursor_id;
let more_rows = result.more_rows;
let cursor = first_cursor_from_result(&result);
let columns: Arc<[ColumnMetadata]> = Arc::from(result.columns.into_boxed_slice());
let batch: VecDeque<OwnedRow> = result.rows.into_iter().collect();
let previous_row = batch.back().cloned();
Self {
connection,
cx,
columns,
cursor,
cursor_id,
arraysize,
deadline,
more_rows,
previous_row,
state: OwnedRowStreamState::Buffered(batch),
}
}
pub fn columns(&self) -> &[ColumnMetadata] {
&self.columns
}
pub fn cursor(&self) -> Option<&Cursor> {
self.cursor.as_ref()
}
pub async fn into_connection(mut self) -> Result<Connection> {
let state = std::mem::replace(&mut self.state, OwnedRowStreamState::Done);
let mut connection = match state {
OwnedRowStreamState::Buffered(_) | OwnedRowStreamState::Done => {
self.connection.take().ok_or_else(|| {
Error::ConnectionClosed(
"owned row stream has no connection to recover".to_string(),
)
})?
}
OwnedRowStreamState::Fetching {
future,
mut cancel_handle,
} => {
cancel_handle.cancel_for_recovery()?;
let (mut connection, _) = future.await;
connection.drain_cancel_response().await?;
connection
}
OwnedRowStreamState::Poisoned => {
return Err(Error::ConnectionClosed(
"owned row stream has no connection to recover".to_string(),
));
}
};
if self.cursor_id != 0 {
connection.release_cursor(self.cursor_id);
self.cursor_id = 0;
}
self.more_rows = false;
Ok(connection)
}
fn release_cursor_if_open(&mut self) {
if self.cursor_id != 0 {
if let Some(connection) = &mut self.connection {
connection.release_cursor(self.cursor_id);
}
self.cursor_id = 0;
self.more_rows = false;
}
}
fn apply_page(&mut self, result: QueryResult) {
if self.cursor.is_none() {
if let Some(cursor) = first_cursor_from_result(&result) {
self.cursor = Some(cursor);
}
}
if result.cursor_id != 0 {
self.cursor_id = result.cursor_id;
}
self.more_rows = result.more_rows;
let columns = result.columns;
if !columns.is_empty() {
self.columns = Arc::from(columns.into_boxed_slice());
}
let batch: VecDeque<OwnedRow> = result.rows.into_iter().collect();
if let Some(last) = batch.back() {
self.previous_row = Some(last.clone());
}
self.state = OwnedRowStreamState::Buffered(batch);
}
}
fn build_fetch_future(
mut connection: Connection,
cx: Cx,
cursor_id: u32,
arraysize: u32,
deadline: QueryDeadline,
columns: Arc<[ColumnMetadata]>,
previous_row: Option<OwnedRow>,
) -> FetchFuture {
Box::pin(async move {
let outcome = match deadline
.run(connection.fetch_rows_with_columns(
&cx,
cursor_id,
arraysize,
&columns,
previous_row.as_deref(),
))
.await
{
Ok(result) => connection.close_cursor_on_error(cursor_id, result),
Err(DeadlineExpiry::BeforeStart) => {
connection.release_cursor(cursor_id);
connection.reject_before_operation_start(&cx, deadline.timeout_ms())
}
Err(DeadlineExpiry::InFlight) => {
let recovered = connection
.recover_from_call_timeout(&cx, deadline.timeout_ms())
.await;
connection.close_cursor_on_error(cursor_id, recovered)
}
};
(connection, outcome)
})
}
impl Stream for OwnedRowStream {
type Item = Result<OwnedRow>;
fn poll_next(self: Pin<&mut Self>, task_cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
let this = self.get_mut();
loop {
match std::mem::replace(&mut this.state, OwnedRowStreamState::Poisoned) {
OwnedRowStreamState::Buffered(mut batch) => {
if let Some(row) = batch.pop_front() {
this.state = OwnedRowStreamState::Buffered(batch);
return Poll::Ready(Some(Ok(row)));
}
if !this.more_rows || this.cursor_id == 0 {
this.release_cursor_if_open();
this.state = OwnedRowStreamState::Done;
return Poll::Ready(None);
}
let Some(connection) = this.connection.take() else {
this.state = OwnedRowStreamState::Poisoned;
return Poll::Ready(Some(Err(Error::ConnectionClosed(
"owned row stream needs another page but holds no connection"
.to_string(),
))));
};
let cancel_handle = match connection.cancel_handle() {
Ok(handle) => handle,
Err(err) => {
this.connection = Some(connection);
this.state = OwnedRowStreamState::Done;
return Poll::Ready(Some(Err(err)));
}
};
let future = build_fetch_future(
connection,
this.cx.clone(),
this.cursor_id,
this.arraysize,
this.deadline,
Arc::clone(&this.columns),
this.previous_row.clone(),
);
this.state = OwnedRowStreamState::Fetching {
future,
cancel_handle,
};
}
OwnedRowStreamState::Fetching {
mut future,
cancel_handle,
} => {
match future.as_mut().poll(task_cx) {
Poll::Pending => {
this.state = OwnedRowStreamState::Fetching {
future,
cancel_handle,
};
return Poll::Pending;
}
Poll::Ready((connection, outcome)) => {
this.connection = Some(connection);
match outcome {
Ok(result) => {
this.apply_page(result);
}
Err(err) => {
this.cursor_id = 0;
this.more_rows = false;
this.state = OwnedRowStreamState::Done;
return Poll::Ready(Some(Err(err)));
}
}
}
}
}
OwnedRowStreamState::Done => {
this.state = OwnedRowStreamState::Done;
return Poll::Ready(None);
}
OwnedRowStreamState::Poisoned => {
this.state = OwnedRowStreamState::Poisoned;
return Poll::Ready(None);
}
}
}
}
}
impl Drop for OwnedRowStream {
fn drop(&mut self) {
self.release_cursor_if_open();
}
}
impl std::fmt::Debug for OwnedRowStream {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let state = match &self.state {
OwnedRowStreamState::Buffered(batch) => {
return f
.debug_struct("OwnedRowStream")
.field("has_connection", &self.connection.is_some())
.field("cursor_id", &self.cursor_id)
.field("more_rows", &self.more_rows)
.field("buffered_rows", &batch.len())
.field("columns", &self.columns.len())
.finish();
}
OwnedRowStreamState::Fetching { .. } => "Fetching",
OwnedRowStreamState::Done => "Done",
OwnedRowStreamState::Poisoned => "Poisoned",
};
f.debug_struct("OwnedRowStream")
.field("has_connection", &self.connection.is_some())
.field("cursor_id", &self.cursor_id)
.field("more_rows", &self.more_rows)
.field("state", &state)
.field("columns", &self.columns.len())
.finish()
}
}
impl Connection {
pub async fn into_row_stream<'q>(self, cx: &Cx, query: Query<'q>) -> Result<OwnedRowStream> {
let Query {
sql,
params,
arraysize,
prefetch,
prefetch_set: _,
materialize_lobs,
scrollable: _,
timeout,
} = query;
let mut connection = self;
let sql_owned = sql.into_owned();
let binds = crate::sql_convert::resolve_params(&sql_owned, params)?;
let bind_rows = if binds.is_empty() {
Vec::new()
} else {
vec![binds]
};
let deadline = QueryDeadline::new(cx, timeout);
let mut result = match deadline
.run(connection.execute_query_with_bind_rows_and_options_core(
cx,
&sql_owned,
prefetch,
&bind_rows,
ExecuteOptions::default(),
))
.await
{
Ok(result) => result?,
Err(DeadlineExpiry::BeforeStart) => {
return connection.reject_before_operation_start(cx, deadline.timeout_ms());
}
Err(DeadlineExpiry::InFlight) => {
return connection
.recover_from_call_timeout(cx, deadline.timeout_ms())
.await
}
};
if materialize_lobs
&& columns_require_define(&result.columns)
&& result.cursor_id != 0
&& result.rows.is_empty()
{
let cursor_id = result.cursor_id;
let columns = result.columns.clone();
let fetched = match deadline
.run(connection.define_and_fetch_rows_with_columns(
cx,
cursor_id,
prefetch.max(1),
&columns,
None,
))
.await
{
Ok(result) => connection.close_cursor_on_error(cursor_id, result)?,
Err(DeadlineExpiry::BeforeStart) => {
connection.release_cursor(cursor_id);
return connection.reject_before_operation_start(cx, deadline.timeout_ms());
}
Err(DeadlineExpiry::InFlight) => {
let recovered = connection
.recover_from_call_timeout(cx, deadline.timeout_ms())
.await;
connection.close_cursor(cursor_id);
return recovered;
}
};
result.rows = fetched.rows;
result.more_rows = fetched.more_rows;
if !fetched.columns.is_empty() {
result.columns = fetched.columns;
}
if result.cursor_id == 0 {
result.cursor_id = cursor_id;
}
}
Ok(OwnedRowStream::from_first_page(
Some(connection),
cx.clone(),
arraysize.get(),
deadline,
result,
))
}
pub async fn into_query_stream<'p>(
self,
cx: &Cx,
sql: &str,
params: impl Into<Params<'p>>,
) -> Result<OwnedRowStream> {
self.into_row_stream(cx, Query::owned_sql(sql.to_string()).bind(params))
.await
}
}
#[cfg(test)]
mod tests {
use std::future::poll_fn;
use std::io::{Read, Write};
use std::pin::{pin, Pin};
use std::task::{Context, Poll, Waker};
use std::thread;
use std::time::Duration;
use asupersync::net::TcpStream;
use asupersync::{CancelKind, Cx};
use oracledb_protocol::thin::{
ColumnMetadata, QueryResult, QueryValue, TNS_DATA_FLAGS_END_OF_RESPONSE,
TNS_PACKET_TYPE_DATA,
};
use oracledb_protocol::wire::{encode_packet, PacketLengthWidth};
use super::*;
fn cols(names: &[&str]) -> Vec<ColumnMetadata> {
names.iter().map(|n| ColumnMetadata::new(*n, 0)).collect()
}
fn text_row(values: &[&str]) -> OwnedRow {
values
.iter()
.map(|v| Some(QueryValue::Text((*v).to_string())))
.collect()
}
fn with_cx<R>(body: impl FnOnce(Cx) -> R) -> R {
let runtime = crate::build_io_runtime().expect("io runtime builds");
runtime.block_on(async {
let cx = Cx::current().expect("block_on installs an ambient Cx");
body(cx)
})
}
fn drain_buffered(stream: &mut OwnedRowStream) -> Vec<Result<OwnedRow>> {
let mut out = Vec::new();
let mut stream = pin!(stream);
let mut task_cx = Context::from_waker(Waker::noop());
loop {
match stream.as_mut().poll_next(&mut task_cx) {
Poll::Ready(Some(item)) => out.push(item),
Poll::Ready(None) => break,
Poll::Pending => panic!("buffered-only stream must never park"),
}
}
out
}
fn offline_stream(cx: Cx, result: QueryResult) -> OwnedRowStream {
let deadline = QueryDeadline::new(&cx, None);
OwnedRowStream::from_first_page(None, cx, 100, deadline, result)
}
fn read_packet(socket: &mut std::net::TcpStream) -> std::io::Result<(u8, Vec<u8>)> {
let mut header = [0u8; 8];
socket.read_exact(&mut header)?;
let declared = u32::from_be_bytes([header[0], header[1], header[2], header[3]]) as usize;
let mut body = vec![0u8; declared - header.len()];
socket.read_exact(&mut body)?;
Ok((header[4], body))
}
fn data_packet(message: &[u8]) -> Vec<u8> {
encode_packet(
TNS_PACKET_TYPE_DATA,
0,
Some(TNS_DATA_FLAGS_END_OF_RESPONSE),
message,
PacketLengthWidth::Large32,
)
.expect("test DATA packet encodes")
}
fn marker_packet(marker_type: u8) -> Vec<u8> {
encode_packet(
crate::TNS_PACKET_TYPE_MARKER,
0,
None,
&[1, 0, marker_type],
PacketLengthWidth::Large32,
)
.expect("test MARKER packet encodes")
}
#[test]
fn continuation_fetch_does_not_rearm_the_query_timeout() -> Result<()> {
const SQL: &str = "select value from drvqa_cursor_cleanup";
const QUERY_TIMEOUT: Duration = Duration::from_millis(40);
const INFLIGHT_BODY: &[u8] = &[0xd1, 0xa1, 0xb1];
const CANCEL_ERROR: &[u8] = &[0x04, 0x01, 0x0d];
let listener = std::net::TcpListener::bind("127.0.0.1:0")?;
let addr = listener.local_addr()?;
let server = thread::spawn(move || -> std::io::Result<bool> {
let (mut socket, _) = listener.accept()?;
socket.set_read_timeout(Some(Duration::from_millis(500)))?;
let (first_type, first_body) = match read_packet(&mut socket) {
Ok(packet) => packet,
Err(err)
if matches!(
err.kind(),
std::io::ErrorKind::UnexpectedEof
| std::io::ErrorKind::WouldBlock
| std::io::ErrorKind::TimedOut
) =>
{
return Ok(false);
}
Err(err) => return Err(err),
};
let saw_fetch = first_type == TNS_PACKET_TYPE_DATA;
let break_body = if saw_fetch {
let (packet_type, body) = read_packet(&mut socket)?;
assert_eq!(packet_type, crate::TNS_PACKET_TYPE_MARKER);
body
} else {
assert_eq!(first_type, crate::TNS_PACKET_TYPE_MARKER);
first_body
};
assert_eq!(break_body, [1, 0, crate::TNS_MARKER_TYPE_BREAK]);
socket.write_all(&data_packet(INFLIGHT_BODY))?;
socket.write_all(&marker_packet(crate::TNS_MARKER_TYPE_BREAK))?;
let (reset_type, reset_body) = read_packet(&mut socket)?;
assert_eq!(reset_type, crate::TNS_PACKET_TYPE_MARKER);
assert_eq!(reset_body, [1, 0, crate::TNS_MARKER_TYPE_RESET]);
socket.write_all(&marker_packet(crate::TNS_MARKER_TYPE_RESET))?;
socket.write_all(&data_packet(CANCEL_ERROR))?;
socket.flush()?;
Ok(saw_fetch)
});
let runtime = crate::build_io_runtime()?;
runtime.block_on(async {
let cx = Cx::current().expect("runtime installs an ambient Cx");
let socket = TcpStream::connect(addr).await?;
let (read, write) = crate::transport::plain_split(socket);
let mut connection = crate::tests::loopback_connection(read, write);
connection.statement_cache_put(SQL, 42, Vec::new());
connection.in_use_cursors.insert(42);
let result = QueryResult {
columns: cols(&["A"]),
rows: vec![text_row(&["first"])],
cursor_id: 42,
more_rows: true,
..QueryResult::default()
};
let deadline = QueryDeadline::new(&cx, Some(QUERY_TIMEOUT));
let mut stream =
OwnedRowStream::from_first_page(Some(connection), cx, 1, deadline, result);
let first = poll_fn(|task_cx| Pin::new(&mut stream).poll_next(task_cx)).await;
assert!(matches!(first, Some(Ok(_))), "first page must be buffered");
thread::sleep(QUERY_TIMEOUT + Duration::from_millis(30));
let second = poll_fn(|task_cx| Pin::new(&mut stream).poll_next(task_cx)).await;
assert!(
matches!(second, Some(Err(Error::CallTimeout(40)))),
"expired logical query must fail immediately, got {second:?}"
);
assert!(
!stream
.connection
.as_ref()
.expect("connection is returned after the failed fetch")
.in_use_cursors
.contains(&42),
"before-start expiry must release the cursor without wire recovery"
);
let connection = stream
.connection
.as_ref()
.expect("connection is returned after the failed fetch");
assert!(
connection
.statement_cache
.iter()
.any(|entry| entry.sql == SQL && entry.cursor_id == 42),
"a fetch proven not to start must leave the valid cached cursor reusable"
);
assert!(
connection.cursors_to_close.is_empty(),
"a fetch proven not to start must not queue the valid cursor for close"
);
Ok::<_, Error>(())
})?;
let saw_fetch = server.join().expect("server thread joins")?;
assert!(
!saw_fetch,
"page 2 was sent after the original query timeout had elapsed; \
the continuation incorrectly received a fresh timeout window"
);
Ok(())
}
#[test]
fn continuation_fetch_error_retires_every_cursor_registry_before_connection_reuse() -> Result<()>
{
const SQL: &str = "select value from drvqa_cursor_cleanup";
const CURSOR_ID: u32 = 42;
const COPIED_CURSOR_ID: u32 = 43;
let listener = std::net::TcpListener::bind("127.0.0.1:0")?;
let addr = listener.local_addr()?;
let server = thread::spawn(move || -> std::io::Result<()> {
let (socket, _) = listener.accept()?;
drop(socket);
Ok(())
});
let runtime = crate::build_io_runtime()?;
runtime.block_on(async {
let cx = Cx::current().expect("runtime installs an ambient Cx");
let socket = TcpStream::connect(addr).await?;
let (read, write) = crate::transport::plain_split(socket);
let mut connection = crate::tests::loopback_connection(read, write);
connection.statement_cache_put(SQL, CURSOR_ID, Vec::new());
connection.in_use_cursors.insert(CURSOR_ID);
connection.cursor_columns.insert(CURSOR_ID, cols(&["A"]));
connection.lob_prefetch_cursors.insert(CURSOR_ID);
let first = QueryResult {
columns: cols(&["A"]),
rows: Vec::new(),
cursor_id: CURSOR_ID,
more_rows: true,
..QueryResult::default()
};
let deadline = QueryDeadline::new(&cx, None);
let mut stream =
OwnedRowStream::from_first_page(Some(connection), cx.clone(), 1, deadline, first);
let connection = stream
.connection
.take()
.expect("stream starts with its connection");
let cancel_handle = connection.cancel_handle()?;
stream.state = OwnedRowStreamState::Fetching {
future: build_fetch_future(
connection,
cx,
CURSOR_ID,
1,
deadline,
Arc::clone(&stream.columns),
None,
),
cancel_handle,
};
let failed = poll_fn(|task_cx| Pin::new(&mut stream).poll_next(task_cx)).await;
assert!(
matches!(failed, Some(Err(_))),
"peer EOF must fail the fetch"
);
assert!(
poll_fn(|task_cx| Pin::new(&mut stream).poll_next(task_cx))
.await
.is_none(),
"a failed stream must surface exactly one error and then terminate"
);
let connection = stream.into_connection().await?;
assert!(
!connection.in_use_cursors.contains(&CURSOR_ID),
"failed cursor must not remain marked in use"
);
assert!(
!connection.cursor_columns.contains_key(&CURSOR_ID),
"failed cursor metadata must be forgotten"
);
assert!(
!connection.lob_prefetch_cursors.contains(&CURSOR_ID),
"failed cursor LOB-prefetch state must be forgotten"
);
assert!(
connection
.statement_cache
.iter()
.all(|entry| entry.cursor_id != CURSOR_ID),
"failed cursor must not remain reusable through the statement cache"
);
assert_eq!(
connection
.cursors_to_close
.iter()
.filter(|cursor_id| **cursor_id == CURSOR_ID)
.count(),
1,
"failed cursor must be queued for close exactly once"
);
let mut connection = connection;
connection.in_use_cursors.insert(COPIED_CURSOR_ID);
connection.copied_cursors.insert(COPIED_CURSOR_ID);
connection
.cursor_columns
.insert(COPIED_CURSOR_ID, cols(&["A"]));
connection.lob_prefetch_cursors.insert(COPIED_CURSOR_ID);
for _ in 0..2 {
let failed: Result<()> = connection.close_cursor_on_error(
COPIED_CURSOR_ID,
Err(Error::Runtime(
"synthetic copied-cursor failure".to_string(),
)),
);
assert!(failed.is_err());
}
assert!(!connection.in_use_cursors.contains(&COPIED_CURSOR_ID));
assert!(!connection.copied_cursors.contains(&COPIED_CURSOR_ID));
assert!(!connection.cursor_columns.contains_key(&COPIED_CURSOR_ID));
assert!(!connection.lob_prefetch_cursors.contains(&COPIED_CURSOR_ID));
assert_eq!(
connection
.cursors_to_close
.iter()
.filter(|cursor_id| **cursor_id == COPIED_CURSOR_ID)
.count(),
1,
"repeated error cleanup must not queue duplicate cursor closes"
);
Ok::<_, Error>(())
})?;
server.join().expect("server thread joins")?;
Ok(())
}
#[test]
fn inflight_continuation_timeout_closes_cursor_after_wire_recovery() -> Result<()> {
const SQL: &str = "select value from drvqa_cursor_timeout";
const CURSOR_ID: u32 = 84;
const QUERY_TIMEOUT: Duration = Duration::from_millis(40);
const INFLIGHT_BODY: &[u8] = &[0xd1, 0xa1, 0xb1];
const CANCEL_ERROR: &[u8] = &[0x04, 0x01, 0x0d];
let listener = std::net::TcpListener::bind("127.0.0.1:0")?;
let addr = listener.local_addr()?;
let server = thread::spawn(move || -> std::io::Result<()> {
let (mut socket, _) = listener.accept()?;
socket.set_read_timeout(Some(Duration::from_secs(2)))?;
let (fetch_type, _) = read_packet(&mut socket)?;
assert_eq!(fetch_type, TNS_PACKET_TYPE_DATA, "page 2 must start");
let (break_type, break_body) = read_packet(&mut socket)?;
assert_eq!(break_type, crate::TNS_PACKET_TYPE_MARKER);
assert_eq!(break_body, [1, 0, crate::TNS_MARKER_TYPE_BREAK]);
socket.write_all(&data_packet(INFLIGHT_BODY))?;
socket.write_all(&marker_packet(crate::TNS_MARKER_TYPE_BREAK))?;
let (reset_type, reset_body) = read_packet(&mut socket)?;
assert_eq!(reset_type, crate::TNS_PACKET_TYPE_MARKER);
assert_eq!(reset_body, [1, 0, crate::TNS_MARKER_TYPE_RESET]);
socket.write_all(&marker_packet(crate::TNS_MARKER_TYPE_RESET))?;
socket.write_all(&data_packet(CANCEL_ERROR))?;
socket.flush()?;
Ok(())
});
let runtime = crate::build_io_runtime()?;
runtime.block_on(async {
let cx = Cx::current().expect("runtime installs an ambient Cx");
let socket = TcpStream::connect(addr).await?;
let (read, write) = crate::transport::plain_split(socket);
let mut connection = crate::tests::loopback_connection(read, write);
connection.statement_cache_put(SQL, CURSOR_ID, Vec::new());
connection.in_use_cursors.insert(CURSOR_ID);
let first = QueryResult {
columns: cols(&["A"]),
rows: Vec::new(),
cursor_id: CURSOR_ID,
more_rows: true,
..QueryResult::default()
};
let deadline = QueryDeadline::new(&cx, Some(QUERY_TIMEOUT));
let mut stream =
OwnedRowStream::from_first_page(Some(connection), cx, 1, deadline, first);
let failed = poll_fn(|task_cx| Pin::new(&mut stream).poll_next(task_cx)).await;
assert!(
matches!(failed, Some(Err(Error::CallTimeout(40)))),
"in-flight page must recover and surface CallTimeout, got {failed:?}"
);
let connection = stream.into_connection().await?;
assert!(!connection.in_use_cursors.contains(&CURSOR_ID));
assert!(
connection
.statement_cache
.iter()
.all(|entry| entry.cursor_id != CURSOR_ID),
"recovered in-flight timeout must evict the uncertain cached cursor"
);
assert_eq!(
connection
.cursors_to_close
.iter()
.filter(|cursor_id| **cursor_id == CURSOR_ID)
.count(),
1,
"recovered in-flight timeout must queue exactly one close"
);
Ok::<_, Error>(())
})?;
server.join().expect("server thread joins")?;
Ok(())
}
#[test]
fn inflight_continuation_into_connection_cancels_and_returns_a_clean_connection() -> Result<()>
{
const CURSOR_ID: u32 = 85;
const INFLIGHT_BODY: &[u8] = &[0xd1, 0xa1, 0xb1];
const CANCEL_ERROR: &[u8] = &[0x04, 0x01, 0x0d];
const FRESH_BODY: &[u8] = &[0x07, 0x05, 0x0c];
let listener = std::net::TcpListener::bind("127.0.0.1:0")?;
let addr = listener.local_addr()?;
let (fetch_seen_tx, fetch_seen_rx) = std::sync::mpsc::channel();
let server = thread::spawn(move || -> std::io::Result<()> {
let (mut socket, _) = listener.accept()?;
socket.set_read_timeout(Some(Duration::from_secs(2)))?;
let (fetch_type, _) = read_packet(&mut socket)?;
assert_eq!(fetch_type, TNS_PACKET_TYPE_DATA, "page 2 must start");
fetch_seen_tx
.send(())
.expect("client waits until the continuation is in flight");
let (break_type, break_body) = read_packet(&mut socket)?;
assert_eq!(break_type, crate::TNS_PACKET_TYPE_MARKER);
assert_eq!(break_body, [1, 0, crate::TNS_MARKER_TYPE_BREAK]);
socket.write_all(&data_packet(INFLIGHT_BODY))?;
socket.write_all(&marker_packet(crate::TNS_MARKER_TYPE_BREAK))?;
let (reset_type, reset_body) = read_packet(&mut socket)?;
assert_eq!(reset_type, crate::TNS_PACKET_TYPE_MARKER);
assert_eq!(reset_body, [1, 0, crate::TNS_MARKER_TYPE_RESET]);
socket.write_all(&marker_packet(crate::TNS_MARKER_TYPE_RESET))?;
socket.write_all(&data_packet(CANCEL_ERROR))?;
socket.flush()?;
let (fresh_type, _) = read_packet(&mut socket)?;
assert_eq!(
fresh_type, TNS_PACKET_TYPE_DATA,
"the reclaimed connection sends its new request only after the drain"
);
socket.write_all(&data_packet(FRESH_BODY))?;
socket.flush()?;
Ok(())
});
let runtime = crate::build_io_runtime()?;
runtime.block_on(async {
let cx = Cx::current().expect("runtime installs an ambient Cx");
let socket = TcpStream::connect(addr).await?;
let (read, write) = crate::transport::plain_split(socket);
let mut connection = crate::tests::loopback_connection(read, write);
connection.in_use_cursors.insert(CURSOR_ID);
let first = QueryResult {
columns: cols(&["A"]),
rows: Vec::new(),
cursor_id: CURSOR_ID,
more_rows: true,
..QueryResult::default()
};
let deadline = QueryDeadline::new(&cx, None);
let mut stream =
OwnedRowStream::from_first_page(Some(connection), cx.clone(), 1, deadline, first);
let mut task_cx = Context::from_waker(Waker::noop());
assert!(
matches!(Pin::new(&mut stream).poll_next(&mut task_cx), Poll::Pending),
"the continuation must be pending before reclamation"
);
fetch_seen_rx
.recv_timeout(Duration::from_secs(2))
.expect("the server observed the in-flight continuation");
let mut connection = stream.into_connection().await?;
assert!(
!connection.in_use_cursors.contains(&CURSOR_ID),
"reclamation releases the abandoned cursor"
);
let sdu = connection.sdu;
connection.core.send_data_packet(&cx, b"fresh", sdu).await?;
let response = connection.core.read_data_response(&cx).await?;
assert_eq!(
response, FRESH_BODY,
"the reclaimed connection reads only its own fresh response"
);
Ok::<_, Error>(())
})?;
server.join().expect("server thread joins")?;
Ok(())
}
#[test]
fn inflight_continuation_reclaims_after_stream_context_is_cancelled() -> Result<()> {
const CURSOR_ID: u32 = 86;
const INFLIGHT_BODY: &[u8] = &[0xd2, 0xa2, 0xb2];
const CANCEL_ERROR: &[u8] = &[0x04, 0x01, 0x0d];
let listener = std::net::TcpListener::bind("127.0.0.1:0")?;
let addr = listener.local_addr()?;
let (fetch_seen_tx, fetch_seen_rx) = std::sync::mpsc::channel();
let server = thread::spawn(move || -> std::io::Result<()> {
let (mut socket, _) = listener.accept()?;
socket.set_read_timeout(Some(Duration::from_secs(2)))?;
let (fetch_type, _) = read_packet(&mut socket)?;
assert_eq!(fetch_type, TNS_PACKET_TYPE_DATA, "page 2 must start");
fetch_seen_tx.send(()).expect("client waits for the fetch");
let (break_type, break_body) = read_packet(&mut socket)?;
assert_eq!(break_type, crate::TNS_PACKET_TYPE_MARKER);
assert_eq!(break_body, [1, 0, crate::TNS_MARKER_TYPE_BREAK]);
socket.write_all(&data_packet(INFLIGHT_BODY))?;
socket.write_all(&marker_packet(crate::TNS_MARKER_TYPE_BREAK))?;
let (reset_type, reset_body) = read_packet(&mut socket)?;
assert_eq!(reset_type, crate::TNS_PACKET_TYPE_MARKER);
assert_eq!(reset_body, [1, 0, crate::TNS_MARKER_TYPE_RESET]);
socket.write_all(&marker_packet(crate::TNS_MARKER_TYPE_RESET))?;
socket.write_all(&data_packet(CANCEL_ERROR))?;
socket.flush()?;
Ok(())
});
let runtime = crate::build_io_runtime()?;
runtime.block_on(async {
let cx = Cx::current().expect("runtime installs an ambient Cx");
let socket = TcpStream::connect(addr).await?;
let (read, write) = crate::transport::plain_split(socket);
let mut connection = crate::tests::loopback_connection(read, write);
connection.in_use_cursors.insert(CURSOR_ID);
let first = QueryResult {
columns: cols(&["A"]),
rows: Vec::new(),
cursor_id: CURSOR_ID,
more_rows: true,
..QueryResult::default()
};
let deadline = QueryDeadline::new(&cx, None);
let mut stream =
OwnedRowStream::from_first_page(Some(connection), cx.clone(), 1, deadline, first);
let mut task_cx = Context::from_waker(Waker::noop());
assert!(matches!(
Pin::new(&mut stream).poll_next(&mut task_cx),
Poll::Pending
));
fetch_seen_rx
.recv_timeout(Duration::from_secs(2))
.expect("the server observed the in-flight continuation");
cx.cancel_fast(CancelKind::User);
let connection = stream.into_connection().await?;
assert!(!connection.in_use_cursors.contains(&CURSOR_ID));
assert_eq!(
connection.core.recovery.phase(),
crate::SessionRecoveryPhase::Ready,
"reclamation must leave the wire clean even when the stream Cx is cancelled"
);
Ok::<_, Error>(())
})?;
server.join().expect("server thread joins")?;
Ok(())
}
#[test]
fn buffered_rows_yield_in_order_then_terminate() {
with_cx(|cx| {
let result = QueryResult {
columns: cols(&["A", "B"]),
rows: vec![text_row(&["1", "x"]), text_row(&["2", "y"])],
cursor_id: 7,
more_rows: false,
..QueryResult::default()
};
let mut stream = offline_stream(cx, result);
assert_eq!(stream.columns().len(), 2);
let drained = drain_buffered(&mut stream);
assert_eq!(drained.len(), 2, "both buffered rows yield");
assert_eq!(drained[0].as_ref().unwrap(), &text_row(&["1", "x"]));
assert_eq!(drained[1].as_ref().unwrap(), &text_row(&["2", "y"]));
assert_eq!(stream.cursor_id, 0, "cursor released on full drain");
assert!(matches!(stream.state, OwnedRowStreamState::Done));
});
}
#[test]
fn empty_single_page_yields_no_rows() {
with_cx(|cx| {
let result = QueryResult {
columns: cols(&["A"]),
rows: Vec::new(),
cursor_id: 0,
more_rows: false,
..QueryResult::default()
};
let mut stream = offline_stream(cx, result);
assert!(drain_buffered(&mut stream).is_empty());
assert!(matches!(stream.state, OwnedRowStreamState::Done));
});
}
#[test]
fn first_page_seeds_duplicate_column_continuation() {
with_cx(|cx| {
let result = QueryResult {
columns: cols(&["A", "B"]),
rows: vec![text_row(&["1", "x"]), text_row(&["2", "y"])],
cursor_id: 7,
more_rows: true,
..QueryResult::default()
};
let stream = offline_stream(cx, result);
assert_eq!(stream.previous_row, Some(text_row(&["2", "y"])));
});
}
#[test]
fn empty_continuation_page_keeps_previous_seed() {
with_cx(|cx| {
let first = QueryResult {
columns: cols(&["A"]),
rows: vec![text_row(&["seed"])],
cursor_id: 7,
more_rows: true,
..QueryResult::default()
};
let mut stream = offline_stream(cx, first);
assert_eq!(stream.previous_row, Some(text_row(&["seed"])));
stream.apply_page(QueryResult {
columns: Vec::new(),
rows: Vec::new(),
cursor_id: 0,
more_rows: true,
..QueryResult::default()
});
assert_eq!(
stream.previous_row,
Some(text_row(&["seed"])),
"empty page keeps the previous seed"
);
assert!(matches!(stream.state, OwnedRowStreamState::Buffered(ref b) if b.is_empty()));
stream.apply_page(QueryResult {
columns: Vec::new(),
rows: vec![text_row(&["p2a"]), text_row(&["p2b"])],
cursor_id: 0,
more_rows: false,
..QueryResult::default()
});
assert_eq!(stream.previous_row, Some(text_row(&["p2b"])));
});
}
#[test]
fn apply_page_adopts_redescribed_cursor_and_columns() {
with_cx(|cx| {
let first = QueryResult {
columns: cols(&["A"]),
rows: vec![text_row(&["1"])],
cursor_id: 7,
more_rows: true,
..QueryResult::default()
};
let mut stream = offline_stream(cx, first);
assert_eq!(stream.cursor_id, 7);
stream.apply_page(QueryResult {
columns: cols(&["A", "B"]),
rows: vec![text_row(&["2", "z"])],
cursor_id: 9,
more_rows: false,
..QueryResult::default()
});
assert_eq!(stream.cursor_id, 9, "adopts re-described cursor id");
assert_eq!(stream.columns().len(), 2, "adopts re-described columns");
assert!(!stream.more_rows);
});
}
#[test]
fn into_connection_on_poisoned_stream_errors() {
let runtime = crate::build_io_runtime().expect("io runtime builds");
runtime.block_on(async {
let cx = Cx::current().expect("block_on installs an ambient Cx");
let result = QueryResult {
columns: cols(&["A"]),
rows: vec![text_row(&["1"])],
cursor_id: 7,
more_rows: false,
..QueryResult::default()
};
let stream = offline_stream(cx, result);
let err = stream.into_connection().await.unwrap_err();
assert!(matches!(err, Error::ConnectionClosed(_)));
});
}
#[test]
fn fetch_without_connection_poisons_cleanly() {
with_cx(|cx| {
let result = QueryResult {
columns: cols(&["A"]),
rows: vec![text_row(&["only"])],
cursor_id: 7,
more_rows: true,
..QueryResult::default()
};
let mut stream = offline_stream(cx, result);
let mut stream = pin!(&mut stream);
let mut task_cx = Context::from_waker(Waker::noop());
let first = stream.as_mut().poll_next(&mut task_cx);
assert!(matches!(first, Poll::Ready(Some(Ok(_)))));
let second = stream.as_mut().poll_next(&mut task_cx);
match second {
Poll::Ready(Some(Err(Error::ConnectionClosed(_)))) => {}
other => panic!("expected clean ConnectionClosed, got {other:?}"),
}
assert!(matches!(
stream.as_mut().poll_next(&mut task_cx),
Poll::Ready(None)
));
});
}
}