std.fs.File.Writer: fix drain implementation

it didn't account for data.len can no longer be zero
This commit is contained in:
Andrew Kelley 2025-06-30 12:17:47 -07:00
parent 77e839e283
commit 435ccf706d

View File

@ -1648,29 +1648,28 @@ pub const Writer = struct {
}; };
} }
pub fn drain(io_writer: *std.io.Writer, data: []const []const u8, splat: usize) std.io.Writer.Error!usize { pub fn drain(io_w: *std.io.Writer, data: []const []const u8, splat: usize) std.io.Writer.Error!usize {
const w: *Writer = @fieldParentPtr("interface", io_writer); const w: *Writer = @fieldParentPtr("interface", io_w);
const handle = w.file.handle; const handle = w.file.handle;
const buffered = io_writer.buffered(); const buffered = io_w.buffered();
var splat_buffer: [256]u8 = undefined;
if (is_windows) { if (is_windows) {
var i: usize = 0; var i: usize = 0;
while (i < buffered.len) { while (i < buffered.len) {
const n = windows.WriteFile(handle, buffered[i..], null) catch |err| { const n = windows.WriteFile(handle, buffered[i..], null) catch |err| {
w.err = err; w.err = err;
w.pos += i; w.pos += i;
_ = io_writer.consume(i); _ = io_w.consume(i);
return error.WriteFailed; return error.WriteFailed;
}; };
i += n; i += n;
if (data.len > 0 and buffered.len - i < n) { if (data.len > 0 and buffered.len - i < n) {
w.pos += i; w.pos += i;
return io_writer.consume(i); return io_w.consume(i);
} }
} }
if (i != 0 or data.len == 0 or (data.len == 1 and splat == 0)) { if (i != 0 or data.len == 0 or (data.len == 1 and splat == 0)) {
w.pos += i; w.pos += i;
return io_writer.consume(i); return io_w.consume(i);
} }
const n = windows.WriteFile(handle, data[0], null) catch |err| { const n = windows.WriteFile(handle, data[0], null) catch |err| {
w.err = err; w.err = err;
@ -1679,19 +1678,6 @@ pub const Writer = struct {
w.pos += n; w.pos += n;
return n; return n;
} }
if (data.len == 0) {
var i: usize = 0;
while (i < buffered.len) {
i += std.posix.write(handle, buffered) catch |err| {
w.err = err;
w.pos += i;
_ = io_writer.consume(i);
return error.WriteFailed;
};
}
w.pos += i;
return io_writer.consumeAll();
}
var iovecs: [max_buffers_len]std.posix.iovec_const = undefined; var iovecs: [max_buffers_len]std.posix.iovec_const = undefined;
var len: usize = 0; var len: usize = 0;
if (buffered.len > 0) { if (buffered.len > 0) {
@ -1700,30 +1686,43 @@ pub const Writer = struct {
} }
for (data) |d| { for (data) |d| {
if (d.len == 0) continue; if (d.len == 0) continue;
if (iovecs.len - len == 0) break;
iovecs[len] = .{ .base = d.ptr, .len = d.len }; iovecs[len] = .{ .base = d.ptr, .len = d.len };
len += 1; len += 1;
if (iovecs.len - len == 0) break;
} }
if (len == 0) return 0;
const pattern = data[data.len - 1];
switch (splat) { switch (splat) {
0 => if (data[data.len - 1].len != 0) { 0 => if (iovecs[len - 1].base == pattern.ptr) {
len -= 1; len -= 1;
}, },
1 => {}, 1 => {},
else => switch (data[data.len - 1].len) { else => switch (pattern.len) {
0 => {}, 0 => {},
1 => { 1 => memset: {
// Replace the 1-byte buffer with a bigger one.
if (iovecs[len - 1].base == pattern.ptr) len -= 1;
if (iovecs.len - len == 0) break :memset;
const splat_buffer_candidate = io_w.buffer[io_w.end..];
var backup_buffer: [64]u8 = undefined;
const splat_buffer = if (splat_buffer_candidate.len >= backup_buffer.len)
splat_buffer_candidate
else
&backup_buffer;
const memset_len = @min(splat_buffer.len, splat); const memset_len = @min(splat_buffer.len, splat);
const buf = splat_buffer[0..memset_len]; const buf = splat_buffer[0..memset_len];
@memset(buf, data[data.len - 1][0]); @memset(buf, pattern[0]);
iovecs[len - 1] = .{ .base = buf.ptr, .len = buf.len }; iovecs[len] = .{ .base = buf.ptr, .len = buf.len };
len += 1;
var remaining_splat = splat - buf.len; var remaining_splat = splat - buf.len;
while (remaining_splat > splat_buffer.len and len < iovecs.len) { while (remaining_splat > splat_buffer.len and iovecs.len - len != 0) {
iovecs[len] = .{ .base = &splat_buffer, .len = splat_buffer.len }; assert(buf.len == splat_buffer.len);
remaining_splat -= splat_buffer.len; iovecs[len] = .{ .base = splat_buffer.ptr, .len = splat_buffer.len };
len += 1; len += 1;
remaining_splat -= splat_buffer.len;
} }
if (remaining_splat > 0 and len < iovecs.len) { if (remaining_splat > 0 and iovecs.len - len != 0) {
iovecs[len] = .{ .base = &splat_buffer, .len = remaining_splat }; iovecs[len] = .{ .base = splat_buffer.ptr, .len = remaining_splat };
len += 1; len += 1;
} }
return std.posix.writev(handle, iovecs[0..len]) catch |err| { return std.posix.writev(handle, iovecs[0..len]) catch |err| {
@ -1733,7 +1732,7 @@ pub const Writer = struct {
}, },
else => for (0..splat - 1) |_| { else => for (0..splat - 1) |_| {
if (iovecs.len - len == 0) break; if (iovecs.len - len == 0) break;
iovecs[len] = .{ .base = data[data.len - 1].ptr, .len = data[data.len - 1].len }; iovecs[len] = .{ .base = pattern.ptr, .len = pattern.len };
len += 1; len += 1;
}, },
}, },
@ -1743,15 +1742,15 @@ pub const Writer = struct {
return error.WriteFailed; return error.WriteFailed;
}; };
w.pos += n; w.pos += n;
return io_writer.consume(n); return io_w.consume(n);
} }
pub fn sendFile( pub fn sendFile(
io_writer: *std.io.Writer, io_w: *std.io.Writer,
file_reader: *Reader, file_reader: *Reader,
limit: std.io.Limit, limit: std.io.Limit,
) std.io.Writer.FileError!usize { ) std.io.Writer.FileError!usize {
const w: *Writer = @fieldParentPtr("interface", io_writer); const w: *Writer = @fieldParentPtr("interface", io_w);
const out_fd = w.file.handle; const out_fd = w.file.handle;
const in_fd = file_reader.file.handle; const in_fd = file_reader.file.handle;
// TODO try using copy_file_range on FreeBSD // TODO try using copy_file_range on FreeBSD
@ -1762,7 +1761,7 @@ pub const Writer = struct {
if (w.sendfile_err != null) break :sf; if (w.sendfile_err != null) break :sf;
// Linux sendfile does not support headers. // Linux sendfile does not support headers.
const buffered = limit.slice(file_reader.interface.buffer); const buffered = limit.slice(file_reader.interface.buffer);
if (io_writer.end != 0 or buffered.len != 0) return drain(io_writer, &.{buffered}, 1); if (io_w.end != 0 or buffered.len != 0) return drain(io_w, &.{buffered}, 1);
const max_count = 0x7ffff000; // Avoid EINVAL. const max_count = 0x7ffff000; // Avoid EINVAL.
var off: std.os.linux.off_t = undefined; var off: std.os.linux.off_t = undefined;
const off_ptr: ?*std.os.linux.off_t, const count: usize = switch (file_reader.mode) { const off_ptr: ?*std.os.linux.off_t, const count: usize = switch (file_reader.mode) {
@ -1813,7 +1812,7 @@ pub const Writer = struct {
if (copy_file_range_fn) |copy_file_range| cfr: { if (copy_file_range_fn) |copy_file_range| cfr: {
if (w.copy_file_range_err != null) break :cfr; if (w.copy_file_range_err != null) break :cfr;
const buffered = limit.slice(file_reader.interface.buffer); const buffered = limit.slice(file_reader.interface.buffer);
if (io_writer.end != 0 or buffered.len != 0) return drain(io_writer, &.{buffered}, 1); if (io_w.end != 0 or buffered.len != 0) return drain(io_w, &.{buffered}, 1);
var off_in: i64 = undefined; var off_in: i64 = undefined;
var off_out: i64 = undefined; var off_out: i64 = undefined;
const off_in_ptr: ?*i64 = switch (file_reader.mode) { const off_in_ptr: ?*i64 = switch (file_reader.mode) {