rework std.Progress again

This time, we preallocate a fixed set of nodes and have the user-visible
Node only be an index into them. This allows for lock-free management of
the node storage.

Only the parent indexes are stored, and the update thread makes a
serialized copy of the state before trying to compute children lists.

The update thread then walks the tree and outputs an entire tree of
progress rather than only one line.

There is a problem with clearing from the cursor to the end of the
screen when the cursor is at the bottom of the terminal.
This commit is contained in:
Andrew Kelley 2024-05-22 20:42:12 -07:00
parent e1e4de2776
commit a3c9511ab9

View File

@ -22,14 +22,10 @@ is_windows_terminal: bool,
/// Whether the terminal supports ANSI escape codes. /// Whether the terminal supports ANSI escape codes.
supports_ansi_escape_codes: bool, supports_ansi_escape_codes: bool,
root: Node,
update_thread: ?std.Thread, update_thread: ?std.Thread,
/// Atomically set by SIGWINCH as well as the root done() function. /// Atomically set by SIGWINCH as well as the root done() function.
redraw_event: std.Thread.ResetEvent, redraw_event: std.Thread.ResetEvent,
/// Ensure there is only 1 global Progress object.
initialized: bool,
/// Indicates a request to shut down and reset global state. /// Indicates a request to shut down and reset global state.
/// Accessed atomically. /// Accessed atomically.
done: bool, done: bool,
@ -43,13 +39,22 @@ cols: u16,
/// Accessed only by the update thread. /// Accessed only by the update thread.
draw_buffer: []u8, draw_buffer: []u8,
/// This is in a separate array from `node_storage` but with the same length so
/// that it can be iterated over efficiently without trashing too much of the
/// CPU cache.
node_parents: []Node.Parent,
node_storage: []Node.Storage,
node_freelist: []Node.OptionalIndex,
node_freelist_first: Node.OptionalIndex,
node_end_index: u32,
pub const Options = struct { pub const Options = struct {
/// User-provided buffer with static lifetime. /// User-provided buffer with static lifetime.
/// ///
/// Used to store the entire write buffer sent to the terminal. Progress output will be truncated if it /// Used to store the entire write buffer sent to the terminal. Progress output will be truncated if it
/// cannot fit into this buffer which will look bad but not cause any malfunctions. /// cannot fit into this buffer which will look bad but not cause any malfunctions.
/// ///
/// Must be at least 100 bytes. /// Must be at least 200 bytes.
draw_buffer: []u8, draw_buffer: []u8,
/// How many nanoseconds between writing updates to the terminal. /// How many nanoseconds between writing updates to the terminal.
refresh_rate_ns: u64 = 50 * std.time.ns_per_ms, refresh_rate_ns: u64 = 50 * std.time.ns_per_ms,
@ -64,66 +69,128 @@ pub const Options = struct {
/// Represents one unit of progress. Each node can have children nodes, or /// Represents one unit of progress. Each node can have children nodes, or
/// one can use integers with `update`. /// one can use integers with `update`.
pub const Node = struct { pub const Node = struct {
mutex: std.Thread.Mutex, index: OptionalIndex,
/// Links to the parent and child nodes.
parent_list_node: std.DoublyLinkedList(void).Node,
/// Links to the prev and next sibling nodes.
sibling_list_node: std.DoublyLinkedList(void).Node,
name: []const u8, pub const max_name_len = 38;
/// Must be handled atomically to be thread-safe. 0 means null.
unprotected_estimated_total_items: usize,
/// Must be handled atomically to be thread-safe.
unprotected_completed_items: usize,
pub const ListNode = std.DoublyLinkedList(void); const Storage = extern struct {
/// Little endian.
completed_count: u32,
/// 0 means unknown.
/// Little endian.
estimated_total_count: u32,
name: [max_name_len]u8,
};
const Parent = enum(u16) {
/// Unallocated storage.
unused = std.math.maxInt(u16) - 1,
/// Indicates root node.
none = std.math.maxInt(u16),
/// Index into `node_storage`.
_,
fn unwrap(i: @This()) ?Index {
return switch (i) {
.unused, .none => return null,
else => @enumFromInt(@intFromEnum(i)),
};
}
};
const OptionalIndex = enum(u16) {
none = std.math.maxInt(u16),
/// Index into `node_storage`.
_,
fn unwrap(i: @This()) ?Index {
if (i == .none) return null;
return @enumFromInt(@intFromEnum(i));
}
fn toParent(i: @This()) Parent {
assert(@intFromEnum(i) != @intFromEnum(Parent.unused));
return @enumFromInt(@intFromEnum(i));
}
};
/// Index into `node_storage`.
const Index = enum(u16) {
_,
fn toParent(i: @This()) Parent {
assert(@intFromEnum(i) != @intFromEnum(Parent.unused));
assert(@intFromEnum(i) != @intFromEnum(Parent.none));
return @enumFromInt(@intFromEnum(i));
}
fn toOptional(i: @This()) OptionalIndex {
return @enumFromInt(@intFromEnum(i));
}
};
/// Create a new child progress node. Thread-safe. /// Create a new child progress node. Thread-safe.
/// ///
/// It is expected for the memory of the result to be stored in the
/// caller's stack and therefore is required to call `activate` immediately
/// on the result after initializing the memory location and `end` when done.
///
/// Passing 0 for `estimated_total_items` means unknown. /// Passing 0 for `estimated_total_items` means unknown.
pub fn start(self: *Node, name: []const u8, estimated_total_items: usize) Node { pub fn start(node: Node, name: []const u8, estimated_total_items: usize) Node {
return .{ const node_index = node.index.unwrap() orelse return .{ .index = .none };
.mutex = .{}, const parent = node_index.toParent();
.parent_list_node = .{
.prev = &self.parent_list_node, const freelist_head = &global_progress.node_freelist_first;
.next = null, var opt_free_index = @atomicLoad(Node.OptionalIndex, freelist_head, .seq_cst);
.data = {}, while (opt_free_index.unwrap()) |free_index| {
}, const freelist_ptr = freelistByIndex(free_index);
.sibling_list_node = .{ .data = {} }, opt_free_index = @cmpxchgWeak(Node.OptionalIndex, freelist_head, opt_free_index, freelist_ptr.*, .seq_cst, .seq_cst) orelse {
.name = name, // We won the allocation race.
.unprotected_estimated_total_items = estimated_total_items, return init(free_index, parent, name, estimated_total_items);
.unprotected_completed_items = 0,
}; };
} }
/// To be called exactly once after `start`. const free_index = @atomicRmw(u32, &global_progress.node_end_index, .Add, 1, .monotonic);
pub fn activate(n: *Node) void { if (free_index >= global_progress.node_storage.len) {
const p = n.parent().?; // Ran out of node storage memory. Progress for this node will not be tracked.
p.mutex.lock(); _ = @atomicRmw(u32, &global_progress.node_end_index, .Sub, 1, .monotonic);
defer p.mutex.unlock(); return .{ .index = .none };
assert(p.parent_list_node.next == null); }
p.parent_list_node.next = &n.parent_list_node;
return init(@enumFromInt(free_index), parent, name, estimated_total_items);
} }
/// This is the same as calling `start` and then `end` on the returned `Node`. Thread-safe. /// This is the same as calling `start` and then `end` on the returned `Node`. Thread-safe.
pub fn completeOne(self: *Node) void { pub fn completeOne(n: Node) void {
_ = @atomicRmw(usize, &self.unprotected_completed_items, .Add, 1, .monotonic); const index = n.index.unwrap() orelse return;
const storage = storageByIndex(index);
_ = @atomicRmw(u32, &storage.completed_count, .Add, 1, .monotonic);
}
/// Thread-safe.
pub fn setCompletedItems(n: Node, completed_items: usize) void {
const index = n.index.unwrap() orelse return;
const storage = storageByIndex(index);
@atomicStore(u32, &storage.completed_count, std.math.lossyCast(u32, completed_items), .monotonic);
}
/// Thread-safe. 0 means unknown.
pub fn setEstimatedTotalItems(n: Node, count: usize) void {
const index = n.index.unwrap() orelse return;
const storage = storageByIndex(index);
@atomicStore(u32, &storage.estimated_total_count, std.math.lossyCast(u32, count), .monotonic);
} }
/// Finish a started `Node`. Thread-safe. /// Finish a started `Node`. Thread-safe.
pub fn end(child: *Node) void { pub fn end(n: Node) void {
if (child.parent()) |p| { const index = n.index.unwrap() orelse return;
// Make sure the other thread doesn't access this memory that is const parent_ptr = parentByIndex(index);
// about to be released. if (parent_ptr.unwrap()) |parent_index| {
child.mutex.lock(); _ = @atomicRmw(u32, &storageByIndex(parent_index).completed_count, .Add, 1, .monotonic);
@atomicStore(Node.Parent, parent_ptr, .unused, .seq_cst);
const other = if (child.sibling_list_node.next) |n| n else child.sibling_list_node.prev; const freelist_head = &global_progress.node_freelist_first;
_ = @cmpxchgStrong(std.DoublyLinkedList(void).Node, &p.parent_list_node.next, child, other, .seq_cst, .seq_cst); var first = @atomicLoad(Node.OptionalIndex, freelist_head, .seq_cst);
p.completeOne(); while (true) {
freelistByIndex(index).* = first;
first = @cmpxchgWeak(Node.OptionalIndex, freelist_head, first, index.toOptional(), .seq_cst, .seq_cst) orelse break;
}
} else { } else {
@atomicStore(bool, &global_progress.done, true, .seq_cst); @atomicStore(bool, &global_progress.done, true, .seq_cst);
global_progress.redraw_event.set(); global_progress.redraw_event.set();
@ -131,19 +198,35 @@ pub const Node = struct {
} }
} }
/// Thread-safe. 0 means unknown. fn storageByIndex(index: Node.Index) *Node.Storage {
pub fn setEstimatedTotalItems(self: *Node, count: usize) void { return &global_progress.node_storage[@intFromEnum(index)];
@atomicStore(usize, &self.unprotected_estimated_total_items, count, .monotonic);
} }
/// Thread-safe. fn parentByIndex(index: Node.Index) *Node.Parent {
pub fn setCompletedItems(self: *Node, completed_items: usize) void { return &global_progress.node_parents[@intFromEnum(index)];
@atomicStore(usize, &self.unprotected_completed_items, completed_items, .monotonic);
} }
fn parent(child: *Node) ?*Node { fn freelistByIndex(index: Node.Index) *Node.OptionalIndex {
const parent_node = child.parent_list_node.prev orelse return null; return &global_progress.node_freelist[@intFromEnum(index)];
return @fieldParentPtr("parent_list_node", parent_node); }
fn init(free_index: Index, parent: Parent, name: []const u8, estimated_total_items: usize) Node {
assert(parent != .unused);
const storage = storageByIndex(free_index);
storage.* = .{
.completed_count = 0,
.estimated_total_count = std.math.lossyCast(u32, estimated_total_items),
.name = [1]u8{0} ** max_name_len,
};
const name_len = @min(max_name_len, name.len);
@memcpy(storage.name[0..name_len], name[0..name_len]);
const parent_ptr = parentByIndex(free_index);
assert(parent_ptr.* == .unused);
@atomicStore(Node.Parent, parent_ptr, parent, .release);
return .{ .index = free_index.toOptional() };
} }
}; };
@ -151,25 +234,36 @@ var global_progress: Progress = .{
.terminal = null, .terminal = null,
.is_windows_terminal = false, .is_windows_terminal = false,
.supports_ansi_escape_codes = false, .supports_ansi_escape_codes = false,
.root = undefined,
.update_thread = null, .update_thread = null,
.redraw_event = .{}, .redraw_event = .{},
.initialized = false,
.refresh_rate_ns = undefined, .refresh_rate_ns = undefined,
.initial_delay_ns = undefined, .initial_delay_ns = undefined,
.rows = 0, .rows = 0,
.cols = 0, .cols = 0,
.draw_buffer = undefined, .draw_buffer = undefined,
.done = false, .done = false,
// TODO: make these configurable and avoid including the globals in .data if unused
.node_parents = &node_parents_buffer,
.node_storage = &node_storage_buffer,
.node_freelist = &node_freelist_buffer,
.node_freelist_first = .none,
.node_end_index = 0,
}; };
const default_node_storage_buffer_len = 100;
var node_parents_buffer: [default_node_storage_buffer_len]Node.Parent = undefined;
var node_storage_buffer: [default_node_storage_buffer_len]Node.Storage = undefined;
var node_freelist_buffer: [default_node_storage_buffer_len]Node.OptionalIndex = undefined;
/// Initializes a global Progress instance. /// Initializes a global Progress instance.
/// ///
/// Asserts there is only one global Progress instance. /// Asserts there is only one global Progress instance.
/// ///
/// Call `Node.end` when done. /// Call `Node.end` when done.
pub fn start(options: Options) *Node { pub fn start(options: Options) Node {
assert(!global_progress.initialized); // Ensure there is only 1 global Progress object.
assert(global_progress.node_end_index == 0);
const stderr = std.io.getStdErr(); const stderr = std.io.getStdErr();
if (stderr.supportsAnsiEscapeCodes()) { if (stderr.supportsAnsiEscapeCodes()) {
global_progress.terminal = stderr; global_progress.terminal = stderr;
@ -181,18 +275,12 @@ pub fn start(options: Options) *Node {
// we are in a "dumb" terminal like in acme or writing to a file // we are in a "dumb" terminal like in acme or writing to a file
global_progress.terminal = stderr; global_progress.terminal = stderr;
} }
global_progress.root = .{ @memset(global_progress.node_parents, .unused);
.mutex = .{}, const root_node = Node.init(@enumFromInt(0), .none, options.root_name, options.estimated_total_items);
.parent_list_node = .{ .data = {} },
.sibling_list_node = .{ .data = {} },
.name = options.root_name,
.unprotected_estimated_total_items = options.estimated_total_items,
.unprotected_completed_items = 0,
};
global_progress.done = false; global_progress.done = false;
global_progress.initialized = true; global_progress.node_end_index = 1;
assert(options.draw_buffer.len >= 100); assert(options.draw_buffer.len >= 200);
global_progress.draw_buffer = options.draw_buffer; global_progress.draw_buffer = options.draw_buffer;
global_progress.refresh_rate_ns = options.refresh_rate_ns; global_progress.refresh_rate_ns = options.refresh_rate_ns;
global_progress.initial_delay_ns = options.initial_delay_ns; global_progress.initial_delay_ns = options.initial_delay_ns;
@ -204,7 +292,7 @@ pub fn start(options: Options) *Node {
}; };
posix.sigaction(posix.SIG.WINCH, &act, null) catch { posix.sigaction(posix.SIG.WINCH, &act, null) catch {
global_progress.terminal = null; global_progress.terminal = null;
return &global_progress.root; return root_node;
}; };
if (global_progress.terminal != null) { if (global_progress.terminal != null) {
@ -215,7 +303,7 @@ pub fn start(options: Options) *Node {
} }
} }
return &global_progress.root; return root_node;
} }
/// Returns whether a resize is needed to learn the terminal size. /// Returns whether a resize is needed to learn the terminal size.
@ -263,11 +351,85 @@ const save = "\x1b7";
const restore = "\x1b8"; const restore = "\x1b8";
const finish_sync = "\x1b[?2026l"; const finish_sync = "\x1b[?2026l";
const tree_tee = "\x1B\x28\x30\x74\x71\x1B\x28\x42 "; //
const tree_line = "\x1B\x28\x30\x78\x1B\x28\x42 "; //
const tree_langle = "\x1B\x28\x30\x6d\x71\x1B\x28\x42 "; //
fn clearTerminal() void { fn clearTerminal() void {
write(clear); write(clear);
} }
const Children = struct {
child: Node.OptionalIndex,
sibling: Node.OptionalIndex,
};
fn computeRedraw() []u8 { fn computeRedraw() []u8 {
// TODO make this configurable
var serialized_node_parents_buffer: [default_node_storage_buffer_len]Node.Parent = undefined;
var serialized_node_storage_buffer: [default_node_storage_buffer_len]Node.Storage = undefined;
var serialized_node_map_buffer: [default_node_storage_buffer_len]Node.Index = undefined;
var serialized_len: usize = 0;
// Iterate all of the nodes and construct a serializable copy of the state that can be examined
// without atomics.
const end_index = @atomicLoad(u32, &global_progress.node_end_index, .monotonic);
const node_parents = global_progress.node_parents[0..end_index];
const node_storage = global_progress.node_storage[0..end_index];
for (node_parents, node_storage, 0..) |*parent_ptr, *storage_ptr, i| {
var begin_parent = @atomicLoad(Node.Parent, parent_ptr, .seq_cst);
while (begin_parent != .unused) {
const dest_storage = &serialized_node_storage_buffer[serialized_len];
@memcpy(&dest_storage.name, &storage_ptr.name);
dest_storage.completed_count = @atomicLoad(u32, &storage_ptr.completed_count, .monotonic);
dest_storage.estimated_total_count = @atomicLoad(u32, &storage_ptr.estimated_total_count, .monotonic);
const end_parent = @atomicLoad(Node.Parent, parent_ptr, .seq_cst);
if (begin_parent == end_parent) {
serialized_node_parents_buffer[serialized_len] = begin_parent;
serialized_node_map_buffer[i] = @enumFromInt(serialized_len);
serialized_len += 1;
break;
}
begin_parent = end_parent;
}
}
// Now we can analyze our copy of the graph without atomics, reconstructing
// children lists which do not exist in the canonical data. These are
// needed for tree traversal below.
const serialized_node_parents = serialized_node_parents_buffer[0..serialized_len];
const serialized_node_storage = serialized_node_storage_buffer[0..serialized_len];
// Remap parents to point inside serialized arrays.
for (serialized_node_parents) |*parent| {
parent.* = switch (parent.*) {
.unused => unreachable,
.none => .none,
_ => |p| serialized_node_map_buffer[@intFromEnum(p)].toParent(),
};
}
var children_buffer: [default_node_storage_buffer_len]Children = undefined;
const children = children_buffer[0..serialized_len];
@memset(children, .{ .child = .none, .sibling = .none });
for (serialized_node_parents, 0..) |parent, child_index_usize| {
const child_index: Node.Index = @enumFromInt(child_index_usize);
assert(parent != .unused);
const parent_index = parent.unwrap() orelse continue;
const children_node = &children[@intFromEnum(parent_index)];
if (children_node.child.unwrap()) |existing_child_index| {
const existing_child = &children[@intFromEnum(existing_child_index)];
existing_child.sibling = child_index.toOptional();
children[@intFromEnum(child_index)].sibling = existing_child.sibling;
} else {
children_node.child = child_index.toOptional();
}
}
// The strategy is: keep the cursor at the beginning, and then with every redraw: // The strategy is: keep the cursor at the beginning, and then with every redraw:
// erase, save, write, restore // erase, save, write, restore
@ -280,27 +442,8 @@ fn computeRedraw() []u8 {
buf[0..prefix.len].* = prefix.*; buf[0..prefix.len].* = prefix.*;
i = prefix.len; i = prefix.len;
// Walk the tree and write the progress output to the buffer. const root_node_index: Node.Index = @enumFromInt(0);
var node: *Node = &global_progress.root; i = computeNode(buf, i, serialized_node_storage, serialized_node_parents, children, root_node_index);
while (true) {
const eti = @atomicLoad(usize, &node.unprotected_estimated_total_items, .monotonic);
const completed_items = @atomicLoad(usize, &node.unprotected_completed_items, .monotonic);
if (node.name.len != 0 or eti > 0) {
if (node.name.len != 0) {
i += (std.fmt.bufPrint(buf[i..], "{s}", .{node.name}) catch @panic("TODO")).len;
}
if (eti > 0) {
i += (std.fmt.bufPrint(buf[i..], "[{d}/{d}] ", .{ completed_items, eti }) catch @panic("TODO")).len;
} else if (completed_items != 0) {
i += (std.fmt.bufPrint(buf[i..], "[{d}] ", .{completed_items}) catch @panic("TODO")).len;
}
}
node = @atomicLoad(?*Node, &node.recently_updated_child, .acquire) orelse break;
}
i = @min(global_progress.cols + prefix.len, i);
buf[i..][0..suffix.len].* = suffix.*; buf[i..][0..suffix.len].* = suffix.*;
i += suffix.len; i += suffix.len;
@ -308,6 +451,84 @@ fn computeRedraw() []u8 {
return buf[0..i]; return buf[0..i];
} }
fn computePrefix(
buf: []u8,
start_i: usize,
serialized_node_storage: []const Node.Storage,
serialized_node_parents: []const Node.Parent,
children: []const Children,
node_index: Node.Index,
) usize {
var i = start_i;
const parent_index = serialized_node_parents[@intFromEnum(node_index)].unwrap() orelse return i;
if (serialized_node_parents[@intFromEnum(parent_index)] == .none) return i;
i = computePrefix(buf, i, serialized_node_storage, serialized_node_parents, children, parent_index);
if (children[@intFromEnum(parent_index)].sibling == .none) {
buf[i..][0..3].* = " ".*;
i += 3;
} else {
buf[i..][0..tree_line.len].* = tree_line.*;
i += tree_line.len;
}
return i;
}
fn computeNode(
buf: []u8,
start_i: usize,
serialized_node_storage: []const Node.Storage,
serialized_node_parents: []const Node.Parent,
children: []const Children,
node_index: Node.Index,
) usize {
var i = start_i;
i = computePrefix(buf, i, serialized_node_storage, serialized_node_parents, children, node_index);
const storage = &serialized_node_storage[@intFromEnum(node_index)];
const estimated_total = storage.estimated_total_count;
const completed_items = storage.completed_count;
const name = if (std.mem.indexOfScalar(u8, &storage.name, 0)) |end| storage.name[0..end] else &storage.name;
const parent = serialized_node_parents[@intFromEnum(node_index)];
if (parent != .none) {
if (children[@intFromEnum(node_index)].sibling == .none) {
buf[i..][0..tree_langle.len].* = tree_langle.*;
i += tree_langle.len;
} else {
buf[i..][0..tree_tee.len].* = tree_tee.*;
i += tree_tee.len;
}
}
if (name.len != 0 or estimated_total > 0) {
if (estimated_total > 0) {
i += (std.fmt.bufPrint(buf[i..], "[{d}/{d}] ", .{ completed_items, estimated_total }) catch &.{}).len;
} else if (completed_items != 0) {
i += (std.fmt.bufPrint(buf[i..], "[{d}] ", .{completed_items}) catch &.{}).len;
}
if (name.len != 0) {
i += (std.fmt.bufPrint(buf[i..], "{s}", .{name}) catch &.{}).len;
}
}
i = @min(global_progress.cols + start_i, i);
buf[i] = '\n';
i += 1;
if (children[@intFromEnum(node_index)].child.unwrap()) |child| {
i = computeNode(buf, i, serialized_node_storage, serialized_node_parents, children, child);
}
{
var opt_sibling = children[@intFromEnum(node_index)].sibling;
while (opt_sibling.unwrap()) |sibling| {
i = computeNode(buf, i, serialized_node_storage, serialized_node_parents, children, sibling);
}
}
return i;
}
fn write(buf: []const u8) void { fn write(buf: []const u8) void {
const tty = global_progress.terminal orelse return; const tty = global_progress.terminal orelse return;
tty.writeAll(buf) catch { tty.writeAll(buf) catch {