lib/tldr/src/formats/elf/archive/prefetch.zig

daab053ee43316e1809a84551d573ddd1e5bf3d2

  1 const std = @import("std");
  2 const allocators = @import("alloc");
  3 const root = @import("../../../root.zig");
  4 const elf = @import("../root.zig");
  5 
  6 const Allocator = std.mem.Allocator;
  7 const archive_format = root.archive;
  8 const elf_object = elf.object;
  9 const model = root.model;
 10 const parallel = root.parallel;
 11 const trace = root.trace;
 12 const ObjectSelectionSummary = elf.parser.ObjectSelectionSummary;
 13 const parseObjectSelectionSummary = elf.parser.parseObjectSelectionSummary;
 14 
 15 pub const speculative_summary_min_bytes = 2 * 1024 * 1024;
 16 pub const summary_bytes_per_worker = 1024 * 1024;
 17 
 18 pub const Prefetched = struct {
 19     parsed: archive_format.ParsedArchive,
 20     summaries: []const ?ObjectSelectionSummary = &.{},
 21     failures: []const ?model.Error = &.{},
 22 };
 23 
 24 pub const Prefetcher = struct {
 25     enabled: bool = false,
 26     selection: Allocator = undefined,
 27     inputs: []const model.Input = &.{},
 28     max_link_jobs: usize = 0,
 29     next_scan: usize = 0,
 30     target: ?usize = null,
 31     parsed: ?archive_format.ParsedArchive = null,
 32     failure: ?model.Error = null,
 33     references: []usize = &.{},
 34     summaries: []?ObjectSelectionSummary = &.{},
 35     failures: []?model.Error = &.{},
 36     table_job: TableJob = .{},
 37     summary_job: SummaryJob = .{},
 38 
 39     const TableJob: type = parallel.Background(*Prefetcher, parseTable);
 40     const SummaryJob: type = parallel.Background(*Prefetcher, parseSummary);
 41 
 42     pub fn init(
 43         selection: Allocator,
 44         inputs: []const model.Input,
 45         options: model.LinkOptions,
 46     ) Prefetcher {
 47         return .{
 48             .enabled = options.max_link_jobs != 1 and parallel.overlapAvailable(),
 49             .selection = selection,
 50             .inputs = inputs,
 51             .max_link_jobs = options.max_link_jobs,
 52         };
 53     }
 54 
 55     pub fn deinit(self: *Prefetcher) void {
 56         self.table_job.wait();
 57         self.summary_job.wait();
 58     }
 59 
 60     pub fn advance(self: *Prefetcher, from_index: usize) void {
 61         if (!self.enabled) return;
 62         if (self.target) |current| {
 63             if (current >= from_index) return;
 64             self.table_job.wait();
 65             self.summary_job.wait();
 66             self.target = null;
 67         }
 68         var index = @max(self.next_scan, from_index);
 69         while (index < self.inputs.len and !archive_format.isArchive(self.inputs[index].bytes)) {
 70             index += 1;
 71         }
 72         if (index >= self.inputs.len) {
 73             self.next_scan = self.inputs.len;
 74             return;
 75         }
 76         self.next_scan = index + 1;
 77         self.target = index;
 78         self.parsed = null;
 79         self.failure = null;
 80         self.references = &.{};
 81         self.summaries = &.{};
 82         self.failures = &.{};
 83         self.table_job.submit(1, 1, self);
 84     }
 85 
 86     pub fn take(self: *Prefetcher, input_index: usize) model.Error!?Prefetched {
 87         const current = self.target orelse return null;
 88         if (current != input_index) return null;
 89         self.table_job.wait();
 90         self.summary_job.wait();
 91         self.target = null;
 92         if (self.failure) |err| {
 93             self.failure = null;
 94             return err;
 95         }
 96         const parsed = self.parsed.?;
 97         self.parsed = null;
 98         return .{
 99             .parsed = parsed,
100             .summaries = self.summaries,
101             .failures = self.failures,
102         };
103     }
104 
105     fn parseTable(self: *Prefetcher, item: usize) void {
106         _ = item;
107         const input = self.inputs[self.target.?];
108         self.parsed = archive_format.parseDetailed(self.selection, input) catch |err| {
109             self.failure = err;
110             return;
111         };
112         const parsed = &self.parsed.?;
113         if (!parsed.has_symbol_index) return;
114         self.chainSummaries(parsed) catch {};
115     }
116 
117     fn chainSummaries(self: *Prefetcher, parsed: *const archive_format.ParsedArchive) Allocator.Error!void {
118         const members = parsed.members;
119         const referenced = try self.selection.alloc(bool, members.len);
120         @memset(referenced, false);
121         var references = std.ArrayListUnmanaged(usize).empty;
122         var total_bytes: usize = 0;
123 
124         var offsets_ordered = true;
125         var previous_member_offset: u64 = 0;
126         var member_cursor: usize = 0;
127         for (parsed.symbol_index) |entry| {
128             if (entry.member_offset < previous_member_offset) offsets_ordered = false;
129             previous_member_offset = entry.member_offset;
130             const member_index = if (offsets_ordered)
131                 archive_format.memberIndexByOffsetWithCursor(members, entry.member_offset, &member_cursor) orelse return
132             else
133                 archive_format.memberIndexByOffset(members, entry.member_offset) orelse return;
134             if (referenced[member_index]) continue;
135             referenced[member_index] = true;
136             try references.append(self.selection, member_index);
137             total_bytes +|= members[member_index].bytes.len;
138         }
139         if (total_bytes < speculative_summary_min_bytes) return;
140 
141         const summaries = try self.selection.alloc(?ObjectSelectionSummary, members.len);
142         @memset(summaries, null);
143         const failures = try self.selection.alloc(?model.Error, members.len);
144         @memset(failures, null);
145         self.summaries = summaries;
146         self.failures = failures;
147         self.references = references.items;
148 
149         const requested_workers = if (self.max_link_jobs != 0)
150             self.max_link_jobs
151         else
152             total_bytes / summary_bytes_per_worker;
153         self.summary_job.submit(self.references.len, requested_workers, self);
154     }
155 
156     fn parseSummary(self: *Prefetcher, item: usize) void {
157         const member_index = self.references[item];
158         const member = self.parsed.?.members[member_index];
159         self.summaries[member_index] = parseObjectSelectionSummary(self.selection, .{
160             .name = member.name,
161             .bytes = member.bytes,
162         }) catch |err| {
163             self.failures[member_index] = err;
164             return;
165         };
166     }
167 };
168 
169 pub fn prepareSummaries(
170     selection: Allocator,
171     candidates: anytype,
172     options: model.LinkOptions,
173 ) void {
174     if (options.max_link_jobs == 1) return;
175     var total_bytes: usize = 0;
176     for (candidates) |candidate| {
177         if (candidate.summary != null or candidate.object != null) continue;
178         total_bytes +|= candidate.member.bytes.len;
179     }
180     if (total_bytes < speculative_summary_min_bytes) return;
181 
182     const requested_workers = if (options.max_link_jobs != 0)
183         options.max_link_jobs
184     else
185         total_bytes / summary_bytes_per_worker;
186     const workers = parallel.chooseWorkers(candidates.len, requested_workers);
187     if (workers <= 1) return;
188 
189     const summary_phase = trace.product(.archive_summary_parse);
190     defer summary_phase.end();
191     const Context = SummaryContext(@TypeOf(candidates));
192     var context = Context{
193         .selection = selection,
194         .candidates = candidates,
195     };
196     parallel.forItems(candidates.len, workers, &context, Context.task);
197 }
198 
199 fn SummaryContext(comptime Candidates: type) type {
200     return struct {
201         selection: Allocator,
202         candidates: Candidates,
203 
204         fn task(context: *@This(), worker: usize, index: usize) void {
205             _ = worker;
206             const candidate = &context.candidates[index];
207             if (candidate.summary != null or candidate.object != null) return;
208             candidate.summary = parseObjectSelectionSummary(context.selection, .{
209                 .name = candidate.member.name,
210                 .bytes = candidate.member.bytes,
211             }) catch |err| {
212                 candidate.summary_failure = err;
213                 return;
214             };
215         }
216     };
217 }
218 
219 test "prefetcher stashes the next archive table and summaries" {
220     const allocator = std.testing.allocator;
221     var selection_state = std.heap.ArenaAllocator.init(allocator);
222     defer selection_state.deinit();
223     var selection_lock = allocators.LockedAllocator.init(selection_state.allocator());
224     const selection = selection_lock.allocator();
225 
226     const payload = try allocator.alloc(u8, speculative_summary_min_bytes);
227     defer allocator.free(payload);
228     @memset(payload, 0x90);
229     const member_object = try elf_object.build(allocator, .{
230         .sections = &.{
231             elf_object.Section.progbits(".text.big", payload, std.elf.SHF_EXECINSTR, 16),
232         },
233         .symbols = &.{
234             elf_object.Symbol.section(1),
235             elf_object.Symbol.function("big", 1, 0, payload.len),
236         },
237     });
238     defer allocator.free(member_object);
239 
240     const archive_bytes = try archive_format.build(allocator, &.{
241         .{ .name = "big.o", .bytes = member_object, .symbols = &.{"big"} },
242     });
243     defer allocator.free(archive_bytes);
244 
245     const inputs = [_]model.Input{
246         .{ .name = "start.o", .bytes = "not an archive" },
247         .{ .name = "libbig.a", .bytes = archive_bytes },
248     };
249 
250     var prefetcher = Prefetcher.init(selection, &inputs, .{});
251     defer prefetcher.deinit();
252     prefetcher.advance(0);
253 
254     try std.testing.expectEqual(@as(?Prefetched, null), try prefetcher.take(0));
255     const ready = (try prefetcher.take(1)) orelse {
256         if (!parallel.overlapAvailable()) return;
257         return error.MissingPrefetch;
258     };
259     try std.testing.expect(ready.parsed.has_symbol_index);
260     try std.testing.expectEqual(@as(usize, 1), ready.parsed.members.len);
261     try std.testing.expectEqualStrings("big.o", ready.parsed.members[0].name);
262     try std.testing.expectEqual(@as(usize, 1), ready.summaries.len);
263     const summary = ready.summaries[0] orelse return error.MissingSummary;
264     try std.testing.expect(summary.symbols.len != 0);
265     try std.testing.expectEqual(@as(?model.Error, null), ready.failures[0]);
266 }
267 
268 test "prefetcher surfaces stashed table failures at take" {
269     const allocator = std.testing.allocator;
270     var selection_state = std.heap.ArenaAllocator.init(allocator);
271     defer selection_state.deinit();
272     var selection_lock = allocators.LockedAllocator.init(selection_state.allocator());
273     const selection = selection_lock.allocator();
274 
275     const corrupt = "!<arch>\ncorrupt archive body";
276     const inputs = [_]model.Input{
277         .{ .name = "libbad.a", .bytes = corrupt },
278     };
279 
280     var prefetcher = Prefetcher.init(selection, &inputs, .{});
281     defer prefetcher.deinit();
282     prefetcher.advance(0);
283     if (!prefetcher.enabled) return;
284 
285     try std.testing.expectError(error.InvalidArchive, prefetcher.take(0));
286     try std.testing.expectEqual(@as(?Prefetched, null), try prefetcher.take(0));
287 }
288 
289 test "prefetcher discards stale targets when the loop passes them" {
290     const allocator = std.testing.allocator;
291     var selection_state = std.heap.ArenaAllocator.init(allocator);
292     defer selection_state.deinit();
293     var selection_lock = allocators.LockedAllocator.init(selection_state.allocator());
294     const selection = selection_lock.allocator();
295 
296     const first = try archive_format.build(allocator, &.{
297         .{ .name = "a.o", .bytes = "xy", .symbols = &.{"a"} },
298     });
299     defer allocator.free(first);
300     const second = try archive_format.build(allocator, &.{
301         .{ .name = "b.o", .bytes = "zw", .symbols = &.{"b"} },
302     });
303     defer allocator.free(second);
304 
305     const inputs = [_]model.Input{
306         .{ .name = "libfirst.a", .bytes = first },
307         .{ .name = "libsecond.a", .bytes = second },
308     };
309 
310     var prefetcher = Prefetcher.init(selection, &inputs, .{});
311     defer prefetcher.deinit();
312     prefetcher.advance(0);
313     if (!prefetcher.enabled) return;
314 
315     prefetcher.advance(1);
316     const ready = (try prefetcher.take(1)) orelse return error.MissingPrefetch;
317     try std.testing.expectEqualStrings("b.o", ready.parsed.members[0].name);
318 }
319 
320 test "disabled prefetcher never targets inputs" {
321     const allocator = std.testing.allocator;
322     var selection_state = std.heap.ArenaAllocator.init(allocator);
323     defer selection_state.deinit();
324     var selection_lock = allocators.LockedAllocator.init(selection_state.allocator());
325     const selection = selection_lock.allocator();
326 
327     const bytes = try archive_format.build(allocator, &.{
328         .{ .name = "a.o", .bytes = "xy", .symbols = &.{"a"} },
329     });
330     defer allocator.free(bytes);
331     const inputs = [_]model.Input{
332         .{ .name = "liba.a", .bytes = bytes },
333     };
334 
335     var prefetcher = Prefetcher.init(selection, &inputs, .{ .max_link_jobs = 1 });
336     defer prefetcher.deinit();
337     prefetcher.advance(0);
338     try std.testing.expectEqual(@as(?Prefetched, null), try prefetcher.take(0));
339 }