opsqueue 0.30.3

lightweight batch processing queue for heavy loads
Documentation 
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

use std::collections::HashSet;
use std::sync::Arc;
use std::sync::Mutex;

use axum_prometheus::metrics::histogram;
use opentelemetry::trace::TraceContextExt;
use tracing::Span;
use tracing_opentelemetry::OpenTelemetrySpanExt;

use crate::common::chunk;
use crate::common::errors::DatabaseError;
use crate::common::{
    chunk::{Chunk, ChunkId},
    submission::Submission,
};
use crate::consumer::strategy;

use super::CompleterMessage;
use super::ServerState;
use crate::common::errors::{IncorrectUsage, LimitIsZero, E};

#[derive(Debug, Clone)]
pub struct ConsumerState {
    server_state: Arc<ServerState>,
    // The following are the consumer-specific chunks that are currently reserved.
    reservations: Arc<Mutex<HashSet<ChunkId>>>,
}

impl Drop for ConsumerState {
    fn drop(&mut self) {
        let reservations = self.reservations.lock().unwrap();

        // We're not tracking chunk durations that are unreserved during consumer shutdown,
        // as those will be by definition unfinished
        self.server_state
            .dispatcher
            .finish_reservations_sync(reservations.iter());
    }
}

impl ConsumerState {
    pub fn new(server_state: &Arc<ServerState>) -> Self {
        Self {
            // pool: server_state.pool.read_pool.clone(),
            // dispatcher: server_state.dispatcher.clone(),
            reservations: Arc::new(Mutex::new(HashSet::new())),
            server_state: server_state.clone(),
        }
    }

    #[tracing::instrument(skip(self, stale_chunks_notifier))]
    #[allow(clippy::type_complexity)]
    pub async fn fetch_and_reserve_chunks(
        &mut self,
        strategy: strategy::Strategy,
        limit: usize,
        stale_chunks_notifier: &tokio::sync::mpsc::UnboundedSender<ChunkId>,
    ) -> Result<Vec<(Chunk, Submission)>, E<DatabaseError, IncorrectUsage<LimitIsZero>>> {
        let start = tokio::time::Instant::now();
        if limit == 0 {
            return Err(E::R(IncorrectUsage(LimitIsZero())));
        }

        let new_reservations = self
            .server_state
            .dispatcher
            .fetch_and_reserve_chunks(
                self.server_state.pool.reader_pool(),
                strategy.clone(),
                limit,
                stale_chunks_notifier,
            )
            .await?;

        self.reservations.lock().expect("No poison").extend(
            new_reservations.iter().map(|(chunk, _submission)| {
                ChunkId::from((chunk.submission_id, chunk.chunk_index))
            }),
        );

        // Link the consumer's trace with the submission's existing trace context
        if new_reservations.len() == 1 {
            let submission = &new_reservations[0].1;
            let context = crate::tracing::json_to_context(&submission.otel_trace_carrier);
            Span::current().set_parent(context);
        } else {
            for (_, submission) in &new_reservations {
                let context = crate::tracing::json_to_context(&submission.otel_trace_carrier);
                Span::current().add_link(context.span().span_context().clone());
            }
        }

        histogram!(
            crate::prometheus::CONSUMER_FETCH_AND_RESERVE_CHUNKS_HISTOGRAM,
            &[
                ("limit", limit.to_string()),
                ("strategy", format!("{strategy:?}"))
            ]
        )
        .record(start.elapsed());
        Ok(new_reservations)
    }

    #[tracing::instrument(skip(self, output_content))]
    pub async fn complete_chunk(&mut self, id: ChunkId, output_content: chunk::Content) {
        // Only possible error indicates sender is closed, which means we're shutting down
        let _ = self
            .server_state
            .completer_tx
            .send(CompleterMessage::Complete {
                id,
                output_content,
                reservations: self.reservations.clone(),
            })
            .await;
    }

    pub async fn fail_chunk(&mut self, id: ChunkId, failure: String) {
        let _ = self
            .server_state
            .completer_tx
            .send(CompleterMessage::Fail {
                id,
                failure,
                reservations: self.reservations.clone(),
            })
            .await;
    }
}