11use std:: ffi:: { CString , c_char} ;
2+ use std:: fmt:: Write as _;
23use std:: net:: SocketAddr ;
34use std:: sync:: Arc ;
45use std:: time:: Duration ;
@@ -16,7 +17,7 @@ use crate::cluster::CassCluster;
1617use crate :: future:: { CassFuture , CassResultValue } ;
1718use crate :: retry_policy:: CassRetryPolicy ;
1819use crate :: statement:: { BoundStatement , CassStatement } ;
19- use crate :: types:: { cass_int32_t, cass_uint16_t, cass_uint64_t, size_t} ;
20+ use crate :: types:: { cass_bool_t , cass_int32_t, cass_uint16_t, cass_uint64_t, size_t} ;
2021
2122#[ unsafe( no_mangle) ]
2223pub unsafe extern "C" fn testing_cluster_get_connect_timeout (
@@ -177,6 +178,156 @@ pub unsafe extern "C" fn testing_batch_set_sleeping_history_listener(
177178 . set_history_listener ( history_listener)
178179}
179180
181+ /// Used to record attempted hosts during a request.
182+ /// This is useful for testing purposes and is used in integration tests.
183+ /// This is enabled by `testing_statement_set_recording_history_listener`
184+ /// and can be queried using `testing_future_get_attempted_hosts`.
185+ #[ derive( Debug ) ]
186+ pub ( crate ) struct RecordingHistoryListener {
187+ attempted_hosts : std:: sync:: Mutex < Vec < SocketAddr > > ,
188+ }
189+
190+ impl RecordingHistoryListener {
191+ pub ( crate ) fn new ( ) -> Self {
192+ RecordingHistoryListener {
193+ attempted_hosts : std:: sync:: Mutex :: new ( Vec :: new ( ) ) ,
194+ }
195+ }
196+
197+ pub ( crate ) fn get_attempted_hosts ( & self ) -> Vec < SocketAddr > {
198+ self . attempted_hosts . lock ( ) . unwrap ( ) . clone ( )
199+ }
200+ }
201+
202+ impl HistoryListener for RecordingHistoryListener {
203+ fn log_request_start ( & self ) -> RequestId {
204+ RequestId ( 0 )
205+ }
206+
207+ fn log_request_success ( & self , _request_id : RequestId ) { }
208+
209+ fn log_request_error ( & self , _request_id : RequestId , _error : & RequestError ) { }
210+
211+ fn log_new_speculative_fiber ( & self , _request_id : RequestId ) -> SpeculativeId {
212+ SpeculativeId ( 0 )
213+ }
214+
215+ fn log_attempt_start (
216+ & self ,
217+ _request_id : RequestId ,
218+ _speculative_id : Option < SpeculativeId > ,
219+ node_addr : SocketAddr ,
220+ ) -> AttemptId {
221+ // Record the host that was attempted.
222+ self . attempted_hosts . lock ( ) . unwrap ( ) . push ( node_addr) ;
223+
224+ AttemptId ( 0 )
225+ }
226+
227+ fn log_attempt_success ( & self , _attempt_id : AttemptId ) { }
228+
229+ fn log_attempt_error (
230+ & self ,
231+ _attempt_id : AttemptId ,
232+ _error : & RequestAttemptError ,
233+ _retry_decision : & RetryDecision ,
234+ ) {
235+ }
236+ }
237+
238+ /// Enables recording of attempted hosts in the statement's history listener.
239+ /// This is useful for testing purposes, allowing us to verify which hosts were attempted
240+ /// during a request.
241+ /// Later, `testing_future_get_attempted_hosts` can be used to retrieve this information.
242+ #[ unsafe( no_mangle) ]
243+ pub unsafe extern "C" fn testing_statement_set_recording_history_listener (
244+ statement_raw : CassBorrowedExclusivePtr < CassStatement , CMut > ,
245+ enable : cass_bool_t ,
246+ ) {
247+ let statement = & mut BoxFFI :: as_mut_ref ( statement_raw) . unwrap ( ) ;
248+
249+ statement. record_hosts = enable != 0 ;
250+ }
251+
252+ /// Retrieves hosts that were attempted during the execution of a future.
253+ /// In order for this to work, the statement must have been configured with
254+ /// `testing_statement_set_recording_history_listener`.
255+ #[ unsafe( no_mangle) ]
256+ pub unsafe extern "C" fn testing_future_get_attempted_hosts (
257+ future_raw : CassBorrowedSharedPtr < CassFuture , CMut > ,
258+ ) -> * mut c_char {
259+ // This function should return a list of attempted hosts.
260+ // Care should be taken to ensure that the list is properly allocated and freed.
261+ // The problem is that the return type must be understandable by C code.
262+ // See testing.cpp:53 (get_attempted_hosts_from_future()) for an example of how
263+ // this is done in C++.
264+ //
265+ // Idea: Create a concatenated string of attempted hosts.
266+ // Return a pointer to that string. Caller is responsible for freeing it.
267+
268+ let future: & CassFuture = ArcFFI :: as_ref ( future_raw) . unwrap ( ) ;
269+
270+ let attempted_hosts = future. attempted_hosts ( ) ;
271+ let concatenated_hosts = attempted_hosts. iter ( ) . fold ( String :: new ( ) , |mut acc, host| {
272+ // Convert the SocketAddr to a string.
273+ // Delimit address strings with '\n' to enable easy parsing in C.
274+
275+ write ! ( & mut acc, "{}\n " , host. ip( ) ) . unwrap ( ) ;
276+ acc
277+ } ) ;
278+
279+ // The caller is responsible for freeing this memory, by calling `testing_free_cstring`.
280+ unsafe { CString :: from_vec_unchecked ( concatenated_hosts. into_bytes ( ) ) } . into_raw ( )
281+ }
282+
283+ /// Ensures that the `testing_future_get_attempted_hosts` function
284+ /// behaves correctly, i.e., it returns a list of attempted hosts as a concatenated string.
285+ #[ test]
286+ fn test_future_get_attempted_hosts ( ) {
287+ use scylla:: observability:: history:: HistoryListener as _;
288+
289+ let listener = Arc :: new ( RecordingHistoryListener :: new ( ) ) ;
290+ let future = CassFuture :: new_from_future ( std:: future:: pending ( ) , Some ( listener. clone ( ) ) ) ;
291+
292+ fn assert_attempted_hosts_eq ( future : & Arc < CassFuture > , hosts : & [ String ] ) {
293+ let hosts_str = unsafe { testing_future_get_attempted_hosts ( ArcFFI :: as_ptr ( future) ) } ;
294+ let hosts_cstr = unsafe { CString :: from_raw ( hosts_str) } ;
295+ let hosts_string = hosts_cstr. to_str ( ) . unwrap ( ) ;
296+
297+ // Split the string by '\n' and collect into a Vec<&str>.
298+ let attempted_hosts: Vec < & str > =
299+ hosts_string. split ( '\n' ) . filter ( |s| !s. is_empty ( ) ) . collect ( ) ;
300+
301+ assert_eq ! ( attempted_hosts, hosts) ;
302+ }
303+
304+ // 1. Test with no attempted hosts.
305+ {
306+ assert_attempted_hosts_eq ( & future, & [ ] ) ;
307+ }
308+
309+ let addr1: SocketAddr = SocketAddr :: from ( ( [ 127 , 0 , 0 , 1 ] , 9042 ) ) ;
310+ let addr2: SocketAddr = SocketAddr :: from ( ( [ 127 , 0 , 0 , 2 ] , 9042 ) ) ;
311+
312+ // 2. Attempt two hosts and see if they are recorded correctly, in order.
313+ {
314+ listener. log_attempt_start ( RequestId ( 0 ) , None , addr1) ;
315+ listener. log_attempt_start ( RequestId ( 0 ) , None , addr2) ;
316+ assert_attempted_hosts_eq ( & future, & [ addr1, addr2] . map ( |addr| addr. ip ( ) . to_string ( ) ) )
317+ }
318+
319+ let addr3: SocketAddr = SocketAddr :: from ( ( [ 127 , 0 , 0 , 3 ] , 9042 ) ) ;
320+
321+ // 3. Attempt one more host and see if all hosts are present, in order.
322+ {
323+ listener. log_attempt_start ( RequestId ( 0 ) , None , addr3) ;
324+ assert_attempted_hosts_eq (
325+ & future,
326+ & [ addr1, addr2, addr3] . map ( |addr| addr. ip ( ) . to_string ( ) ) ,
327+ )
328+ }
329+ }
330+
180331/// A retry policy that always ignores all errors.
181332///
182333/// Useful for testing purposes.
0 commit comments