vik 0.1.2

Vik is an issue-driven coding workflow automation tool.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139

//! Channel-driven runtime orchestrator.
//!
//! The top-level orchestrator owns intake, shutdown, and drain. Stage
//! matching, issue preparation, hook execution, and session state live
//! behind [`session_manager::StageSessionManager`].

mod event;
mod intake;
mod session_manager;

use std::sync::Arc;

use thiserror::Error;
use tokio_util::sync::CancellationToken;

use crate::workflow::Workflow;

use self::event::{IntakeEvent, OrchestratorEvent, event_channel};
use self::intake::IntakeLoop;
use self::session_manager::StageSessionManager;

#[derive(Debug, Error)]
pub enum OrchestratorError {
  #[error("orchestrator event channel closed while work was still running")]
  EventChannelClosed,
}

pub struct Orchestrator {
  workflow: Arc<Workflow>,
  sessions: StageSessionManager,
}

impl Orchestrator {
  pub fn new(workflow: Workflow) -> Self {
    let workflow = Arc::new(workflow);
    let sessions = StageSessionManager::new(Arc::clone(&workflow));

    Self { workflow, sessions }
  }

  /// Drive intake until shutdown or natural drain.
  ///
  /// Natural exit requires intake to stop and all stage sessions to drain.
  /// Hard shutdown cancels sessions and aborts intake without waiting for
  /// the manager channel to drain.
  pub async fn run(&mut self, shutdown: CancellationToken) -> Result<(), OrchestratorError> {
    let (producer, mut consumer) = event_channel();
    let intake = IntakeLoop::new(Arc::clone(&self.workflow), producer);
    let intake_handle = intake.start(shutdown.clone());
    let mut intake_stopped = false;
    let mut intake_closed = false;

    loop {
      if intake_stopped && self.sessions.is_empty() {
        return Ok(());
      }

      tokio::select! {
        biased;

        _ = shutdown.cancelled() => {
          self.sessions.cancel_all().await;
          intake_handle.abort();
          return Ok(());
        }

        event = consumer.recv(), if !intake_closed => {
          match event {
            Some(event) => {
              if self.handle_event(event) {
                intake_stopped = true;
              }
            }
            None => {
              intake_closed = true;
              if self.sessions.is_empty() {
                return Ok(());
              }
              if !intake_stopped {
                return Err(OrchestratorError::EventChannelClosed);
              }
            },
          }
        }

        received = self.sessions.recv(), if !self.sessions.is_empty() => {
          if received.is_none() {
            return Err(OrchestratorError::EventChannelClosed);
          }
        }
      }
    }
  }

  /// Returns `true` when intake has stopped. The caller still waits for
  /// `StageSessionManager` to drain before exiting.
  fn handle_event(&mut self, event: OrchestratorEvent) -> bool {
    match event {
      OrchestratorEvent::Intake(IntakeEvent::Issue(issue)) => {
        self.sessions.try_run_issue(issue);
        false
      },
      OrchestratorEvent::Intake(IntakeEvent::Failed(error)) => {
        tracing::info_span!("intake").in_scope(|| {
          tracing::error!(error = %error, "intake cycle failed");
        });
        false
      },
      OrchestratorEvent::Intake(IntakeEvent::Stopped) => true,
    }
  }
}

#[cfg(test)]
mod tests {
  use tracing_subscriber::{Registry, layer::SubscriberExt};

  use super::*;
  use crate::logging::tests::{CaptureLayer, captured_event};

  #[test]
  fn intake_failure_logs_inside_intake_span() {
    let (layer, events) = CaptureLayer::new();
    let subscriber = Registry::default().with(layer);
    let _default = tracing::subscriber::set_default(subscriber);

    let mut orchestrator = Orchestrator::new(Workflow::builder().workspace_root("workspace").build());

    let stopped = orchestrator.handle_event(OrchestratorEvent::Intake(IntakeEvent::Failed(
      "pull failed".to_string(),
    )));

    assert!(!stopped);
    let events = events.lock().expect("events mutex");
    let event = captured_event(&events, "intake cycle failed");
    assert_eq!(event["spans"][0]["name"], "intake");
    assert!(event.get("phase").is_none());
  }
}