Line data Source code
1 : use std::sync::Arc;
2 : use std::{thread, time::Duration};
3 :
4 : use chrono::{DateTime, Utc};
5 : use postgres::{Client, NoTls};
6 : use tracing::{debug, error, info, warn};
7 :
8 : use crate::compute::ComputeNode;
9 : use compute_api::responses::ComputeStatus;
10 : use compute_api::spec::ComputeFeature;
11 :
12 : const MONITOR_CHECK_INTERVAL: Duration = Duration::from_millis(500);
13 :
14 : // Spin in a loop and figure out the last activity time in the Postgres.
15 : // Then update it in the shared state. This function never errors out.
16 : // NB: the only expected panic is at `Mutex` unwrap(), all other errors
17 : // should be handled gracefully.
18 0 : fn watch_compute_activity(compute: &ComputeNode) {
19 0 : // Suppose that `connstr` doesn't change
20 0 : let mut connstr = compute.connstr.clone();
21 0 : connstr
22 0 : .query_pairs_mut()
23 0 : .append_pair("application_name", "compute_activity_monitor");
24 0 : let connstr = connstr.as_str();
25 0 :
26 0 : // During startup and configuration we connect to every Postgres database,
27 0 : // but we don't want to count this as some user activity. So wait until
28 0 : // the compute fully started before monitoring activity.
29 0 : wait_for_postgres_start(compute);
30 0 :
31 0 : // Define `client` outside of the loop to reuse existing connection if it's active.
32 0 : let mut client = Client::connect(connstr, NoTls);
33 0 :
34 0 : let mut sleep = false;
35 0 : let mut prev_active_time: Option<f64> = None;
36 0 : let mut prev_sessions: Option<i64> = None;
37 0 :
38 0 : if compute.has_feature(ComputeFeature::ActivityMonitorExperimental) {
39 0 : info!("starting experimental activity monitor for {}", connstr);
40 : } else {
41 0 : info!("starting activity monitor for {}", connstr);
42 : }
43 :
44 0 : loop {
45 0 : // We use `continue` a lot, so it's more convenient to sleep at the top of the loop.
46 0 : // But skip the first sleep, so we can connect to Postgres immediately.
47 0 : if sleep {
48 0 : // Should be outside of the mutex lock to allow others to read while we sleep.
49 0 : thread::sleep(MONITOR_CHECK_INTERVAL);
50 0 : } else {
51 0 : sleep = true;
52 0 : }
53 :
54 0 : match &mut client {
55 0 : Ok(cli) => {
56 0 : if cli.is_closed() {
57 0 : info!("connection to Postgres is closed, trying to reconnect");
58 :
59 : // Connection is closed, reconnect and try again.
60 0 : client = Client::connect(connstr, NoTls);
61 0 : continue;
62 0 : }
63 0 :
64 0 : // This is a new logic, only enable if the feature flag is set.
65 0 : // TODO: remove this once we are sure that it works OR drop it altogether.
66 0 : if compute.has_feature(ComputeFeature::ActivityMonitorExperimental) {
67 : // First, check if the total active time or sessions across all databases has changed.
68 : // If it did, it means that user executed some queries. In theory, it can even go down if
69 : // some databases were dropped, but it's still a user activity.
70 0 : match get_database_stats(cli) {
71 0 : Ok((active_time, sessions)) => {
72 0 : let mut detected_activity = false;
73 0 :
74 0 : prev_active_time = match prev_active_time {
75 0 : Some(prev_active_time) => {
76 0 : if active_time != prev_active_time {
77 0 : detected_activity = true;
78 0 : }
79 0 : Some(active_time)
80 : }
81 0 : None => Some(active_time),
82 : };
83 0 : prev_sessions = match prev_sessions {
84 0 : Some(prev_sessions) => {
85 0 : if sessions != prev_sessions {
86 0 : detected_activity = true;
87 0 : }
88 0 : Some(sessions)
89 : }
90 0 : None => Some(sessions),
91 : };
92 :
93 0 : if detected_activity {
94 : // Update the last active time and continue, we don't need to
95 : // check backends state change.
96 0 : compute.update_last_active(Some(Utc::now()));
97 0 : continue;
98 0 : }
99 : }
100 0 : Err(e) => {
101 0 : error!("could not get database statistics: {}", e);
102 0 : continue;
103 : }
104 : }
105 0 : }
106 :
107 : // Second, if database statistics is the same, check all backends state change,
108 : // maybe there is some with more recent activity. `get_backends_state_change()`
109 : // can return None or stale timestamp, so it's `compute.update_last_active()`
110 : // responsibility to check if the new timestamp is more recent than the current one.
111 : // This helps us to discover new sessions, that did nothing yet.
112 0 : match get_backends_state_change(cli) {
113 0 : Ok(last_active) => {
114 0 : compute.update_last_active(last_active);
115 0 : }
116 0 : Err(e) => {
117 0 : error!("could not get backends state change: {}", e);
118 : }
119 : }
120 :
121 : // Finally, if there are existing (logical) walsenders, do not suspend.
122 : //
123 : // walproposer doesn't currently show up in pg_stat_replication,
124 : // but protect if it will be
125 0 : let ws_count_query = "select count(*) from pg_stat_replication where application_name != 'walproposer';";
126 0 : match cli.query_one(ws_count_query, &[]) {
127 0 : Ok(r) => match r.try_get::<&str, i64>("count") {
128 0 : Ok(num_ws) => {
129 0 : if num_ws > 0 {
130 0 : compute.update_last_active(Some(Utc::now()));
131 0 : continue;
132 0 : }
133 : }
134 0 : Err(e) => {
135 0 : warn!("failed to parse walsenders count: {:?}", e);
136 0 : continue;
137 : }
138 : },
139 0 : Err(e) => {
140 0 : warn!("failed to get list of walsenders: {:?}", e);
141 0 : continue;
142 : }
143 : }
144 : //
145 : // Don't suspend compute if there is an active logical replication subscription
146 : //
147 : // `where pid is not null` – to filter out read only computes and subscription on branches
148 : //
149 0 : let logical_subscriptions_query =
150 0 : "select count(*) from pg_stat_subscription where pid is not null;";
151 0 : match cli.query_one(logical_subscriptions_query, &[]) {
152 0 : Ok(row) => match row.try_get::<&str, i64>("count") {
153 0 : Ok(num_subscribers) => {
154 0 : if num_subscribers > 0 {
155 0 : compute.update_last_active(Some(Utc::now()));
156 0 : continue;
157 0 : }
158 : }
159 0 : Err(e) => {
160 0 : warn!("failed to parse `pg_stat_subscription` count: {:?}", e);
161 0 : continue;
162 : }
163 : },
164 0 : Err(e) => {
165 0 : warn!(
166 0 : "failed to get list of active logical replication subscriptions: {:?}",
167 : e
168 : );
169 0 : continue;
170 : }
171 : }
172 : //
173 : // Do not suspend compute if autovacuum is running
174 : //
175 0 : let autovacuum_count_query = "select count(*) from pg_stat_activity where backend_type = 'autovacuum worker'";
176 0 : match cli.query_one(autovacuum_count_query, &[]) {
177 0 : Ok(r) => match r.try_get::<&str, i64>("count") {
178 0 : Ok(num_workers) => {
179 0 : if num_workers > 0 {
180 0 : compute.update_last_active(Some(Utc::now()));
181 0 : continue;
182 0 : }
183 : }
184 0 : Err(e) => {
185 0 : warn!("failed to parse autovacuum workers count: {:?}", e);
186 0 : continue;
187 : }
188 : },
189 0 : Err(e) => {
190 0 : warn!("failed to get list of autovacuum workers: {:?}", e);
191 0 : continue;
192 : }
193 : }
194 : }
195 0 : Err(e) => {
196 0 : debug!("could not connect to Postgres: {}, retrying", e);
197 :
198 : // Establish a new connection and try again.
199 0 : client = Client::connect(connstr, NoTls);
200 : }
201 : }
202 : }
203 : }
204 :
205 : // Hang on condition variable waiting until the compute status is `Running`.
206 0 : fn wait_for_postgres_start(compute: &ComputeNode) {
207 0 : let mut state = compute.state.lock().unwrap();
208 0 : while state.status != ComputeStatus::Running {
209 0 : info!("compute is not running, waiting before monitoring activity");
210 0 : state = compute.state_changed.wait(state).unwrap();
211 0 :
212 0 : if state.status == ComputeStatus::Running {
213 0 : break;
214 0 : }
215 : }
216 0 : }
217 :
218 : // Figure out the total active time and sessions across all non-system databases.
219 : // Returned tuple is `(active_time, sessions)`.
220 : // It can return `0.0` active time or `0` sessions, which means no user databases exist OR
221 : // it was a start with skipped `pg_catalog` updates and user didn't do any queries
222 : // (or open any sessions) yet.
223 0 : fn get_database_stats(cli: &mut Client) -> anyhow::Result<(f64, i64)> {
224 0 : // Filter out `postgres` database as `compute_ctl` and other monitoring tools
225 0 : // like `postgres_exporter` use it to query Postgres statistics.
226 0 : // Use explicit 8 bytes type casts to match Rust types.
227 0 : let stats = cli.query_one(
228 0 : "SELECT coalesce(sum(active_time), 0.0)::float8 AS total_active_time,
229 0 : coalesce(sum(sessions), 0)::bigint AS total_sessions
230 0 : FROM pg_stat_database
231 0 : WHERE datname NOT IN (
232 0 : 'postgres',
233 0 : 'template0',
234 0 : 'template1'
235 0 : );",
236 0 : &[],
237 0 : );
238 0 : let stats = match stats {
239 0 : Ok(stats) => stats,
240 0 : Err(e) => {
241 0 : return Err(anyhow::anyhow!("could not query active_time: {}", e));
242 : }
243 : };
244 :
245 0 : let active_time: f64 = match stats.try_get("total_active_time") {
246 0 : Ok(active_time) => active_time,
247 0 : Err(e) => return Err(anyhow::anyhow!("could not get total_active_time: {}", e)),
248 : };
249 :
250 0 : let sessions: i64 = match stats.try_get("total_sessions") {
251 0 : Ok(sessions) => sessions,
252 0 : Err(e) => return Err(anyhow::anyhow!("could not get total_sessions: {}", e)),
253 : };
254 :
255 0 : Ok((active_time, sessions))
256 0 : }
257 :
258 : // Figure out the most recent state change time across all client backends.
259 : // If there is currently active backend, timestamp will be `Utc::now()`.
260 : // It can return `None`, which means no client backends exist or we were
261 : // unable to parse the timestamp.
262 0 : fn get_backends_state_change(cli: &mut Client) -> anyhow::Result<Option<DateTime<Utc>>> {
263 0 : let mut last_active: Option<DateTime<Utc>> = None;
264 0 : // Get all running client backends except ourself, use RFC3339 DateTime format.
265 0 : let backends = cli.query(
266 0 : "SELECT state, to_char(state_change, 'YYYY-MM-DD\"T\"HH24:MI:SS.US\"Z\"') AS state_change
267 0 : FROM pg_stat_activity
268 0 : WHERE backend_type = 'client backend'
269 0 : AND pid != pg_backend_pid()
270 0 : AND usename != 'cloud_admin';", // XXX: find a better way to filter other monitors?
271 0 : &[],
272 0 : );
273 0 :
274 0 : match backends {
275 0 : Ok(backs) => {
276 0 : let mut idle_backs: Vec<DateTime<Utc>> = vec![];
277 :
278 0 : for b in backs.into_iter() {
279 0 : let state: String = match b.try_get("state") {
280 0 : Ok(state) => state,
281 0 : Err(_) => continue,
282 : };
283 :
284 0 : if state == "idle" {
285 0 : let change: String = match b.try_get("state_change") {
286 0 : Ok(state_change) => state_change,
287 0 : Err(_) => continue,
288 : };
289 0 : let change = DateTime::parse_from_rfc3339(&change);
290 0 : match change {
291 0 : Ok(t) => idle_backs.push(t.with_timezone(&Utc)),
292 0 : Err(e) => {
293 0 : info!("cannot parse backend state_change DateTime: {}", e);
294 0 : continue;
295 : }
296 : }
297 : } else {
298 : // Found non-idle backend, so the last activity is NOW.
299 : // Return immediately, no need to check other backends.
300 0 : return Ok(Some(Utc::now()));
301 : }
302 : }
303 :
304 : // Get idle backend `state_change` with the max timestamp.
305 0 : if let Some(last) = idle_backs.iter().max() {
306 0 : last_active = Some(*last);
307 0 : }
308 : }
309 0 : Err(e) => {
310 0 : return Err(anyhow::anyhow!("could not query backends: {}", e));
311 : }
312 : }
313 :
314 0 : Ok(last_active)
315 0 : }
316 :
317 : /// Launch a separate compute monitor thread and return its `JoinHandle`.
318 0 : pub fn launch_monitor(compute: &Arc<ComputeNode>) -> thread::JoinHandle<()> {
319 0 : let compute = Arc::clone(compute);
320 0 :
321 0 : thread::Builder::new()
322 0 : .name("compute-monitor".into())
323 0 : .spawn(move || watch_compute_activity(&compute))
324 0 : .expect("cannot launch compute monitor thread")
325 0 : }
|
