pipeline/manager/
pipeline_cache.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::future::Future;
16use std::sync::Arc;
17use std::time::Duration;
18
19use common_frontend::metrics;
20use datatypes::timestamp::TimestampNanosecond;
21use moka::future::Cache;
22
23use crate::error::{CacheLoadSnafu, MultiPipelineWithDiffSchemaSnafu, Result};
24use crate::etl::Pipeline;
25use crate::manager::PipelineVersion;
26use crate::table::EMPTY_SCHEMA_NAME;
27use crate::util::{generate_pipeline_cache_key, generate_pipeline_cache_key_suffix};
28
29/// Pipeline table cache size.
30const PIPELINES_CACHE_SIZE: u64 = 10000;
31
32/// Pipeline cache is located on a separate file on purpose,
33/// to encapsulate inner cache. Only public methods are exposed.
34///
35/// `pipelines` and `original_pipelines` are keyed by the *requested* schema so
36/// a lookup is a single key probe through [`Cache::entry`];
37/// resolving it to a stored schema is the loader's job. `failover_cache` has no
38/// loader and keeps the stored-schema key.
39pub(crate) struct PipelineCache {
40    pipelines: Cache<String, Arc<Pipeline>>,
41    original_pipelines: Cache<String, PipelineContent>,
42    /// If the pipeline table is invalid, we can use this cache to prevent failures when writing logs through the pipeline
43    /// The failover cache never expires, but it will be updated when the pipelines cache is updated.
44    failover_cache: Cache<String, PipelineContent>,
45}
46
47#[derive(Clone, Debug, PartialEq, Eq)]
48pub struct PipelineContent {
49    pub name: String,
50    pub content: String,
51    pub version: TimestampNanosecond,
52    pub schema: String,
53}
54
55impl PipelineCache {
56    pub(crate) fn new(ttl: Duration) -> Self {
57        Self {
58            pipelines: Cache::builder()
59                .max_capacity(PIPELINES_CACHE_SIZE)
60                .time_to_live(ttl)
61                .name("pipelines")
62                .build(),
63            original_pipelines: Cache::builder()
64                .max_capacity(PIPELINES_CACHE_SIZE)
65                .time_to_live(ttl)
66                .name("original_pipelines")
67                .build(),
68            failover_cache: Cache::builder()
69                .max_capacity(PIPELINES_CACHE_SIZE)
70                .name("failover_cache")
71                .build(),
72        }
73    }
74
75    /// Concurrent misses share one `init` call; successful waiters count as hits.
76    pub(crate) async fn get_pipeline_with(
77        &self,
78        schema: &str,
79        name: &str,
80        version: PipelineVersion,
81        init: impl Future<Output = Result<Arc<Pipeline>>>,
82    ) -> Result<Arc<Pipeline>> {
83        let key = generate_pipeline_cache_key(schema, name, version);
84        let entry = self.pipelines.entry(key).or_try_insert_with(init).await;
85        metrics::record_cache_lookup(
86            "pipeline",
87            entry.as_ref().is_ok_and(|entry| !entry.is_fresh()),
88        );
89        entry
90            .map(|entry| entry.into_value())
91            .map_err(|error| CacheLoadSnafu { error }.build())
92    }
93
94    /// Concurrent misses share one `init` call; successful waiters count as hits.
95    pub(crate) async fn get_pipeline_str_with(
96        &self,
97        schema: &str,
98        name: &str,
99        version: PipelineVersion,
100        init: impl Future<Output = Result<PipelineContent>>,
101    ) -> Result<PipelineContent> {
102        let key = generate_pipeline_cache_key(schema, name, version);
103        let entry = self
104            .original_pipelines
105            .entry(key)
106            .or_try_insert_with(init)
107            .await;
108        metrics::record_cache_lookup(
109            "pipeline_original",
110            entry.as_ref().is_ok_and(|entry| !entry.is_fresh()),
111        );
112        entry
113            .map(|entry| entry.into_value())
114            .map_err(|error| CacheLoadSnafu { error }.build())
115    }
116
117    /// Resolves across schemas, unlike the loaded caches: a pipeline stored
118    /// under the empty schema is reachable from any schema.
119    pub(crate) async fn get_failover_cache(
120        &self,
121        schema: &str,
122        name: &str,
123        version: PipelineVersion,
124    ) -> Result<Option<PipelineContent>> {
125        for key in [
126            generate_pipeline_cache_key(EMPTY_SCHEMA_NAME, name, version),
127            generate_pipeline_cache_key(schema, name, version),
128        ] {
129            if let Some(content) = self.failover_cache.get(&key).await {
130                metrics::record_cache_lookup("pipeline_failover", true);
131                return Ok(Some(content));
132            }
133        }
134
135        // Stored under some other schema; unambiguous only if exactly one has it.
136        let suffix = generate_pipeline_cache_key_suffix(name, version);
137        let mut found = self
138            .failover_cache
139            .iter()
140            .filter(|(k, _)| k.ends_with(&suffix))
141            .collect::<Vec<_>>();
142
143        let result = match found.len() {
144            0 => Ok(None),
145            1 => Ok(Some(found.remove(0).1)),
146            _ => MultiPipelineWithDiffSchemaSnafu {
147                name: name.to_string(),
148                current_schema: schema.to_string(),
149                schemas: found
150                    .iter()
151                    .filter_map(|(k, _)| k.split_once('/').map(|k| k.0))
152                    .collect::<Vec<_>>()
153                    .join(","),
154            }
155            .fail(),
156        };
157        metrics::record_cache_lookup(
158            "pipeline_failover",
159            result.as_ref().is_ok_and(|content| content.is_some()),
160        );
161        result
162    }
163
164    pub(crate) async fn insert_failover_cache(&self, content: PipelineContent, with_latest: bool) {
165        let versioned =
166            generate_pipeline_cache_key(&content.schema, &content.name, Some(content.version));
167        let latest = generate_pipeline_cache_key(&content.schema, &content.name, None);
168
169        self.failover_cache.insert(versioned, content.clone()).await;
170        if with_latest {
171            self.failover_cache.insert(latest, content).await;
172        }
173    }
174
175    /// Dropping the stale `latest` aliases also clears the failover entries, so
176    /// the new version is written back: an outage before the first read-back
177    /// would otherwise have nothing to fall back on.
178    pub(crate) async fn on_pipeline_created(&self, content: PipelineContent) {
179        self.invalidate(&content.name, None).await;
180        self.insert_failover_cache(content, true).await;
181    }
182
183    /// Sweeps every schema and all three caches: the `latest` alias always,
184    /// plus `version` when given.
185    pub(crate) async fn invalidate(&self, name: &str, version: PipelineVersion) {
186        let mut suffixes = vec![generate_pipeline_cache_key_suffix(name, None)];
187        if version.is_some() {
188            suffixes.push(generate_pipeline_cache_key_suffix(name, version));
189        }
190
191        let ks = self
192            .pipelines
193            .iter()
194            .map(|(k, _)| k)
195            .chain(self.original_pipelines.iter().map(|(k, _)| k))
196            .chain(self.failover_cache.iter().map(|(k, _)| k))
197            .filter(|k| suffixes.iter().any(|suffix| k.ends_with(suffix)))
198            .collect::<Vec<_>>();
199
200        for k in ks {
201            let k = k.as_str();
202            self.pipelines.invalidate(k).await;
203            self.original_pipelines.invalidate(k).await;
204            self.failover_cache.invalidate(k).await;
205        }
206    }
207}
208
209#[cfg(test)]
210mod tests {
211    use std::sync::atomic::{AtomicUsize, Ordering};
212
213    use common_frontend::metrics::{CACHE_HIT, CACHE_MISS};
214    use tokio::sync::Notify;
215
216    use super::*;
217
218    /// Stored under the empty schema, i.e. visible from every schema.
219    fn content_at(version: i64) -> PipelineContent {
220        PipelineContent {
221            name: "p".to_string(),
222            content: "transform:".to_string(),
223            version: TimestampNanosecond::new(version),
224            schema: EMPTY_SCHEMA_NAME.to_string(),
225        }
226    }
227
228    #[tokio::test]
229    async fn test_concurrent_misses_run_one_loader() {
230        const CONCURRENCY: usize = 8;
231
232        let hits = CACHE_HIT.with_label_values(&["pipeline_original"]);
233        let misses = CACHE_MISS.with_label_values(&["pipeline_original"]);
234        let before = (hits.get(), misses.get());
235        let cache = PipelineCache::new(Duration::from_secs(60));
236        let loads = AtomicUsize::new(0);
237        let release = Notify::new();
238        let mut requests = (0..CONCURRENCY)
239            .map(|_| {
240                Box::pin(cache.get_pipeline_str_with("db", "p", None, async {
241                    loads.fetch_add(1, Ordering::SeqCst);
242                    release.notified().await;
243                    Ok(content_at(1))
244                }))
245            })
246            .collect::<Vec<_>>();
247        for request in &mut requests {
248            assert!(futures::poll!(request.as_mut()).is_pending());
249        }
250        assert_eq!(loads.load(Ordering::SeqCst), 1);
251        release.notify_waiters();
252        for request in requests {
253            assert_eq!(request.await.unwrap(), content_at(1));
254        }
255        assert_eq!(loads.load(Ordering::SeqCst), 1);
256        assert_eq!(
257            (hits.get(), misses.get()),
258            (before.0 + CONCURRENCY as u64 - 1, before.1 + 1)
259        );
260    }
261
262    #[tokio::test]
263    async fn test_pipeline_lookup_metrics() {
264        let cache = PipelineCache::new(Duration::from_secs(60));
265        for compiled in [false, true] {
266            let cache_type = if compiled {
267                "pipeline"
268            } else {
269                "pipeline_original"
270            };
271            let hits = CACHE_HIT.with_label_values(&[cache_type]);
272            let misses = CACHE_MISS.with_label_values(&[cache_type]);
273            let before = (hits.get(), misses.get());
274            let load = async |fail| {
275                if compiled {
276                    cache
277                        .get_pipeline_with("db", "p", None, async {
278                            if fail {
279                                return Err(crate::error::PipelineNotFoundSnafu {
280                                    name: "p",
281                                    version: None,
282                                }
283                                .build());
284                            }
285                            Ok(Arc::new(crate::parse(&crate::Content::Yaml(
286                                "transform:\n  - fields: [message]\n    type: string",
287                            ))?))
288                        })
289                        .await
290                        .map(|_| ())
291                } else {
292                    cache
293                        .get_pipeline_str_with("db", "p", None, async {
294                            if fail {
295                                return Err(crate::error::PipelineNotFoundSnafu {
296                                    name: "p",
297                                    version: None,
298                                }
299                                .build());
300                            }
301                            Ok(content_at(1))
302                        })
303                        .await
304                        .map(|_| ())
305                }
306            };
307            assert!(load(true).await.is_err());
308            load(false).await.unwrap();
309            // A warm lookup must not evaluate the failing initializer.
310            load(true).await.unwrap();
311            assert_eq!((hits.get(), misses.get()), (before.0 + 1, before.1 + 2));
312        }
313    }
314
315    #[tokio::test]
316    async fn test_failover_lookup_metrics() {
317        let hits = CACHE_HIT.with_label_values(&["pipeline_failover"]);
318        let misses = CACHE_MISS.with_label_values(&["pipeline_failover"]);
319        let before = (hits.get(), misses.get());
320        let cache = PipelineCache::new(Duration::from_secs(60));
321        assert!(
322            cache
323                .get_failover_cache("db", "p", None)
324                .await
325                .unwrap()
326                .is_none()
327        );
328        let local = PipelineContent {
329            schema: "a".into(),
330            ..content_at(1)
331        };
332        cache.insert_failover_cache(local.clone(), true).await;
333        // Direct schema lookup and cross-schema fallback each count one hit.
334        assert_eq!(
335            cache.get_failover_cache("a", "p", None).await.unwrap(),
336            Some(local.clone())
337        );
338        assert_eq!(
339            cache.get_failover_cache("db", "p", None).await.unwrap(),
340            Some(local)
341        );
342        cache
343            .insert_failover_cache(
344                PipelineContent {
345                    schema: "b".into(),
346                    ..content_at(2)
347                },
348                true,
349            )
350            .await;
351        assert!(cache.get_failover_cache("db", "p", None).await.is_err());
352        // Global pipelines win over local entries inserted before or after them.
353        for (schema, version) in [(EMPTY_SCHEMA_NAME, 3), ("c", 4)] {
354            cache
355                .insert_failover_cache(
356                    PipelineContent {
357                        schema: schema.into(),
358                        ..content_at(version)
359                    },
360                    true,
361                )
362                .await;
363            assert_eq!(
364                cache.get_failover_cache("db", "p", None).await.unwrap(),
365                Some(content_at(3))
366            );
367        }
368        assert_eq!((hits.get(), misses.get()), (before.0 + 4, before.1 + 2));
369    }
370
371    #[tokio::test]
372    async fn test_delete_drops_version_pinned_entry() {
373        let cache = PipelineCache::new(Duration::from_secs(60));
374        let content = content_at(1);
375        let version = Some(content.version);
376
377        cache
378            .get_pipeline_str_with("db", "p", version, async { Ok(content.clone()) })
379            .await
380            .unwrap();
381
382        cache.invalidate("p", version).await;
383
384        let loads = AtomicUsize::new(0);
385        cache
386            .get_pipeline_str_with("db", "p", version, async {
387                loads.fetch_add(1, Ordering::SeqCst);
388                Ok(content.clone())
389            })
390            .await
391            .unwrap();
392        assert_eq!(loads.load(Ordering::SeqCst), 1);
393    }
394
395    #[tokio::test]
396    async fn test_create_drops_stale_latest_and_primes_failover() {
397        let cache = PipelineCache::new(Duration::from_secs(60));
398        let v2 = content_at(2);
399
400        cache
401            .get_pipeline_str_with("a", "p", None, async { Ok(content_at(1)) })
402            .await
403            .unwrap();
404        cache.insert_failover_cache(content_at(1), true).await;
405
406        cache.on_pipeline_created(v2.clone()).await;
407
408        let loads = AtomicUsize::new(0);
409        let cached = cache
410            .get_pipeline_str_with("a", "p", None, async {
411                loads.fetch_add(1, Ordering::SeqCst);
412                Ok(v2.clone())
413            })
414            .await
415            .unwrap();
416        assert_eq!(loads.load(Ordering::SeqCst), 1);
417        assert_eq!(cached.version, v2.version);
418
419        let failover = cache.get_failover_cache("b", "p", None).await.unwrap();
420        assert_eq!(failover.map(|c| c.version), Some(v2.version));
421    }
422}
pipeline/manager/pipeline_cache.rs

pipeline/manager/
pipeline_cache.rs
