@@ -84,29 +84,38 @@ pub fn checkout_streaming_chunk_buffer(
8484 let size = streaming_chunk_size ( ) ;
8585 role. with_cache ( |cache| {
8686 let mut slot = cache. borrow_mut ( ) ;
87- let replace_cached = slot
88- . as_ref ( )
89- . is_none_or ( |cached| cached. size != size && !cached. checked_out ) ;
90-
91- if replace_cached {
92- * slot = Some ( CachedStreamingChunkBuffer {
93- size,
94- array : new_streaming_chunk_buffer ( env, size) ?,
95- checked_out : false ,
96- } ) ;
97- }
98-
99- let Some ( cached) = slot. as_mut ( ) else {
100- return Ok ( ( new_streaming_chunk_buffer ( env, size) ?, None ) ) ;
101- } ;
102-
103- if cached. size != size || cached. checked_out {
104- return Ok ( ( new_streaming_chunk_buffer ( env, size) ?, None ) ) ;
87+ // Three outcomes, decided by the cached slot's state:
88+ match slot. as_mut ( ) {
89+ // Still checked out — a concurrent dispatch holds it, or a prior
90+ // dispatch panicked mid-stream and never returned its lease. Hand
91+ // back a throwaway, unpooled buffer and leave the cache untouched
92+ // so we never alias a Java array that may still be in flight.
93+ Some ( cached) if cached. checked_out => {
94+ return Ok ( ( new_streaming_chunk_buffer ( env, size) ?, None ) ) ;
95+ }
96+ // Free to reuse — refresh the backing array only if the configured
97+ // chunk size changed, then lease it back to the caller.
98+ Some ( cached) => {
99+ if cached. size != size {
100+ cached. array = new_streaming_chunk_buffer ( env, size) ?;
101+ cached. size = size;
102+ }
103+ let cached_array: & JByteArray < ' static > = cached. array . as_ref ( ) ;
104+ let dispatch_array = env. new_global_ref ( cached_array) ?;
105+ cached. checked_out = true ;
106+ return Ok ( ( dispatch_array, Some ( StreamingChunkBufferLease :: new ( role) ) ) ) ;
107+ }
108+ // Empty slot — fall through to install a fresh cached buffer.
109+ None => { }
105110 }
106-
107- let cached_array: & JByteArray < ' static > = cached. array . as_ref ( ) ;
108- let dispatch_array = env. new_global_ref ( cached_array) ?;
109- cached. checked_out = true ;
111+ let array = new_streaming_chunk_buffer ( env, size) ?;
112+ let array_ref: & JByteArray < ' static > = array. as_ref ( ) ;
113+ let dispatch_array = env. new_global_ref ( array_ref) ?;
114+ * slot = Some ( CachedStreamingChunkBuffer {
115+ size,
116+ array,
117+ checked_out : true ,
118+ } ) ;
110119 Ok ( ( dispatch_array, Some ( StreamingChunkBufferLease :: new ( role) ) ) )
111120 } )
112121}
@@ -116,3 +125,39 @@ pub fn mark_streaming_buffer_reusable(lease: Option<StreamingChunkBufferLease>)
116125 lease. mark_reusable ( ) ;
117126 }
118127}
128+
129+ /// The pull + push per-thread chunk buffers (and their leases) acquired
130+ /// together for one bidirectional streaming dispatch.
131+ pub struct PullPushBuffers {
132+ pub pull_buf : StreamingChunkBuffer ,
133+ pub pull_buf_lease : Option < StreamingChunkBufferLease > ,
134+ pub push_buf : StreamingChunkBuffer ,
135+ pub push_buf_lease : Option < StreamingChunkBufferLease > ,
136+ }
137+
138+ /// Check out the pull + push chunk buffers for a bidirectional stream in
139+ /// one step. Pull and push run concurrently on different threads, so each
140+ /// direction gets its own per-thread cached buffer.
141+ ///
142+ /// If the push checkout fails after the pull buffer was already leased, the
143+ /// pull lease is released before returning the error so a half-acquired pair
144+ /// never leaks a leased buffer (which would force the next dispatch to
145+ /// allocate a fresh array). Centralising this cleanup keeps the invariant in
146+ /// one place instead of duplicating it across every bidirectional entry point.
147+ pub fn checkout_pull_push_buffers ( env : & mut jni:: Env < ' _ > ) -> jni:: errors:: Result < PullPushBuffers > {
148+ let ( pull_buf, pull_buf_lease) = checkout_streaming_chunk_buffer ( env, StreamingBufferRole :: Pull ) ?;
149+ let ( push_buf, push_buf_lease) =
150+ match checkout_streaming_chunk_buffer ( env, StreamingBufferRole :: Push ) {
151+ Ok ( checked_out) => checked_out,
152+ Err ( err) => {
153+ mark_streaming_buffer_reusable ( pull_buf_lease) ;
154+ return Err ( err) ;
155+ }
156+ } ;
157+ Ok ( PullPushBuffers {
158+ pull_buf,
159+ pull_buf_lease,
160+ push_buf,
161+ push_buf_lease,
162+ } )
163+ }
0 commit comments