1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
use std::process;
use clap::Parser;
use istat::bar::Bar;
use istat::cli::Cli;
use istat::config::AppConfig;
use istat::context::{Context, SharedState, StopAction};
use istat::dispatcher::Dispatcher;
use istat::error::Result;
use istat::i3::header::I3BarHeader;
use istat::i3::ipc::handle_click_events;
use istat::i3::I3Item;
use istat::ipc::{create_ipc_socket, handle_ipc_events, IpcContext};
use istat::signals::handle_signals;
use istat::util::{local_block_on, RcCell, UrgentTimer};
use tokio::sync::mpsc::{self, Receiver};
use tokio::time::Instant;
use tokio_util::sync::CancellationToken;
enum RuntimeStopReason {
Shutdown,
}
fn main() {
match start_runtime() {
Ok(RuntimeStopReason::Shutdown) => {}
Err(err) => {
log::error!("{}", err);
process::exit(1);
}
}
}
fn start_runtime() -> Result<RuntimeStopReason> {
pretty_env_logger::try_init_timed()?;
let args = Cli::parse();
let (result, runtime) = local_block_on(async_main(args))?;
// NOTE: since we use tokio's stdin implementation which spawns a background thread and blocks,
// we have to shutdown the runtime ourselves here. If we didn't, then when the runtime is
// dropped it would block indefinitely until that background thread unblocked (i.e., another
// JSON line from i3).
// Thus, if anything other than the stdin task fails, we have to manually shut it down here.
// See: https://github.com/tokio-rs/tokio/discussions/5684
runtime.shutdown_background();
result
}
async fn async_main(args: Cli) -> Result<RuntimeStopReason> {
let config = RcCell::new(AppConfig::read(args).await?);
// create socket first, so it's ready before anything is written to stdout
let socket = create_ipc_socket(&config).await?;
// create i3 bar and spawn tasks for each bar item
let (bar, dispatcher) = setup_i3_bar(&config)?;
// handle incoming signals
let signal_handle = handle_signals(config.clone(), dispatcher.clone())?;
// used to handle app shutdown
let token = CancellationToken::new();
// ipc context
let ipc_ctx = IpcContext::new(
bar.clone(),
token.clone(),
config.clone(),
dispatcher.clone(),
);
// handle our inputs: i3's IPC and our own IPC
let result = tokio::select! {
Err(err) = handle_ipc_events(socket, ipc_ctx) => Err(err),
Err(err) = handle_click_events(dispatcher.clone()) => Err(err),
_ = token.cancelled() => Ok(RuntimeStopReason::Shutdown),
};
// if we reach here, then something went wrong, so clean up
signal_handle.close();
return result;
}
fn setup_i3_bar(config: &RcCell<AppConfig>) -> Result<(RcCell<Bar>, RcCell<Dispatcher>)> {
let item_count = config.items.len();
// shared state
let state = SharedState::new();
// A list of items which represents the i3 bar
let bar = RcCell::new(Bar::new(item_count));
// Used to send events to each bar item, and also to trigger updates of the bar
let (update_tx, update_rx) = mpsc::channel(1);
let dispatcher = RcCell::new(Dispatcher::new(update_tx, item_count));
// Used by items to send updates back to the bar
let (item_tx, item_rx) = mpsc::channel(item_count + 1);
// Iterate config and create bar items
for (idx, item) in config.items.iter().enumerate() {
let bar_item = item.to_bar_item();
// all cheaply cloneable (smart pointers, senders, etc)
let mut bar = bar.clone();
let state = state.clone();
let config = config.clone();
let item_tx = item_tx.clone();
let mut dispatcher = dispatcher.clone();
tokio::task::spawn_local(async move {
let mut retries = 0;
let mut last_start;
loop {
last_start = Instant::now();
let (event_tx, event_rx) = mpsc::channel(32);
dispatcher.set(idx, event_tx);
let ctx = Context::new(
config.clone(),
state.clone(),
item_tx.clone(),
event_rx,
idx,
);
let fut = bar_item.start(ctx);
match fut.await {
Ok(StopAction::Restart) => {
// reset retries if no retries have occurred in the last 5 minutes
if last_start.elapsed().as_secs() > 60 * 5 {
retries = 0;
}
// restart if we haven't exceeded limit
if retries < 3 {
log::warn!("item[{}] requested restart...", idx);
retries += 1;
continue;
}
// we exceeded the limit, so error out
log::error!("item[{}] stopped, exceeded max retries", idx);
let theme = config.theme.clone();
bar[idx] = I3Item::new("MAX RETRIES")
.color(theme.bg)
.background_color(theme.red);
break;
}
// since this item has terminated, remove its entry from the bar
action @ Ok(StopAction::Complete) | action @ Ok(StopAction::Remove) => {
log::info!("item[{}] finished running", idx);
dispatcher.remove(idx);
// Remove this item if requested
if matches!(action, Ok(StopAction::Remove)) {
// NOTE: wait for all tasks in queue so any remaining item updates are flushed and processed
// before we set it for the last time here
tokio::task::yield_now().await;
// replace with an empty item
bar[idx] = I3Item::empty();
}
break;
}
// unexpected error, log and display an error block
Err(e) => {
log::error!("item[{}] exited with error: {}", idx, e);
// replace with an error item
let theme = config.theme.clone();
bar[idx] = I3Item::new("ERROR")
.color(theme.bg)
.background_color(theme.red);
break;
}
}
}
});
}
// setup listener for handling item updates and printing the bar to STDOUT
handle_item_updates(config.clone(), item_rx, update_rx, bar.clone())?;
Ok((bar, dispatcher))
}
// task to manage updating the bar and printing it as JSON
fn handle_item_updates(
config: RcCell<AppConfig>,
mut item_rx: Receiver<(I3Item, usize)>,
mut update_rx: Receiver<()>,
mut bar: RcCell<Bar>,
) -> Result<()> {
// output first parts of the i3 bar protocol - the header
println!("{}", serde_json::to_string(&I3BarHeader::default())?);
// and the opening bracket for the "infinite array"
println!("[");
tokio::task::spawn_local(async move {
let item_names = config.item_idx_to_name();
let mut urgent_timer = UrgentTimer::new();
loop {
// enable urgent timer if any item is urgent
urgent_timer.toggle(bar.any_urgent());
tokio::select! {
// the urgent timer triggered, so update the timer and start it again
// this logic makes urgent items "flash" between two coloured states
() = urgent_timer.wait() => urgent_timer.reset(),
// a manual update was requested
Some(()) = update_rx.recv() => {}
// an item is requesting an update, update the bar state
Some((i3_item, idx)) = item_rx.recv() => {
let mut i3_item = i3_item
// the name of the item
.name(item_names[idx].clone())
// always override the bar item's `instance`, since we track that ourselves
.instance(idx.to_string());
if let Some(separator) = config.items[idx].common.separator {
i3_item = i3_item.separator(separator);
}
// don't bother doing anything if the item hasn't changed
if bar[idx] == i3_item {
continue;
}
// update item in bar
bar[idx] = i3_item;
}
}
// style urgent colours differently based on the urgent_timer's status
let mut theme = config.theme.clone();
if urgent_timer.swapped() {
theme.urgent_bg = config.theme.urgent_fg;
theme.urgent_fg = config.theme.urgent_bg;
}
// print bar to STDOUT for i3
match bar.to_json(&theme) {
// make sure to include the trailing comma `,` as part of the protocol
Ok(json) => println!("{},", json),
// on any serialisation error, emit an error that will be drawn to the status bar
Err(e) => {
log::error!("failed to serialise bar to json: {}", e);
println!(
r#"[{{"full_text":"FATAL ERROR: see logs in stderr","color":"black","background":"red"}}],"#
);
}
}
}
});
Ok(())
}