@@ -17,7 +17,7 @@ use std::{
1717} ;
1818
1919use arc_swap:: ArcSwap ;
20- use tokio:: { fs, sync:: mpsc, time} ;
20+ use tokio:: { fs, sync:: mpsc, sync :: oneshot , time} ;
2121
2222use crate :: {
2323 accumulator,
@@ -445,6 +445,27 @@ impl CaptureManager<formats::jsonl::Format<BufWriter<std::fs::File>>, RealClock>
445445 }
446446}
447447
448+ /// Request to rotate to a new output file
449+ ///
450+ /// Contains the path for the new file and a channel to send the result.
451+ pub struct RotationRequest {
452+ /// Path for the new output file
453+ pub path : PathBuf ,
454+ /// Channel to send rotation result (Ok on success, Err on failure)
455+ pub response : oneshot:: Sender < Result < ( ) , formats:: Error > > ,
456+ }
457+
458+ impl std:: fmt:: Debug for RotationRequest {
459+ fn fmt ( & self , f : & mut std:: fmt:: Formatter < ' _ > ) -> std:: fmt:: Result {
460+ f. debug_struct ( "RotationRequest" )
461+ . field ( "path" , & self . path )
462+ . finish_non_exhaustive ( )
463+ }
464+ }
465+
466+ /// Handle for sending rotation requests to a running CaptureManager
467+ pub type RotationSender = mpsc:: Sender < RotationRequest > ;
468+
448469impl CaptureManager < formats:: parquet:: Format < BufWriter < std:: fs:: File > > , RealClock > {
449470 /// Create a new [`CaptureManager`] with file-based Parquet writer
450471 ///
@@ -478,6 +499,166 @@ impl CaptureManager<formats::parquet::Format<BufWriter<std::fs::File>>, RealCloc
478499 RealClock :: default ( ) ,
479500 ) )
480501 }
502+
503+ /// Run [`CaptureManager`] with file rotation support
504+ ///
505+ /// Similar to [`start`](CaptureManager::start), but also provides a channel
506+ /// for rotation requests. When a rotation request is received, the current
507+ /// Parquet file is finalized (footer written) and a new file is created at
508+ /// the specified path.
509+ ///
510+ /// Returns a tuple of ([`RotationSender`], [`JoinHandle`](tokio::task::JoinHandle))
511+ /// immediately. The `RotationSender` can be used to trigger rotations while
512+ /// the event loop runs. The `JoinHandle` can be awaited to ensure the
513+ /// CaptureManager has fully drained and closed before shutdown.
514+ ///
515+ /// # Errors
516+ ///
517+ /// Returns an error if there is already a global recorder set.
518+ #[ allow( clippy:: cast_possible_truncation) ]
519+ pub async fn start_with_rotation (
520+ mut self ,
521+ ) -> Result < ( RotationSender , tokio:: task:: JoinHandle < ( ) > ) , Error > {
522+ // Create rotation channel - return the sender immediately
523+ let ( rotation_tx, rotation_rx) = mpsc:: channel :: < RotationRequest > ( 4 ) ;
524+
525+ // Initialize historical sender
526+ HISTORICAL_SENDER . store ( Arc :: new ( Some ( Arc :: new ( Sender {
527+ snd : self . snd . clone ( ) ,
528+ } ) ) ) ) ;
529+
530+ self . install ( ) ?;
531+ info ! ( "Capture manager installed with rotation support, recording to capture file." ) ;
532+
533+ // Wait until the target is running then mark time-zero
534+ self . target_running . recv ( ) . await ;
535+ self . clock . mark_start ( ) ;
536+
537+ let compression_level = self . format . compression_level ( ) ;
538+
539+ // Run the event loop in a spawned task so we can return the sender immediately
540+ let expiration = self . expiration ;
541+ let format = self . format ;
542+ let flush_seconds = self . flush_seconds ;
543+ let registry = self . registry ;
544+ let accumulator = self . accumulator ;
545+ let global_labels = self . global_labels ;
546+ let clock = self . clock ;
547+ let recv = self . recv ;
548+ let shutdown = self . shutdown . take ( ) . expect ( "shutdown watcher must be present" ) ;
549+
550+ let handle = tokio:: spawn ( async move {
551+ if let Err ( e) = Self :: rotation_event_loop (
552+ expiration,
553+ format,
554+ flush_seconds,
555+ registry,
556+ accumulator,
557+ global_labels,
558+ clock,
559+ recv,
560+ shutdown,
561+ rotation_rx,
562+ compression_level,
563+ )
564+ . await
565+ {
566+ error ! ( error = %e, "CaptureManager rotation event loop error" ) ;
567+ }
568+ } ) ;
569+
570+ Ok ( ( rotation_tx, handle) )
571+ }
572+
573+ /// Internal event loop with rotation support
574+ #[ allow( clippy:: too_many_arguments) ]
575+ async fn rotation_event_loop (
576+ expiration : Duration ,
577+ format : formats:: parquet:: Format < BufWriter < std:: fs:: File > > ,
578+ flush_seconds : u64 ,
579+ registry : Arc < Registry < Key , AtomicStorage > > ,
580+ accumulator : Accumulator ,
581+ global_labels : FxHashMap < String , String > ,
582+ clock : RealClock ,
583+ mut recv : mpsc:: Receiver < Metric > ,
584+ shutdown : lading_signal:: Watcher ,
585+ mut rotation_rx : mpsc:: Receiver < RotationRequest > ,
586+ compression_level : i32 ,
587+ ) -> Result < ( ) , Error > {
588+ let mut flush_interval = clock. interval ( Duration :: from_millis ( TICK_DURATION_MS as u64 ) ) ;
589+ let shutdown_wait = shutdown. recv ( ) ;
590+ tokio:: pin!( shutdown_wait) ;
591+
592+ // Create state machine with owned state
593+ let mut state_machine = StateMachine :: new (
594+ expiration,
595+ format,
596+ flush_seconds,
597+ registry,
598+ accumulator,
599+ global_labels,
600+ clock,
601+ ) ;
602+
603+ // Event loop with rotation support
604+ loop {
605+ let event = tokio:: select! {
606+ val = recv. recv( ) => {
607+ match val {
608+ Some ( metric) => Event :: MetricReceived ( metric) ,
609+ None => Event :: ChannelClosed ,
610+ }
611+ }
612+ ( ) = flush_interval. tick( ) => Event :: FlushTick ,
613+ Some ( rotation_req) = rotation_rx. recv( ) => {
614+ // Handle rotation inline since it's not a state machine event
615+ let result = Self :: handle_rotation(
616+ & mut state_machine,
617+ rotation_req. path,
618+ compression_level,
619+ ) . await ;
620+ // Send result back to caller (ignore send error if receiver dropped)
621+ let _ = rotation_req. response. send( result) ;
622+ continue ;
623+ }
624+ ( ) = & mut shutdown_wait => Event :: ShutdownSignaled ,
625+ } ;
626+
627+ match state_machine. next ( event) ? {
628+ Operation :: Continue => { }
629+ Operation :: Exit => return Ok ( ( ) ) ,
630+ }
631+ }
632+ }
633+
634+ /// Handle a rotation request
635+ async fn handle_rotation (
636+ state_machine : & mut StateMachine <
637+ formats:: parquet:: Format < BufWriter < std:: fs:: File > > ,
638+ RealClock ,
639+ > ,
640+ new_path : PathBuf ,
641+ compression_level : i32 ,
642+ ) -> Result < ( ) , formats:: Error > {
643+ // Create new file and format
644+ let fp = fs:: File :: create ( & new_path)
645+ . await
646+ . map_err ( formats:: Error :: Io ) ?;
647+ let fp = fp. into_std ( ) . await ;
648+ let writer = BufWriter :: new ( fp) ;
649+ let new_format = parquet:: Format :: new ( writer, compression_level) ?;
650+
651+ // Swap formats - this flushes any buffered data
652+ let old_format = state_machine
653+ . replace_format ( new_format)
654+ . map_err ( |e| formats:: Error :: Io ( io:: Error :: new ( io:: ErrorKind :: Other , e. to_string ( ) ) ) ) ?;
655+
656+ // Close old format to write Parquet footer
657+ old_format. close ( ) ?;
658+
659+ info ! ( path = %new_path. display( ) , "Rotated to new capture file" ) ;
660+ Ok ( ( ) )
661+ }
481662}
482663
483664impl
0 commit comments