tools

various tools I have been using throughout the years
Log | Files | Refs | README | LICENSE

main.zig (8206B)


      1 const std = @import("std");
      2 
      3 const default_download_url = "https://proof.ovh.net/files/100Mb.dat";
      4 const default_upload_url = "https://speed.cloudflare.com/__up";
      5 const default_upload_mb: usize = 100;
      6 
      7 const Config = struct {
      8     download_url: []const u8 = default_download_url,
      9     upload_url: []const u8 = default_upload_url,
     10     upload_bytes: usize = default_upload_mb * 1024 * 1024,
     11     run_download: bool = true,
     12     run_upload: bool = true,
     13 };
     14 
     15 const usage =
     16     \\Usage: speediness [options]
     17     \\
     18     \\Options:
     19     \\  --download-url <url>  URL to download from   (default: {s})
     20     \\  --upload-url <url>    URL to upload to       (default: {s})
     21     \\  --upload-mb <n>       Megabytes to upload    (default: {d})
     22     \\  --both                Run both tests
     23     \\  --download-only       Run the download test only
     24     \\  --upload-only         Run the upload test only
     25     \\  -h, --help            Show this help and exit
     26     \\
     27 ;
     28 
     29 fn printUsage() void {
     30     std.debug.print(usage, .{ default_download_url, default_upload_url, default_upload_mb });
     31 }
     32 
     33 fn mbps(bytes: usize, ns: u64) f64 {
     34     const secs = @as(f64, @floatFromInt(ns)) / 1e9;
     35     if (secs < 0.001) return 0;
     36     return (@as(f64, @floatFromInt(bytes)) / secs) / (1024.0 * 1024.0);
     37 }
     38 
     39 fn downloadTest(io: std.Io, allocator: std.mem.Allocator, url: []const u8) !void {
     40     var client: std.http.Client = .{
     41         .allocator = allocator,
     42         .io = io,
     43     };
     44     defer client.deinit();
     45 
     46     const uri = try std.Uri.parse(url);
     47 
     48     var req = try client.request(.GET, uri, .{});
     49     defer req.deinit();
     50 
     51     try req.sendBodiless();
     52 
     53     var redirect_buffer: [8192]u8 = undefined;
     54     var response = try req.receiveHead(&redirect_buffer);
     55 
     56     if (response.head.status != .ok) {
     57         std.debug.print("{}\n", .{response.head.status});
     58         return error.HttpError;
     59     }
     60 
     61     var transfer_buffer: [8192]u8 = undefined;
     62     var reader = response.reader(&transfer_buffer);
     63 
     64     var buf: [131072]u8 = undefined;
     65     var total: usize = 0;
     66     var iter: usize = 0;
     67     const start = std.Io.Timestamp.now(io, .awake);
     68     const start_ns: u64 = @intCast(start.nanoseconds);
     69 
     70     while (true) {
     71         const n = reader.readSliceShort(&buf) catch |err| {
     72             std.debug.print("\nRead error: {}\n", .{err});
     73             break;
     74         };
     75         if (n == 0) break;
     76         total += n;
     77         iter += 1;
     78         if (iter % 8 == 0) {
     79             const now_ns: u64 = @intCast(std.Io.Timestamp.now(io, .awake).nanoseconds);
     80             const elapsed = now_ns - start_ns;
     81             const speed = mbps(total, elapsed);
     82             std.debug.print("\r{d:6.1} MB {d:6.2} MB/s {d:6.1} Mbps ", .{
     83                 @as(f64, @floatFromInt(total)) / 1e6,
     84                 speed,
     85                 speed * 8,
     86             });
     87         }
     88     }
     89 
     90     const end_ns: u64 = @intCast(std.Io.Timestamp.now(io, .awake).nanoseconds);
     91     const elapsed = end_ns - start_ns;
     92     const speed = mbps(total, elapsed);
     93     const secs = @as(f64, @floatFromInt(elapsed)) / 1e9;
     94     std.debug.print("\n{d:.1} MB in {d:.2}s => {d:.2} MB/s  ({d:.1} Mbps)\n", .{
     95         @as(f64, @floatFromInt(total)) / 1e6, secs, speed, speed * 8,
     96     });
     97 }
     98 
     99 fn uploadTest(io: std.Io, allocator: std.mem.Allocator, url: []const u8, upload_bytes: usize) !void {
    100     var client: std.http.Client = .{
    101         .allocator = allocator,
    102         .io = io,
    103     };
    104     defer client.deinit();
    105 
    106     const uri = try std.Uri.parse(url);
    107 
    108     var req = try client.request(.POST, uri, .{
    109         .headers = .{
    110             .content_type = .{ .override = "application/octet-stream" },
    111         },
    112     });
    113     defer req.deinit();
    114     req.transfer_encoding = .{ .content_length = upload_bytes };
    115 
    116     var body = try req.sendBodyUnflushed(&.{});
    117 
    118     const chunk_size: usize = 131072;
    119     const chunk = try allocator.alloc(u8, chunk_size);
    120     defer allocator.free(chunk);
    121     @memset(chunk, 0x55);
    122 
    123     var sent: usize = 0;
    124     const start = std.Io.Timestamp.now(io, .awake);
    125     const start_ns: u64 = @intCast(start.nanoseconds);
    126 
    127     while (sent < upload_bytes) {
    128         const to_send = @min(chunk_size, upload_bytes - sent);
    129         try body.writer.writeAll(chunk[0..to_send]);
    130         sent += to_send;
    131 
    132         const now_ns: u64 = @intCast(std.Io.Timestamp.now(io, .awake).nanoseconds);
    133         const elapsed = now_ns - start_ns;
    134         const speed = mbps(sent, elapsed);
    135         std.debug.print("\r {d:6.1} MB {d:6.2} MB/s {d:6.1} Mbps ", .{
    136             @as(f64, @floatFromInt(sent)) / 1e6,
    137             speed,
    138             speed * 8,
    139         });
    140     }
    141 
    142     try body.end();
    143     if (req.connection) |conn| {
    144         try conn.flush();
    145     }
    146 
    147     const end_ns: u64 = @intCast(std.Io.Timestamp.now(io, .awake).nanoseconds);
    148 
    149     var redirect_buffer: [8192]u8 = undefined;
    150     _ = try req.receiveHead(&redirect_buffer);
    151 
    152     const elapsed = end_ns - start_ns;
    153     const speed = mbps(sent, elapsed);
    154     const secs = @as(f64, @floatFromInt(elapsed)) / 1e9;
    155     std.debug.print("\n{d:.1} MB in {d:.2}s => {d:.2} MB/s  ({d:.1} Mbps)\n", .{
    156         @as(f64, @floatFromInt(sent)) / 1e6, secs, speed, speed * 8,
    157     });
    158 }
    159 
    160 /// Returns the value for an option, taking it either from the rest of the
    161 /// current argument after '=' or from the next argument.
    162 fn optionValue(args: []const [:0]const u8, i: *usize, eq: ?usize) ![]const u8 {
    163     if (eq) |pos| return args[i.*][pos + 1 ..];
    164     if (i.* + 1 >= args.len) {
    165         std.debug.print("missing value for '{s}'\n", .{args[i.*]});
    166         return error.MissingValue;
    167     }
    168     i.* += 1;
    169     return args[i.*];
    170 }
    171 
    172 fn parseArgs(args: []const [:0]const u8) !?Config {
    173     // With no arguments, show usage instead of running the tests.
    174     if (args.len <= 1) {
    175         printUsage();
    176         return null;
    177     }
    178 
    179     var cfg: Config = .{};
    180 
    181     // Tracks which of the mutually exclusive mode flags has been seen.
    182     var mode: ?[]const u8 = null;
    183 
    184     var i: usize = 1;
    185     while (i < args.len) : (i += 1) {
    186         const arg = args[i];
    187 
    188         // Split "--flag=value" into name and value position.
    189         const eq = std.mem.indexOfScalar(u8, arg, '=');
    190         const name = if (eq) |pos| arg[0..pos] else arg;
    191 
    192         if (std.mem.eql(u8, name, "-h") or std.mem.eql(u8, name, "--help")) {
    193             printUsage();
    194             return null;
    195         } else if (std.mem.eql(u8, name, "--download-url")) {
    196             cfg.download_url = try optionValue(args, &i, eq);
    197         } else if (std.mem.eql(u8, name, "--upload-url")) {
    198             cfg.upload_url = try optionValue(args, &i, eq);
    199         } else if (std.mem.eql(u8, name, "--upload-mb")) {
    200             const value = try optionValue(args, &i, eq);
    201             const mb = std.fmt.parseInt(usize, value, 10) catch {
    202                 std.debug.print("invalid --upload-mb value: '{s}'\n", .{value});
    203                 return error.InvalidArgument;
    204             };
    205             cfg.upload_bytes = mb * 1024 * 1024;
    206         } else if (std.mem.eql(u8, name, "--both") or
    207             std.mem.eql(u8, name, "--download-only") or
    208             std.mem.eql(u8, name, "--upload-only"))
    209         {
    210             if (mode) |prev| {
    211                 std.debug.print("{s} and {s} are mutually exclusive\n", .{ prev, name });
    212                 return error.InvalidArgument;
    213             }
    214             mode = name;
    215             cfg.run_download = !std.mem.eql(u8, name, "--upload-only");
    216             cfg.run_upload = !std.mem.eql(u8, name, "--download-only");
    217         } else {
    218             std.debug.print("unknown argument: '{s}'\n", .{arg});
    219             printUsage();
    220             return error.UnknownArgument;
    221         }
    222     }
    223 
    224     return cfg;
    225 }
    226 
    227 pub fn main(init: std.process.Init) !void {
    228     const allocator = init.gpa;
    229 
    230     const args = try init.minimal.args.toSlice(init.arena.allocator());
    231     const cfg = parseArgs(args) catch std.process.exit(2) orelse return;
    232 
    233     if (cfg.run_download) {
    234         downloadTest(init.io, allocator, cfg.download_url) catch |err| {
    235             std.debug.print("Download error: {}\n", .{err});
    236         };
    237     }
    238 
    239     if (cfg.run_upload) {
    240         uploadTest(init.io, allocator, cfg.upload_url, cfg.upload_bytes) catch |err| {
    241             std.debug.print("Upload error: {}\n", .{err});
    242         };
    243     }
    244 }