partition/
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::sync::Arc;
16
17use common_base::hash::partition_expr_version;
18use common_error::ext::BoxedError;
19use common_meta::cache::{
20    CacheContainer, InitStrategy, Initializer, TableRoute, TableRouteCacheRef,
21};
22use common_meta::instruction::CacheIdent;
23use common_meta::rpc::router::RegionRoute;
24use moka::future::Cache;
25use snafu::ResultExt;
26use store_api::storage::{RegionId, TableId};
27
28use crate::expr::PartitionExpr;
29use crate::manager::PartitionInfoWithVersion;
30
31#[derive(Debug, Clone)]
32pub struct PhysicalPartitionInfo {
33    pub partitions: Vec<PartitionInfoWithVersion>,
34}
35
36#[derive(Debug, Clone)]
37pub enum CachedPartitionInfo {
38    Physical(Arc<PhysicalPartitionInfo>),
39    Logical(TableId),
40}
41
42impl CachedPartitionInfo {
43    /// Returns the physical partitions if the cached partition info is physical.
44    pub fn into_physical(self) -> Option<Arc<PhysicalPartitionInfo>> {
45        match self {
46            CachedPartitionInfo::Physical(partitions) => Some(partitions),
47            CachedPartitionInfo::Logical(_) => None,
48        }
49    }
50}
51
52pub type PartitionInfoCache = CacheContainer<TableId, CachedPartitionInfo, CacheIdent>;
53
54pub type PartitionInfoCacheRef = Arc<PartitionInfoCache>;
55
56pub fn create_partitions_with_version_from_region_routes(
57    table_id: TableId,
58    region_routes: &[RegionRoute],
59) -> common_meta::error::Result<Vec<PartitionInfoWithVersion>> {
60    let mut partitions = Vec::with_capacity(region_routes.len());
61    for r in region_routes {
62        // Ignore regions marked as reject-all-writes; they should not participate
63        // in writable partition-cache construction.
64        if r.is_ignore_all_writes() {
65            continue;
66        }
67
68        let expr_json = r.region.partition_expr();
69        let partition_expr_version = if expr_json.is_empty() {
70            None
71        } else {
72            Some(partition_expr_version(Some(expr_json.as_str())))
73        };
74        let partition_expr = PartitionExpr::from_json_str(expr_json.as_str())
75            .map_err(BoxedError::new)
76            .context(common_meta::error::ExternalSnafu)?;
77        let id = RegionId::new(table_id, r.region.id.region_number());
78        partitions.push(PartitionInfoWithVersion {
79            id,
80            partition_expr,
81            partition_expr_version,
82        });
83    }
84
85    Ok(partitions)
86}
87
88fn init_factory(
89    table_route_cache: TableRouteCacheRef,
90) -> Initializer<TableId, CachedPartitionInfo> {
91    Arc::new(move |table_id: &TableId| {
92        let table_route_cache = table_route_cache.clone();
93        Box::pin(async move {
94            let Some(table_route) = table_route_cache.get(*table_id).await? else {
95                return Ok(None);
96            };
97
98            let cached = match table_route.as_ref() {
99                TableRoute::Physical(physical) => {
100                    let partitions = create_partitions_with_version_from_region_routes(
101                        *table_id,
102                        &physical.region_routes,
103                    )?;
104
105                    CachedPartitionInfo::Physical(Arc::new(PhysicalPartitionInfo { partitions }))
106                }
107                TableRoute::Logical(logical) => {
108                    let table_id = logical.physical_table_id();
109                    CachedPartitionInfo::Logical(table_id)
110                }
111            };
112
113            Ok(Some(cached))
114        })
115    })
116}
117
118pub fn new_partition_info_cache(
119    name: String,
120    cache: Cache<TableId, CachedPartitionInfo>,
121    table_route_cache: TableRouteCacheRef,
122) -> PartitionInfoCache {
123    CacheContainer::with_strategy(
124        name,
125        cache,
126        Box::new(|cache, idents| {
127            Box::pin(async move {
128                for ident in idents {
129                    if let CacheIdent::TableId(table_id) = ident {
130                        cache.invalidate(table_id).await
131                    }
132                }
133                Ok(())
134            })
135        }),
136        init_factory(table_route_cache),
137        |ident| matches!(ident, CacheIdent::TableId(_)),
138        InitStrategy::VersionChecked,
139    )
140}
141
142#[cfg(test)]
143mod tests {
144    use common_base::hash::partition_expr_version;
145    use common_meta::rpc::router::{Region, RegionRoute, WriteRoutePolicy};
146    use store_api::storage::RegionId;
147
148    use super::create_partitions_with_version_from_region_routes;
149
150    #[test]
151    fn test_create_partitions_with_version_excludes_reject_all_writes() {
152        let table_id = 1024;
153        let expr_json =
154            r#"{"Expr":{"lhs":{"Column":"a"},"op":"GtEq","rhs":{"Value":{"UInt32":10}}}}"#;
155        let region_routes = vec![
156            RegionRoute {
157                region: Region {
158                    id: RegionId::new(table_id, 1),
159                    partition_expr: expr_json.to_string(),
160                    ..Default::default()
161                },
162                leader_peer: None,
163                follower_peers: vec![],
164                leader_state: None,
165                leader_down_since: None,
166                write_route_policy: Some(WriteRoutePolicy::IgnoreAllWrites),
167            },
168            RegionRoute {
169                region: Region {
170                    id: RegionId::new(table_id, 2),
171                    partition_expr: expr_json.to_string(),
172                    ..Default::default()
173                },
174                ..Default::default()
175            },
176        ];
177
178        let partitions =
179            create_partitions_with_version_from_region_routes(table_id, &region_routes).unwrap();
180        assert_eq!(1, partitions.len());
181        assert_eq!(RegionId::new(table_id, 2), partitions[0].id);
182        assert_eq!(
183            Some(partition_expr_version(Some(expr_json))),
184            partitions[0].partition_expr_version
185        );
186    }
187}
partition/cache.rs

partition/
cache.rs
