Fix refactor recovery and metrics

src/runtime.rs

@@ -62,24 +62,11 @@ enum Operation {
     Measure,
 }
 
-#[derive(Clone, Copy)]
-enum TrackingPolicy {
-    None,
-    Queued,
-    QueuedAndRejected,
-}
-
-impl TrackingPolicy {
-    fn records_queue(self) -> bool {
-        matches!(self, Self::Queued | Self::QueuedAndRejected)
-    }
-
-    fn records_rejection(self) -> bool {
-        matches!(self, Self::QueuedAndRejected)
-    }
-}
-
 impl Operation {
+    fn tracks_metrics(&self) -> bool {
+        !matches!(self, Self::Measure)
+    }
+
     fn error_handler(&self) -> fn(String) -> GenerationEvent {
         match self {
             Self::Generate => |error| GenerationEvent::Finished(Err(error)),
@@ -137,39 +124,33 @@ impl GenerationService {
         source: WorkSource,
         idle_timeout: Duration,
     ) -> Result<ActiveGeneration, String> {
-        self.submit(
-            CommandRequest {
-                engine,
-                turn,
-                messages,
-                checkpoint: Some(checkpoint),
-                source,
-                operation: Operation::Generate,
-                idle_timeout,
-            },
-            TrackingPolicy::QueuedAndRejected,
-        )
+        self.submit(CommandRequest {
+            engine,
+            turn,
+            messages,
+            checkpoint: Some(checkpoint),
+            source,
+            operation: Operation::Generate,
+            idle_timeout,
+        })
     }
 
     pub(crate) fn compact(&self, input: CompactionInput) -> Result<ActiveGeneration, String> {
-        self.submit(
-            CommandRequest {
-                engine: input.engine,
-                turn: input.turn,
-                messages: input.messages,
-                checkpoint: Some(CheckpointTarget::Local {
-                    checkpoint: input.checkpoint,
-                    bootstrap: None,
-                }),
-                source: WorkSource::LocalChat,
-                operation: Operation::Compact {
-                    reason: input.reason,
-                    rebuild_system_prompt: input.rebuild_system_prompt,
-                },
-                idle_timeout: input.idle_timeout,
+        self.submit(CommandRequest {
+            engine: input.engine,
+            turn: input.turn,
+            messages: input.messages,
+            checkpoint: Some(CheckpointTarget::Local {
+                checkpoint: input.checkpoint,
+                bootstrap: None,
+            }),
+            source: WorkSource::LocalChat,
+            operation: Operation::Compact {
+                reason: input.reason,
+                rebuild_system_prompt: input.rebuild_system_prompt,
             },
-            TrackingPolicy::None,
-        )
+            idle_timeout: input.idle_timeout,
+        })
     }
 
     pub(crate) fn measure_context(
@@ -179,28 +160,22 @@ impl GenerationService {
         messages: Vec<ChatTurn>,
         idle_timeout: Duration,
     ) -> Result<ActiveGeneration, String> {
-        self.submit(
-            CommandRequest {
-                engine,
-                turn,
-                messages,
-                checkpoint: None,
-                source: WorkSource::LocalChat,
-                operation: Operation::Measure,
-                idle_timeout,
-            },
-            TrackingPolicy::Queued,
-        )
+        self.submit(CommandRequest {
+            engine,
+            turn,
+            messages,
+            checkpoint: None,
+            source: WorkSource::LocalChat,
+            operation: Operation::Measure,
+            idle_timeout,
+        })
     }
 
-    fn submit(
-        &self,
-        request: CommandRequest,
-        tracking: TrackingPolicy,
-    ) -> Result<ActiveGeneration, String> {
+    fn submit(&self, request: CommandRequest) -> Result<ActiveGeneration, String> {
         let cancel = Arc::new(AtomicBool::new(false));
         let (events, receiver) = mpsc::channel();
-        if tracking.records_queue() {
+        let tracked = request.operation.tracks_metrics();
+        if tracked {
             self.metrics.request_queued(request.source);
         }
         let command = Command {
@@ -215,7 +190,7 @@ impl GenerationService {
             events,
         };
         if self.commands.send(command).is_err() {
-            if tracking.records_rejection() {
+            if tracked {
                 self.metrics.request_rejected();
             }
             return Err("The model runtime stopped unexpectedly.".to_owned());
@@ -250,7 +225,7 @@ fn run(commands: Receiver<Command>, metrics: Arc<Metrics>) {
                 let source = command.source;
                 let events = command.events.clone();
                 let error_event = command.operation.error_handler();
-                let tracked = !matches!(command.operation, Operation::Measure);
+                let tracked = command.operation.tracks_metrics();
                 if tracked {
                     metrics.request_started(source);
                 }
@@ -285,7 +260,7 @@ fn run_command(
     source: WorkSource,
 ) {
     let error_event = command.operation.error_handler();
-    let tracked = !matches!(command.operation, Operation::Measure);
+    let tracked = command.operation.tracks_metrics();
     state.idle_timeout = command.idle_timeout;
     if command.cancel.load(Ordering::Relaxed) {
         if tracked {
@@ -483,4 +458,17 @@ mod tests {
             GenerationEvent::Measured(Err(error)) if error == "stopped"
         ));
     }
+
+    #[test]
+    fn operation_tracking_is_consistent() {
+        assert!(Operation::Generate.tracks_metrics());
+        assert!(
+            Operation::Compact {
+                reason: String::new(),
+                rebuild_system_prompt: String::new(),
+            }
+            .tracks_metrics()
+        );
+        assert!(!Operation::Measure.tracks_metrics());
+    }
 }