datanode/region_server/
registrations.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::collections::hash_map::Entry;
16use std::collections::{HashMap, HashSet};
17use std::fmt::{Display, Formatter};
18use std::sync::{Arc, Mutex, OnceLock};
19
20use common_query::request::{
21    DynFilterPayload, INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY, InitialDynFilterReg,
22    InitialDynFilterRegs, REMOTE_DYN_FILTER_PAYLOAD_MAX_BYTES,
23};
24use common_telemetry::warn;
25use dashmap::DashMap;
26use datafusion::arrow::datatypes::{Schema, SchemaRef};
27use datafusion::execution::{SessionStateBuilder, TaskContext};
28use datafusion::physical_plan::PhysicalExpr;
29use datafusion::physical_plan::expressions::{DynamicFilterPhysicalExpr, lit};
30use datafusion_common::Result as DataFusionResult;
31use session::context::QueryContextRef;
32use session::query_id::QueryId;
33use store_api::storage::RegionId;
34
35use crate::metrics::{
36    REMOTE_DYN_FILTER_PAYLOAD_BYTES, REMOTE_DYN_FILTER_REGISTRY_ENTRIES,
37    REMOTE_DYN_FILTER_UNREGISTER_TOTAL, REMOTE_DYN_FILTER_UPDATE_DROP_TOTAL,
38    REMOTE_DYN_FILTER_UPDATE_OUTCOME_TOTAL,
39};
40
41type QueryRemoteDynFilterRegs = HashMap<RemoteDynFilterId, RegisteredDynFilter>;
42
43#[derive(Debug, Default)]
44pub(super) struct RemoteDynFilterRegistry {
45    // Keep cross-query concurrency while making each query's RDF state machine a
46    // single critical section. RDF count per query is small
47    queries: DashMap<QueryId, Arc<Mutex<QueryRemoteDynFilterRegs>>>,
48}
49
50impl RemoteDynFilterRegistry {
51    pub(super) fn new() -> Self {
52        Self::default()
53    }
54
55    fn get_or_insert_query(&self, query_id: QueryId) -> Arc<Mutex<QueryRemoteDynFilterRegs>> {
56        self.queries
57            .entry(query_id)
58            .or_insert_with(|| Arc::new(Mutex::new(HashMap::new())))
59            .clone()
60    }
61
62    fn get_query(&self, query_id: &QueryId) -> Option<Arc<Mutex<QueryRemoteDynFilterRegs>>> {
63        self.queries
64            .get(query_id)
65            .map(|query_regs| query_regs.clone())
66    }
67
68    fn remove_query_if_empty(
69        &self,
70        query_id: &QueryId,
71        expected: &Arc<Mutex<QueryRemoteDynFilterRegs>>,
72    ) {
73        // Protect against stale cleanup of an old per-query state: remove the
74        // outer entry only if it still points to the same inner mutex and that
75        // inner map is still empty.
76        self.queries.remove_if(query_id, |_, query_regs| {
77            Arc::ptr_eq(query_regs, expected) && query_regs.lock().unwrap().is_empty()
78        });
79    }
80
81    #[cfg(test)]
82    pub(super) fn inspect_query<R>(
83        &self,
84        query_id: &QueryId,
85        inspect: impl FnOnce(&QueryRemoteDynFilterRegs) -> R,
86    ) -> Option<R> {
87        self.get_query(query_id)
88            .map(|query_regs| inspect(&query_regs.lock().unwrap()))
89    }
90}
91
92#[derive(Clone, Debug, PartialEq, Eq, Hash)]
93pub(super) struct RemoteDynFilterId(String);
94
95impl RemoteDynFilterId {
96    pub(super) fn new(value: impl Into<String>) -> Self {
97        Self(value.into())
98    }
99}
100
101impl Display for RemoteDynFilterId {
102    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
103        self.0.fmt(f)
104    }
105}
106
107#[derive(Debug, Clone, Copy, PartialEq, Eq)]
108pub(super) enum RemoteDynFilterUpdateOutcome {
109    MissingRegistration,
110    Buffered,
111    Applied,
112    Idempotent,
113    Stale,
114    AlreadyComplete,
115    PayloadTooLarge,
116    DecodeFailed,
117}
118
119impl RemoteDynFilterUpdateOutcome {
120    pub fn as_str(&self) -> &'static str {
121        match self {
122            Self::MissingRegistration => "missing_registration",
123            Self::Buffered => "buffered",
124            Self::Applied => "applied",
125            Self::Idempotent => "idempotent",
126            Self::Stale => "stale",
127            Self::AlreadyComplete => "already_complete",
128            Self::PayloadTooLarge => "payload_too_large",
129            Self::DecodeFailed => "decode_failed",
130        }
131    }
132
133    fn is_drop(&self) -> bool {
134        matches!(
135            self,
136            Self::MissingRegistration
137                | Self::Stale
138                | Self::AlreadyComplete
139                | Self::PayloadTooLarge
140                | Self::DecodeFailed
141        )
142    }
143}
144
145fn record_update_outcome(outcome: RemoteDynFilterUpdateOutcome) {
146    REMOTE_DYN_FILTER_UPDATE_OUTCOME_TOTAL
147        .with_label_values(&[outcome.as_str()])
148        .inc();
149    if outcome.is_drop() {
150        REMOTE_DYN_FILTER_UPDATE_DROP_TOTAL
151            .with_label_values(&[outcome.as_str()])
152            .inc();
153    }
154}
155
156fn record_unregister_outcome(outcome: RemoteDynFilterUpdateOutcome) {
157    REMOTE_DYN_FILTER_UNREGISTER_TOTAL
158        .with_label_values(&[outcome.as_str()])
159        .inc();
160}
161
162#[derive(Debug, Clone)]
163struct PendingDynFilterUpdate {
164    payload: Vec<u8>,
165    generation: u64,
166    is_complete: bool,
167}
168
169impl PendingDynFilterUpdate {
170    fn from_initial_reg(reg: &InitialDynFilterReg) -> Option<Self> {
171        let snapshot = reg.initial_snapshot.as_ref()?;
172        match &snapshot.payload {
173            DynFilterPayload::Datafusion(payload) => Some(Self {
174                payload: payload.clone(),
175                generation: snapshot.generation,
176                is_complete: snapshot.is_complete,
177            }),
178            _ => None,
179        }
180    }
181}
182
183#[derive(Debug)]
184struct RemoteDynFilterEpochState {
185    generation: Option<u64>,
186    is_complete: bool,
187}
188
189#[derive(Debug)]
190struct RemoteDynFilterState {
191    filter: Arc<DynamicFilterPhysicalExpr>,
192    input_schema: SchemaRef,
193    epoch: Mutex<RemoteDynFilterEpochState>,
194}
195
196impl RemoteDynFilterState {
197    fn new(filter: Arc<DynamicFilterPhysicalExpr>, input_schema: SchemaRef) -> Self {
198        Self {
199            filter,
200            input_schema,
201            epoch: Mutex::new(RemoteDynFilterEpochState {
202                generation: None,
203                is_complete: false,
204            }),
205        }
206    }
207
208    fn filter(&self) -> Arc<DynamicFilterPhysicalExpr> {
209        self.filter.clone()
210    }
211
212    fn apply_update(
213        &self,
214        payload: &[u8],
215        generation: u64,
216        is_complete: bool,
217    ) -> RemoteDynFilterUpdateOutcome {
218        if !validate_update_payload_size(payload) {
219            return RemoteDynFilterUpdateOutcome::PayloadTooLarge;
220        }
221
222        let mut epoch = self.epoch.lock().unwrap();
223        if let Some(current_generation) = epoch.generation {
224            if generation < current_generation {
225                return RemoteDynFilterUpdateOutcome::Stale;
226            }
227
228            if generation == current_generation {
229                if is_complete && !epoch.is_complete {
230                    self.filter.mark_complete();
231                    epoch.is_complete = true;
232                    return RemoteDynFilterUpdateOutcome::Applied;
233                }
234                return RemoteDynFilterUpdateOutcome::Idempotent;
235            }
236        }
237
238        if epoch.is_complete {
239            return RemoteDynFilterUpdateOutcome::AlreadyComplete;
240        }
241
242        let expr = match decode_update_payload(payload, self.input_schema.as_ref()) {
243            Ok(expr) => expr,
244            Err(error) => {
245                warn!(error; "Failed to decode remote dynamic filter update payload");
246                return RemoteDynFilterUpdateOutcome::DecodeFailed;
247            }
248        };
249
250        if let Err(error) = self.filter.update(expr) {
251            warn!(error; "Failed to apply remote dynamic filter update");
252            return RemoteDynFilterUpdateOutcome::DecodeFailed;
253        }
254
255        epoch.generation = Some(generation);
256        if is_complete {
257            self.filter.mark_complete();
258            epoch.is_complete = true;
259        }
260
261        RemoteDynFilterUpdateOutcome::Applied
262    }
263}
264
265#[derive(Debug)]
266pub(super) struct RegisteredDynFilter {
267    pub(super) filter_id: RemoteDynFilterId,
268    pub(super) child_exprs_datafusion_proto: Vec<Vec<u8>>,
269    pub(super) subscriber_regions: HashSet<RegionId>,
270    runtime: Option<Arc<RemoteDynFilterState>>,
271    pending_update: Option<PendingDynFilterUpdate>,
272}
273
274impl RegisteredDynFilter {
275    fn new(
276        filter_id: RemoteDynFilterId,
277        child_exprs_datafusion_proto: Vec<Vec<u8>>,
278        pending_update: Option<PendingDynFilterUpdate>,
279        region_id: RegionId,
280    ) -> Self {
281        let mut subscriber_regions = HashSet::new();
282        subscriber_regions.insert(region_id);
283
284        Self {
285            filter_id,
286            child_exprs_datafusion_proto,
287            subscriber_regions,
288            runtime: None,
289            pending_update,
290        }
291    }
292
293    fn apply_initial_snapshot(
294        &mut self,
295        reg: &InitialDynFilterReg,
296    ) -> RemoteDynFilterUpdateOutcome {
297        let Some(snapshot) = PendingDynFilterUpdate::from_initial_reg(reg) else {
298            return RemoteDynFilterUpdateOutcome::Idempotent;
299        };
300
301        self.apply_or_buffer_update(&snapshot.payload, snapshot.generation, snapshot.is_complete)
302    }
303
304    fn register_subscriber(&mut self, region_id: RegionId) -> bool {
305        if !self.subscriber_regions.insert(region_id) {
306            warn!(
307                "Duplicate remote dynamic filter subscriber region, filter_id: {}, region_id: {}",
308                self.filter_id, region_id
309            );
310            return false;
311        }
312
313        true
314    }
315
316    fn has_subscribers(&self) -> bool {
317        !self.subscriber_regions.is_empty()
318    }
319
320    fn should_drop_after_remove(&mut self, region_id: RegionId) -> bool {
321        self.subscriber_regions.remove(&region_id);
322        !self.has_subscribers()
323    }
324
325    fn buffer_update(
326        &mut self,
327        payload: &[u8],
328        generation: u64,
329        is_complete: bool,
330    ) -> RemoteDynFilterUpdateOutcome {
331        if !validate_update_payload_size(payload) {
332            return RemoteDynFilterUpdateOutcome::PayloadTooLarge;
333        }
334
335        if let Some(pending) = self.pending_update.as_mut() {
336            if generation < pending.generation {
337                return RemoteDynFilterUpdateOutcome::Stale;
338            }
339
340            if generation == pending.generation {
341                pending.is_complete |= is_complete;
342                return RemoteDynFilterUpdateOutcome::Idempotent;
343            }
344
345            if pending.is_complete {
346                return RemoteDynFilterUpdateOutcome::AlreadyComplete;
347            }
348        }
349
350        self.pending_update = Some(PendingDynFilterUpdate {
351            payload: payload.to_vec(),
352            generation,
353            is_complete,
354        });
355        RemoteDynFilterUpdateOutcome::Buffered
356    }
357
358    fn apply_or_buffer_update(
359        &mut self,
360        payload: &[u8],
361        generation: u64,
362        is_complete: bool,
363    ) -> RemoteDynFilterUpdateOutcome {
364        if let Some(runtime) = &self.runtime {
365            return runtime.apply_update(payload, generation, is_complete);
366        }
367
368        self.buffer_update(payload, generation, is_complete)
369    }
370
371    fn decode_children(
372        &self,
373        input_schema: &Schema,
374    ) -> DataFusionResult<Vec<Arc<dyn PhysicalExpr>>> {
375        InitialDynFilterReg::new(
376            self.filter_id.to_string(),
377            self.child_exprs_datafusion_proto.clone(),
378        )
379        .decode_children(
380            remote_dyn_filter_task_context(),
381            input_schema,
382            REMOTE_DYN_FILTER_PAYLOAD_MAX_BYTES,
383        )
384    }
385
386    fn dyn_filter(&mut self, input_schema: &Schema) -> Option<Arc<dyn PhysicalExpr>> {
387        let children = match self.decode_children(input_schema) {
388            Ok(children) => children,
389            Err(error) => {
390                warn!(error; "Failed to decode remote dynamic filter initial children");
391                return None;
392            }
393        };
394
395        let runtime = match &self.runtime {
396            Some(runtime) => runtime.clone(),
397            None => {
398                let filter = Arc::new(DynamicFilterPhysicalExpr::new(children.clone(), lit(true)));
399                let runtime = Arc::new(RemoteDynFilterState::new(
400                    filter,
401                    Arc::new(input_schema.clone()),
402                ));
403                if let Some(pending) = self.pending_update.take() {
404                    let outcome = runtime.apply_update(
405                        &pending.payload,
406                        pending.generation,
407                        pending.is_complete,
408                    );
409                    if matches!(outcome, RemoteDynFilterUpdateOutcome::DecodeFailed) {
410                        warn!(
411                            "Dropped buffered remote dynamic filter update after decode failure, filter_id: {}, generation: {}",
412                            self.filter_id, pending.generation
413                        );
414                    }
415                }
416                self.runtime = Some(runtime.clone());
417                runtime
418            }
419        };
420
421        match runtime.filter().with_new_children(children) {
422            Ok(expr) => Some(expr),
423            Err(error) => {
424                warn!(error; "Failed to remap remote dynamic filter children for scan");
425                None
426            }
427        }
428    }
429
430    fn deactivate(&self) {
431        if let Some(runtime) = &self.runtime {
432            runtime.filter.mark_complete();
433        }
434    }
435}
436
437impl Drop for RegisteredDynFilter {
438    fn drop(&mut self) {
439        self.deactivate();
440    }
441}
442
443pub(super) fn initial_dyn_filter_regs_from_query_ctx(
444    query_ctx: &QueryContextRef,
445) -> Option<InitialDynFilterRegs> {
446    let registrations =
447        query_ctx.extension(INITIAL_REMOTE_DYN_FILTER_REGISTRATIONS_EXTENSION_KEY)?;
448    match InitialDynFilterRegs::from_extension_value(registrations) {
449        Ok(registrations) => match registrations.validate_default_bounds() {
450            Ok(()) => Some(registrations),
451            Err(error) => {
452                warn!(error; "Initial remote dyn filter registrations exceeded Task 03 bounds");
453                None
454            }
455        },
456        Err(error) => {
457            warn!(error; "Failed to decode initial remote dyn filter registrations from query context");
458            None
459        }
460    }
461}
462
463pub(super) fn register_initial_dyn_filter_regs(
464    regs_by_query: &RemoteDynFilterRegistry,
465    query_id: &QueryId,
466    region_id: RegionId,
467    regs: &InitialDynFilterRegs,
468) -> Vec<RemoteDynFilterId> {
469    if regs.is_empty() {
470        return Vec::new();
471    }
472
473    if let Err(error) = regs.validate_default_bounds() {
474        warn!(error; "Ignored invalid initial dyn filter registrations for query_id {}", query_id);
475        return Vec::new();
476    }
477
478    let query_regs = regs_by_query.get_or_insert_query(*query_id);
479    let mut query_regs = query_regs.lock().unwrap();
480    let mut registered_filter_ids = Vec::with_capacity(regs.regs.len());
481
482    for reg in &regs.regs {
483        let filter_id = RemoteDynFilterId::new(reg.filter_id.clone());
484        match query_regs.entry(filter_id.clone()) {
485            Entry::Occupied(mut entry) => {
486                let registered = entry.get_mut();
487                if registered.child_exprs_datafusion_proto != reg.child_exprs_datafusion_proto {
488                    warn!(
489                        "Remote dynamic filter registration reused filter_id with different children, query_id: {}, filter_id: {}, region_id: {}",
490                        query_id, filter_id, region_id
491                    );
492                }
493                if registered.register_subscriber(region_id) {
494                    registered_filter_ids.push(filter_id);
495                }
496                let _ = registered.apply_initial_snapshot(reg);
497            }
498            Entry::Vacant(entry) => {
499                entry.insert(RegisteredDynFilter::new(
500                    filter_id.clone(),
501                    reg.child_exprs_datafusion_proto.clone(),
502                    PendingDynFilterUpdate::from_initial_reg(reg),
503                    region_id,
504                ));
505                REMOTE_DYN_FILTER_REGISTRY_ENTRIES.inc();
506                registered_filter_ids.push(filter_id);
507            }
508        }
509    }
510
511    registered_filter_ids
512}
513
514pub(super) fn remote_dyn_filter_exprs_for_initial_regs(
515    regs_by_query: &RemoteDynFilterRegistry,
516    query_id: &QueryId,
517    initial_regs: &InitialDynFilterRegs,
518    input_schema: &Schema,
519) -> Vec<Arc<dyn PhysicalExpr>> {
520    let Some(query_regs) = regs_by_query.get_query(query_id) else {
521        return Vec::new();
522    };
523
524    let mut query_regs = query_regs.lock().unwrap();
525    initial_regs
526        .regs
527        .iter()
528        .filter_map(|reg| {
529            let filter_id = RemoteDynFilterId::new(reg.filter_id.clone());
530            let registered = query_regs.get_mut(&filter_id)?;
531            registered.dyn_filter(input_schema)
532        })
533        .collect()
534}
535
536pub(super) fn apply_remote_dyn_filter_update(
537    regs_by_query: &RemoteDynFilterRegistry,
538    query_id: &QueryId,
539    filter_id: &RemoteDynFilterId,
540    payload: &[u8],
541    generation: u64,
542    is_complete: bool,
543) -> RemoteDynFilterUpdateOutcome {
544    if !validate_update_payload_size(payload) {
545        warn!(
546            "Ignored oversized remote dynamic filter update, query_id: {}, filter_id: {}, payload_size: {}, max_payload_size: {}",
547            query_id,
548            filter_id,
549            payload.len(),
550            REMOTE_DYN_FILTER_PAYLOAD_MAX_BYTES
551        );
552        let outcome = RemoteDynFilterUpdateOutcome::PayloadTooLarge;
553        record_update_outcome(outcome);
554        return outcome;
555    }
556
557    REMOTE_DYN_FILTER_PAYLOAD_BYTES.observe(payload.len() as f64);
558
559    let Some(query_regs) = regs_by_query.get_query(query_id) else {
560        warn!(
561            "Ignored remote dynamic filter update without query registration, query_id: {}, filter_id: {}",
562            query_id, filter_id
563        );
564        let outcome = RemoteDynFilterUpdateOutcome::MissingRegistration;
565        record_update_outcome(outcome);
566        return outcome;
567    };
568
569    let mut query_regs = query_regs.lock().unwrap();
570    let Some(registered) = query_regs.get_mut(filter_id) else {
571        warn!(
572            "Ignored remote dynamic filter update without filter registration, query_id: {}, filter_id: {}",
573            query_id, filter_id
574        );
575        let outcome = RemoteDynFilterUpdateOutcome::MissingRegistration;
576        record_update_outcome(outcome);
577        return outcome;
578    };
579
580    let outcome = registered.apply_or_buffer_update(payload, generation, is_complete);
581    record_update_outcome(outcome);
582
583    outcome
584}
585
586pub(super) fn unregister_remote_dyn_filter(
587    regs_by_query: &RemoteDynFilterRegistry,
588    query_id: &QueryId,
589    filter_id: &RemoteDynFilterId,
590) -> RemoteDynFilterUpdateOutcome {
591    let Some(query_regs) = regs_by_query.get_query(query_id) else {
592        warn!(
593            "Ignored remote dynamic filter unregister without query registration, query_id: {}, filter_id: {}",
594            query_id, filter_id
595        );
596        let outcome = RemoteDynFilterUpdateOutcome::MissingRegistration;
597        record_unregister_outcome(outcome);
598        return outcome;
599    };
600
601    let (registered, should_remove_query) = {
602        let mut locked = query_regs.lock().unwrap();
603        let Some(registered) = locked.remove(filter_id) else {
604            warn!(
605                "Ignored remote dynamic filter unregister without filter registration, query_id: {}, filter_id: {}",
606                query_id, filter_id
607            );
608            let outcome = RemoteDynFilterUpdateOutcome::MissingRegistration;
609            record_unregister_outcome(outcome);
610            return outcome;
611        };
612        REMOTE_DYN_FILTER_REGISTRY_ENTRIES.dec();
613        let should_remove_query = locked.is_empty();
614        (registered, should_remove_query)
615    };
616
617    drop(registered);
618    if should_remove_query {
619        regs_by_query.remove_query_if_empty(query_id, &query_regs);
620    }
621
622    let outcome = RemoteDynFilterUpdateOutcome::Applied;
623    record_unregister_outcome(outcome);
624
625    outcome
626}
627
628pub(super) fn remove_initial_dyn_filter_regs(
629    regs_by_query: &RemoteDynFilterRegistry,
630    query_id: &QueryId,
631    region_id: RegionId,
632    filter_ids: &[RemoteDynFilterId],
633) {
634    if filter_ids.is_empty() {
635        return;
636    }
637
638    let Some(query_regs) = regs_by_query.get_query(query_id) else {
639        return;
640    };
641
642    let (removed_filters, should_remove_query) = {
643        let mut locked = query_regs.lock().unwrap();
644        let mut removed_filters = Vec::new();
645
646        for filter_id in filter_ids {
647            let should_remove_filter = locked
648                .get_mut(filter_id)
649                .map(|registered| registered.should_drop_after_remove(region_id))
650                .unwrap_or(false);
651
652            if should_remove_filter && let Some(registered) = locked.remove(filter_id) {
653                REMOTE_DYN_FILTER_REGISTRY_ENTRIES.dec();
654                removed_filters.push(registered);
655            }
656        }
657
658        let should_remove_query = locked.is_empty();
659        (removed_filters, should_remove_query)
660    };
661
662    drop(removed_filters);
663    if should_remove_query {
664        regs_by_query.remove_query_if_empty(query_id, &query_regs);
665    }
666}
667
668fn decode_update_payload(
669    payload: &[u8],
670    input_schema: &Schema,
671) -> DataFusionResult<Arc<dyn PhysicalExpr>> {
672    DynFilterPayload::Datafusion(payload.to_vec()).decode_datafusion_expr(
673        remote_dyn_filter_task_context(),
674        input_schema,
675        REMOTE_DYN_FILTER_PAYLOAD_MAX_BYTES,
676    )
677}
678
679fn remote_dyn_filter_task_context() -> &'static TaskContext {
680    static TASK_CONTEXT: OnceLock<TaskContext> = OnceLock::new();
681
682    TASK_CONTEXT.get_or_init(|| {
683        // RDF payloads can contain DataFusion built-in scalar functions. For
684        // example, multi-column join dynamic filters use `struct(...) IN (...)`.
685        // `TaskContext::default()` has an empty function registry and cannot
686        // decode those expressions.
687        let session_state = SessionStateBuilder::new().with_default_features().build();
688        TaskContext::from(&session_state)
689    })
690}
691
692fn validate_update_payload_size(payload: &[u8]) -> bool {
693    payload.len() <= REMOTE_DYN_FILTER_PAYLOAD_MAX_BYTES
694}
695
696#[cfg(test)]
697mod tests {
698    use super::*;
699
700    #[test]
701    fn remote_dyn_filter_same_query_reuses_one_inner_lock() {
702        let regs_by_query = RemoteDynFilterRegistry::new();
703        let query_id = QueryId::new();
704
705        let first = regs_by_query.get_or_insert_query(query_id);
706        let second = regs_by_query.get_or_insert_query(query_id);
707
708        assert!(Arc::ptr_eq(&first, &second));
709    }
710
711    #[test]
712    fn remote_dyn_filter_stale_query_remove_does_not_remove_new_query_state() {
713        let regs_by_query = RemoteDynFilterRegistry::new();
714        let query_id = QueryId::new();
715        let stale_query_regs = regs_by_query.get_or_insert_query(query_id);
716
717        regs_by_query.remove_query_if_empty(&query_id, &stale_query_regs);
718        assert!(regs_by_query.get_query(&query_id).is_none());
719
720        let new_query_regs = regs_by_query.get_or_insert_query(query_id);
721        let filter_id = RemoteDynFilterId::new("filter-1");
722        let region_id = RegionId::new(1024, 7);
723        new_query_regs.lock().unwrap().insert(
724            filter_id.clone(),
725            RegisteredDynFilter::new(filter_id.clone(), vec![], None, region_id),
726        );
727
728        regs_by_query.remove_query_if_empty(&query_id, &stale_query_regs);
729
730        let current_query_regs = regs_by_query.get_query(&query_id).unwrap();
731        assert!(Arc::ptr_eq(&current_query_regs, &new_query_regs));
732        assert_eq!(current_query_regs.lock().unwrap().len(), 1);
733    }
734
735    #[test]
736    fn remote_dyn_filter_register_uses_entry_to_merge_same_filter_subscribers() {
737        let regs_by_query = RemoteDynFilterRegistry::new();
738        let query_id = QueryId::new();
739        let first_region_id = RegionId::new(1024, 7);
740        let second_region_id = RegionId::new(1024, 8);
741        let regs = InitialDynFilterRegs::new(vec![InitialDynFilterReg::new("filter-1", vec![])]);
742
743        register_initial_dyn_filter_regs(&regs_by_query, &query_id, first_region_id, &regs);
744        register_initial_dyn_filter_regs(&regs_by_query, &query_id, second_region_id, &regs);
745
746        let query_regs = regs_by_query.get_query(&query_id).unwrap();
747        let query_regs = query_regs.lock().unwrap();
748        assert_eq!(query_regs.len(), 1);
749        let registered = query_regs.get(&RemoteDynFilterId::new("filter-1")).unwrap();
750        assert_eq!(registered.subscriber_regions.len(), 2);
751        assert!(registered.subscriber_regions.contains(&first_region_id));
752        assert!(registered.subscriber_regions.contains(&second_region_id));
753    }
754}
datanode/region_server/registrations.rs

datanode/region_server/
registrations.rs
