memcached.zig (7385B)
1 const std = @import("std"); 2 const io = @import("io.zig"); 3 4 pub const Memcached = struct { 5 allocator: std.mem.Allocator, 6 stream: std.Io.net.Stream, 7 read_buf: [65536]u8 = @splat(0), 8 write_buf: [4096]u8 = @splat(0), 9 // `null` until first use; the Stream's buffered Reader/Writer hold pointers 10 // into this struct's own buffers, so they can only be created once the 11 // struct sits at its final address (i.e. after `init` returns by value). 12 reader: ?std.Io.net.Stream.Reader = null, 13 writer: ?std.Io.net.Stream.Writer = null, 14 15 pub fn init(allocator: std.mem.Allocator, host: []const u8, port: u16) !Memcached { 16 const address = try std.Io.net.IpAddress.parseIp4(host, port); 17 const stream = try address.connect(io.rt(), .{ .mode = .stream }); 18 return .{ 19 .allocator = allocator, 20 .stream = stream, 21 }; 22 } 23 24 pub fn deinit(self: *Memcached) void { 25 if (self.writer) |*writer| { 26 writer.interface.flush() catch |err| { 27 io.warn("failed to flush memcached on deinit: {}\n", .{err}); 28 }; 29 } 30 self.stream.close(io.rt()); 31 } 32 33 fn ensureStarted(self: *Memcached) void { 34 if (self.reader != null) return; 35 self.reader = self.stream.reader(io.rt(), &self.read_buf); 36 self.writer = self.stream.writer(io.rt(), &self.write_buf); 37 } 38 39 fn r(self: *Memcached) *std.Io.Reader { 40 self.ensureStarted(); 41 return &self.reader.?.interface; 42 } 43 44 fn w(self: *Memcached) *std.Io.Writer { 45 self.ensureStarted(); 46 return &self.writer.?.interface; 47 } 48 49 fn flushWrite(self: *Memcached) !void { 50 try self.w().flush(); 51 } 52 53 fn bufferedWrite(self: *Memcached, data: []const u8) !void { 54 try self.w().writeAll(data); 55 } 56 57 const chunk_size = 900 * 1024; 58 const chunk_prefix = "\x00chunked:"; 59 60 pub fn set(self: *Memcached, key: []const u8, value: []const u8, exptime: u32) !void { 61 if (value.len <= chunk_size) { 62 try self.sendSet(key, value, exptime); 63 return; 64 } 65 66 const num_chunks = (value.len + chunk_size - 1) / chunk_size; 67 var meta_buf: [128]u8 = undefined; 68 const meta = try std.fmt.bufPrint(&meta_buf, chunk_prefix ++ "{d}:{d}", .{ num_chunks, value.len }); 69 try self.sendSet(key, meta, exptime); 70 71 var i: usize = 0; 72 while (i < num_chunks) : (i += 1) { 73 const start = i * chunk_size; 74 const end = @min((i + 1) * chunk_size, value.len); 75 var chunk_key_buf: [512]u8 = undefined; 76 const chunk_key = try std.fmt.bufPrint(&chunk_key_buf, "{s}:{d}", .{ key, i }); 77 try self.sendSet(chunk_key, value[start..end], exptime); 78 } 79 } 80 81 fn sendSet(self: *Memcached, key: []const u8, value: []const u8, exptime: u32) !void { 82 var header_buf: [512]u8 = undefined; 83 const header = try std.fmt.bufPrint(&header_buf, "ms {s} {d} T{d} q\r\n", .{ key, value.len, exptime }); 84 85 try self.bufferedWrite(header); 86 if (value.len > 0) { 87 try self.bufferedWrite(value); 88 } 89 try self.bufferedWrite("\r\n"); 90 } 91 92 fn readLine(self: *Memcached, buf: []u8) !?[]const u8 { 93 // Inclusive so the '\n' is consumed from the stream; exclusive would 94 // leave it buffered and desync the following `readExact`. 95 const line = self.r().takeDelimiterInclusive('\n') catch |err| switch (err) { 96 error.EndOfStream => return null, 97 error.StreamTooLong => return error.LineTooLong, 98 else => |e| return e, 99 }; 100 const trimmed = std.mem.trim(u8, line, "\r\n"); 101 if (buf.len < trimmed.len) return error.BufferTooSmall; 102 @memcpy(buf[0..trimmed.len], trimmed); 103 return buf[0..trimmed.len]; 104 } 105 106 fn readExact(self: *Memcached, buf: []u8) !void { 107 self.r().readSliceAll(buf) catch |err| switch (err) { 108 error.EndOfStream => return error.IncompleteRead, 109 else => |e| return e, 110 }; 111 } 112 113 fn consumeCrlf(self: *Memcached) void { 114 const reader = self.r(); 115 const first = reader.peekByte() catch return; 116 if (first == '\r') { 117 reader.toss(1); 118 const second = reader.peekByte() catch return; 119 if (second == '\n') reader.toss(1); 120 } else if (first == '\n') { 121 reader.toss(1); 122 } 123 } 124 125 pub fn get(self: *Memcached, key: []const u8) !?[]u8 { 126 var cmd_buf: [512]u8 = undefined; 127 const cmd = try std.fmt.bufPrint(&cmd_buf, "mg {s} v\r\n", .{key}); 128 try self.bufferedWrite(cmd); 129 try self.flushWrite(); 130 131 var line_buf: [512]u8 = undefined; 132 const line = (try self.readLine(&line_buf)) orelse return null; 133 134 if (std.mem.startsWith(u8, line, "VA ")) { 135 const rest = line[3..]; 136 const space_idx = std.mem.indexOfScalar(u8, rest, ' ') orelse rest.len; 137 const size = try std.fmt.parseInt(usize, rest[0..space_idx], 10); 138 139 const data = try self.allocator.alloc(u8, size); 140 errdefer self.allocator.free(data); 141 142 if (size > 0) { 143 try self.readExact(data); 144 } 145 146 self.consumeCrlf(); 147 148 if (std.mem.startsWith(u8, data, chunk_prefix)) { 149 defer self.allocator.free(data); 150 return try self.reassembleChunks(key, data); 151 } 152 153 return data; 154 } 155 156 return null; 157 } 158 159 fn reassembleChunks(self: *Memcached, key: []const u8, meta: []const u8) ![]u8 { 160 var it = std.mem.splitScalar(u8, meta[chunk_prefix.len..], ':'); 161 const num_chunks_str = it.next() orelse return error.InvalidCacheData; 162 const total_size_str = it.next() orelse return error.InvalidCacheData; 163 164 const num_chunks = try std.fmt.parseInt(usize, num_chunks_str, 10); 165 const total_size = try std.fmt.parseInt(usize, total_size_str, 10); 166 167 const result = try self.allocator.alloc(u8, total_size); 168 errdefer self.allocator.free(result); 169 170 var i: usize = 0; 171 while (i < num_chunks) : (i += 1) { 172 var chunk_key_buf: [512]u8 = undefined; 173 const chunk_key = try std.fmt.bufPrint(&chunk_key_buf, "{s}:{d}", .{ key, i }); 174 175 var cmd_buf: [512]u8 = undefined; 176 const cmd = try std.fmt.bufPrint(&cmd_buf, "mg {s} v\r\n", .{chunk_key}); 177 try self.bufferedWrite(cmd); 178 try self.flushWrite(); 179 180 var line_buf: [512]u8 = undefined; 181 const line = (try self.readLine(&line_buf)) orelse return error.CacheChunkMissing; 182 183 if (!std.mem.startsWith(u8, line, "VA ")) return error.CacheChunkMissing; 184 185 const rest = line[3..]; 186 const space_idx = std.mem.indexOfScalar(u8, rest, ' ') orelse rest.len; 187 const size = try std.fmt.parseInt(usize, rest[0..space_idx], 10); 188 189 const chunk_start = i * chunk_size; 190 const chunk_end = @min((i + 1) * chunk_size, total_size); 191 if (size != chunk_end - chunk_start) return error.InvalidChunkSize; 192 193 if (size > 0) { 194 try self.readExact(result[chunk_start..chunk_end]); 195 } 196 197 self.consumeCrlf(); 198 } 199 200 return result; 201 } 202 };