Skip to main content

safer_ring/future/
operation_future.rs

1//! Generic future for any operation type.
2
3use std::future::Future;
4use std::io;
5use std::marker::PhantomData;
6use std::pin::Pin as StdPin;
7use std::sync::Arc;
8use std::task::{Context, Poll};
9
10use crate::future::waker::WakerRegistry;
11use crate::operation::BufferType;
12use crate::operation::{Operation, Submitted};
13use crate::ring::Ring;
14
15/// Generic future for any operation type.
16///
17/// This provides a unified interface for different operation types while
18/// maintaining type safety and proper buffer ownership semantics.
19/// Unlike the specialized I/O futures, this returns the raw i32 result
20/// from io_uring without conversion to usize.
21///
22/// # Type Parameters
23///
24/// * `'ring` - Lifetime of the io_uring instance
25/// * `'buf` - Lifetime of the buffer being used for the operation
26///
27/// # Returns
28///
29/// Returns `(i32, Option<Pin<&'buf mut [u8]>>)` where:
30/// - `i32` is the raw result from io_uring (can be negative for errors)
31/// - `Option<Pin<&'buf mut [u8]>>` is the buffer if the operation used one
32pub struct OperationFuture<'ring, 'buf> {
33    /// The underlying operation (None after completion)
34    /// Using Option to allow taking ownership during completion
35    operation: Option<Operation<'ring, 'buf, Submitted>>,
36    /// Reference to the ring for polling completions
37    ring: &'ring mut Ring<'ring>,
38    /// Waker registry for async notification
39    waker_registry: Arc<WakerRegistry>,
40    /// Phantom data for lifetime tracking
41    _phantom: PhantomData<(&'ring (), &'buf ())>,
42}
43
44impl<'ring, 'buf> OperationFuture<'ring, 'buf> {
45    /// Create a new operation future from a submitted operation.
46    pub(crate) fn new(
47        operation: Operation<'ring, 'buf, Submitted>,
48        ring: &'ring mut Ring<'ring>,
49        waker_registry: Arc<WakerRegistry>,
50    ) -> Self {
51        Self {
52            operation: Some(operation),
53            ring,
54            waker_registry,
55            _phantom: PhantomData,
56        }
57    }
58}
59
60impl<'ring, 'buf> Future for OperationFuture<'ring, 'buf> {
61    type Output = io::Result<(i32, Option<StdPin<&'buf mut [u8]>>)>;
62
63    fn poll(mut self: StdPin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
64        let operation = match self.operation.as_ref() {
65            Some(op) => op,
66            None => {
67                // Future has already been polled to completion
68                // This is a programming error - futures should not be polled after Ready
69                panic!("OperationFuture polled after completion");
70            }
71        };
72
73        let operation_id = operation.id();
74
75        // Check if the operation has completed
76        match self.ring.try_complete_by_id(operation_id) {
77            Ok(Some(result)) => {
78                // Operation completed, extract the result and buffer
79                let operation = self.operation.take().unwrap();
80                let completed = operation.complete_with_result(result);
81                let (io_result, buffer) = completed.into_result();
82
83                // Clean up waker registration to prevent memory leaks
84                self.waker_registry.remove_waker(operation_id);
85
86                // Return the raw result and buffer without conversion
87                // This allows callers to handle negative results as needed
88                // Convert BufferType to Option<Pin<&mut [u8]>> for compatibility
89                let buffer_option = match buffer {
90                    BufferType::Pinned(buf) => Some(buf),
91                    _ => None,
92                };
93                Poll::Ready(io_result.map(|bytes| (bytes, buffer_option)))
94            }
95            Ok(None) => {
96                // Operation still in flight, register waker and return Pending
97                // Clone the waker to avoid borrowing issues
98                self.waker_registry
99                    .register_waker(operation_id, cx.waker().clone());
100                Poll::Pending
101            }
102            Err(e) => {
103                // Error checking completion status
104                self.waker_registry.remove_waker(operation_id);
105                Poll::Ready(Err(io::Error::other(format!(
106                    "Error checking operation completion: {e}"
107                ))))
108            }
109        }
110    }
111}
112
113impl<'ring, 'buf> Drop for OperationFuture<'ring, 'buf> {
114    fn drop(&mut self) {
115        // Clean up waker registration if the future is dropped before completion
116        // This prevents memory leaks in the waker registry
117        if let Some(operation) = &self.operation {
118            self.waker_registry.remove_waker(operation.id());
119        }
120    }
121}