sui_indexer_alt_reader/
system_package_task.rs

1// Copyright (c) Mysten Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4use std::sync::Arc;
5use std::time::Duration;
6
7use diesel::QueryableByName;
8use diesel::sql_types::BigInt;
9use diesel::sql_types::Bytea;
10use move_core_types::account_address::AccountAddress;
11use sui_futures::service::Service;
12use sui_sql_macro::query;
13use tokio::time;
14use tracing::info;
15use tracing::warn;
16
17use crate::package_resolver::PackageCache;
18use crate::pg_reader::PgReader;
19
20#[derive(clap::Args, Debug, Clone)]
21pub struct SystemPackageTaskArgs {
22    /// How long to wait between checking for epoch changes.
23    #[clap(long, default_value_t = Self::default().epoch_polling_interval_ms)]
24    epoch_polling_interval_ms: u64,
25}
26
27/// Background task responsible for evicting system package from the package resolver's cache after
28/// detecting an epoch boundary.
29pub struct SystemPackageTask {
30    /// Access to the database
31    pg_reader: PgReader,
32
33    /// The cached store underlying the package resolver.
34    package_cache: Arc<PackageCache>,
35
36    /// How long to wait between checks.
37    interval: Duration,
38}
39
40impl SystemPackageTaskArgs {
41    pub fn epoch_polling_interval(&self) -> Duration {
42        Duration::from_millis(self.epoch_polling_interval_ms)
43    }
44}
45
46impl SystemPackageTask {
47    pub fn new(
48        args: SystemPackageTaskArgs,
49        pg_reader: PgReader,
50        package_cache: Arc<PackageCache>,
51    ) -> Self {
52        Self {
53            pg_reader,
54            package_cache,
55            interval: args.epoch_polling_interval(),
56        }
57    }
58
59    /// Start a new task that regularly polls the database for the latest epoch and evicts system
60    /// packages if it detects that the epoch has changed (which means that a framework upgrade
61    /// could have happened).
62    ///
63    /// This operation consumes the `self` and returns a Service handle.
64    pub fn run(self) -> Service {
65        Service::new().spawn_aborting(async move {
66            let Self {
67                pg_reader,
68                package_cache,
69                interval,
70            } = self;
71
72            let mut last_epoch: i64 = 0;
73            let mut interval = time::interval(interval);
74
75            loop {
76                interval.tick().await;
77
78                let mut conn = match pg_reader.connect().await {
79                    Ok(conn) => conn,
80                    Err(e) => {
81                        warn!("Failed to connect to database: {:?}", e);
82                        continue;
83                    }
84                };
85
86                #[derive(QueryableByName, Copy, Clone)]
87                struct Watermark {
88                    #[diesel(sql_type = BigInt)]
89                    epoch_hi_inclusive: i64,
90
91                    #[diesel(sql_type = BigInt)]
92                    checkpoint_hi_inclusive: i64,
93                }
94
95                let query = query!(
96                    r#"
97                    SELECT
98                        epoch_hi_inclusive,
99                        checkpoint_hi_inclusive
100                    FROM
101                        watermarks
102                    WHERE
103                        pipeline = 'kv_packages'
104                    "#
105                );
106
107                let Watermark {
108                    epoch_hi_inclusive: next_epoch,
109                    checkpoint_hi_inclusive,
110                } = match conn.results(query).await.as_deref() {
111                    Ok([watermark]) => *watermark,
112
113                    Ok([]) => {
114                        info!("Package index isn't populated yet, no epoch information");
115                        continue;
116                    }
117
118                    Ok(_) => {
119                        warn!("Expected exactly one row from the watermarks table");
120                        continue;
121                    }
122
123                    Err(e) => {
124                        warn!("Failed to fetch latest epoch: {e}");
125                        continue;
126                    }
127                };
128
129                if next_epoch <= last_epoch {
130                    continue;
131                }
132
133                info!(last_epoch, next_epoch, "Detected epoch boundary");
134                last_epoch = next_epoch;
135
136                #[derive(QueryableByName, Clone)]
137                struct SystemPackage {
138                    #[diesel(sql_type = Bytea)]
139                    original_id: Vec<u8>,
140                }
141
142                let query = query!(
143                    r#"
144                    SELECT DISTINCT
145                        original_id
146                    FROM
147                        kv_packages
148                    WHERE
149                        is_system_package
150                    AND cp_sequence_number <= {BigInt}
151                    "#,
152                    checkpoint_hi_inclusive
153                );
154
155                let system_packages: Vec<SystemPackage> = match conn.results(query).await {
156                    Ok(system_packages) => system_packages,
157
158                    Err(e) => {
159                        warn!("Failed to fetch system packages: {e}");
160                        continue;
161                    }
162                };
163
164                let Ok(system_packages) = system_packages
165                    .into_iter()
166                    .map(|pkg| AccountAddress::from_bytes(pkg.original_id))
167                    .collect::<Result<Vec<_>, _>>()
168                else {
169                    warn!("Failed to deserialize system package addresses");
170                    continue;
171                };
172
173                info!(system_packages = ?system_packages, "Evicting...");
174                package_cache.evict(system_packages)
175            }
176        })
177    }
178}
179
180impl Default for SystemPackageTaskArgs {
181    fn default() -> Self {
182        Self {
183            epoch_polling_interval_ms: 10_000,
184        }
185    }
186}
sui_indexer_alt_reader/system_package_task.rs

sui_indexer_alt_reader/
system_package_task.rs
