Merge pull request #15125 from ziglang/hcs-win-poc

coff: add hot-code swapping PoC

4 files changed, 162 insertions, 41 deletions

diff --git a/src/link.zig b/src/link.zig
index 239dc646b2..dc114eb0ad 100644
--- a/src/link.zig
+++ b/src/link.zig
@@ -379,24 +379,30 @@ pub const File = struct {
                 if (base.file != null) return;
                 const emit = base.options.emit orelse return;
                 if (base.child_pid) |pid| {
-                    // If we try to open the output file in write mode while it is running,
-                    // it will return ETXTBSY. So instead, we copy the file, atomically rename it
-                    // over top of the exe path, and then proceed normally. This changes the inode,
-                    // avoiding the error.
-                    const tmp_sub_path = try std.fmt.allocPrint(base.allocator, "{s}-{x}", .{
-                        emit.sub_path, std.crypto.random.int(u32),
-                    });
-                    try emit.directory.handle.copyFile(emit.sub_path, emit.directory.handle, tmp_sub_path, .{});
-                    try emit.directory.handle.rename(tmp_sub_path, emit.sub_path);
-                    switch (builtin.os.tag) {
-                        .linux => std.os.ptrace(std.os.linux.PTRACE.ATTACH, pid, 0, 0) catch |err| {
-                            log.warn("ptrace failure: {s}", .{@errorName(err)});
-                        },
-                        .macos => base.cast(MachO).?.ptraceAttach(pid) catch |err| {
+                    if (builtin.os.tag == .windows) {
+                        base.cast(Coff).?.ptraceAttach(pid) catch |err| {
                             log.warn("attaching failed with error: {s}", .{@errorName(err)});
-                        },
-                        .windows => {},
-                        else => return error.HotSwapUnavailableOnHostOperatingSystem,
+                        };
+                    } else {
+                        // If we try to open the output file in write mode while it is running,
+                        // it will return ETXTBSY. So instead, we copy the file, atomically rename it
+                        // over top of the exe path, and then proceed normally. This changes the inode,
+                        // avoiding the error.
+                        const tmp_sub_path = try std.fmt.allocPrint(base.allocator, "{s}-{x}", .{
+                            emit.sub_path, std.crypto.random.int(u32),
+                        });
+                        try emit.directory.handle.copyFile(emit.sub_path, emit.directory.handle, tmp_sub_path, .{});
+                        try emit.directory.handle.rename(tmp_sub_path, emit.sub_path);
+                        switch (builtin.os.tag) {
+                            .linux => std.os.ptrace(std.os.linux.PTRACE.ATTACH, pid, 0, 0) catch |err| {
+                                log.warn("ptrace failure: {s}", .{@errorName(err)});
+                            },
+                            .macos => base.cast(MachO).?.ptraceAttach(pid) catch |err| {
+                                log.warn("attaching failed with error: {s}", .{@errorName(err)});
+                            },
+                            .windows => unreachable,
+                            else => return error.HotSwapUnavailableOnHostOperatingSystem,
+                        }
                     }
                 }
                 base.file = try emit.directory.handle.createFile(emit.sub_path, .{
@@ -437,7 +443,7 @@ pub const File = struct {
                         .macos => base.cast(MachO).?.ptraceDetach(pid) catch |err| {
                             log.warn("detaching failed with error: {s}", .{@errorName(err)});
                         },
-                        .windows => {},
+                        .windows => base.cast(Coff).?.ptraceDetach(pid),
                         else => return error.HotSwapUnavailableOnHostOperatingSystem,
                     }
                 }
diff --git a/src/link/Coff.zig b/src/link/Coff.zig
index 7b5539287d..f3068f01a9 100644
--- a/src/link/Coff.zig
+++ b/src/link/Coff.zig
@@ -89,6 +89,20 @@ relocs: RelocTable = .{},
 /// this will be a table indexed by index into the list of Atoms.
 base_relocs: BaseRelocationTable = .{},
 
+/// Hot-code swapping state.
+hot_state: if (is_hot_update_compatible) HotUpdateState else struct {} = .{},
+
+const is_hot_update_compatible = switch (builtin.target.os.tag) {
+    .windows => true,
+    else => false,
+};
+
+const HotUpdateState = struct {
+    /// Base address at which the process (image) got loaded.
+    /// We need this info to correctly slide pointers when relocating.
+    loaded_base_address: ?std.os.windows.HMODULE = null,
+};
+
 const Entry = struct {
     target: SymbolWithLoc,
     // Index into the synthetic symbol table (i.e., file == null).
@@ -772,13 +786,87 @@ fn writeAtom(self: *Coff, atom_index: Atom.Index, code: []u8) !void {
     const sym = atom.getSymbol(self);
     const section = self.sections.get(@enumToInt(sym.section_number) - 1);
     const file_offset = section.header.pointer_to_raw_data + sym.value - section.header.virtual_address;
+
     log.debug("writing atom for symbol {s} at file offset 0x{x} to 0x{x}", .{
         atom.getName(self),
         file_offset,
         file_offset + code.len,
     });
-    self.resolveRelocs(atom_index, code);
+
+    const gpa = self.base.allocator;
+
+    // Gather relocs which can be resolved.
+    // We need to do this as we will be applying different slide values depending
+    // if we are running in hot-code swapping mode or not.
+    // TODO: how crazy would it be to try and apply the actual image base of the loaded
+    // process for the in-file values rather than the Windows defaults?
+    var relocs = std.ArrayList(*Relocation).init(gpa);
+    defer relocs.deinit();
+
+    if (self.relocs.getPtr(atom_index)) |rels| {
+        try relocs.ensureTotalCapacityPrecise(rels.items.len);
+        for (rels.items) |*reloc| {
+            if (reloc.isResolvable(self)) relocs.appendAssumeCapacity(reloc);
+        }
+    }
+
+    if (is_hot_update_compatible) {
+        if (self.base.child_pid) |handle| {
+            const slide = @ptrToInt(self.hot_state.loaded_base_address.?);
+
+            const mem_code = try gpa.dupe(u8, code);
+            defer gpa.free(mem_code);
+            self.resolveRelocs(atom_index, relocs.items, mem_code, slide);
+
+            const vaddr = sym.value + slide;
+            const pvaddr = @intToPtr(*anyopaque, vaddr);
+
+            log.debug("writing to memory at address {x}", .{vaddr});
+
+            if (build_options.enable_logging) {
+                try debugMem(gpa, handle, pvaddr, mem_code);
+            }
+
+            if (section.header.flags.MEM_WRITE == 0) {
+                writeMemProtected(handle, pvaddr, mem_code) catch |err| {
+                    log.warn("writing to protected memory failed with error: {s}", .{@errorName(err)});
+                };
+            } else {
+                writeMem(handle, pvaddr, mem_code) catch |err| {
+                    log.warn("writing to protected memory failed with error: {s}", .{@errorName(err)});
+                };
+            }
+        }
+    }
+
+    self.resolveRelocs(atom_index, relocs.items, code, self.getImageBase());
     try self.base.file.?.pwriteAll(code, file_offset);
+
+    // Now we can mark the relocs as resolved.
+    while (relocs.popOrNull()) |reloc| {
+        reloc.dirty = false;
+    }
+}
+
+fn debugMem(allocator: Allocator, handle: std.ChildProcess.Id, pvaddr: std.os.windows.LPVOID, code: []const u8) !void {
+    var buffer = try allocator.alloc(u8, code.len);
+    defer allocator.free(buffer);
+    const memread = try std.os.windows.ReadProcessMemory(handle, pvaddr, buffer);
+    log.debug("to write: {x}", .{std.fmt.fmtSliceHexLower(code)});
+    log.debug("in memory: {x}", .{std.fmt.fmtSliceHexLower(memread)});
+}
+
+fn writeMemProtected(handle: std.ChildProcess.Id, pvaddr: std.os.windows.LPVOID, code: []const u8) !void {
+    const old_prot = try std.os.windows.VirtualProtectEx(handle, pvaddr, code.len, std.os.windows.PAGE_EXECUTE_WRITECOPY);
+    try writeMem(handle, pvaddr, code);
+    // TODO: We can probably just set the pages writeable and leave it at that without having to restore the attributes.
+    // For that though, we want to track which page has already been modified.
+    _ = try std.os.windows.VirtualProtectEx(handle, pvaddr, code.len, old_prot);
+}
+
+fn writeMem(handle: std.ChildProcess.Id, pvaddr: std.os.windows.LPVOID, code: []const u8) !void {
+    const amt = try std.os.windows.WriteProcessMemory(handle, pvaddr, code);
+    if (amt != code.len) return error.InputOutput;
 }
 
 fn writePtrWidthAtom(self: *Coff, atom_index: Atom.Index) !void {
@@ -814,19 +902,30 @@ fn markRelocsDirtyByAddress(self: *Coff, addr: u32) void {
     }
 }
 
-fn resolveRelocs(self: *Coff, atom_index: Atom.Index, code: []u8) void {
-    const relocs = self.relocs.getPtr(atom_index) orelse return;
-
+fn resolveRelocs(self: *Coff, atom_index: Atom.Index, relocs: []*const Relocation, code: []u8, image_base: u64) void {
     log.debug("relocating '{s}'", .{self.getAtom(atom_index).getName(self)});
-
-    for (relocs.items) |*reloc| {
-        if (!reloc.dirty) continue;
-        if (reloc.resolve(atom_index, code, self)) {
-            reloc.dirty = false;
-        }
+    for (relocs) |reloc| {
+        reloc.resolve(atom_index, code, image_base, self);
     }
 }
 
+pub fn ptraceAttach(self: *Coff, handle: std.ChildProcess.Id) !void {
+    if (!is_hot_update_compatible) return;
+
+    log.debug("attaching to process with handle {*}", .{handle});
+    self.hot_state.loaded_base_address = std.os.windows.ProcessBaseAddress(handle) catch |err| {
+        log.warn("failed to get base address for the process with error: {s}", .{@errorName(err)});
+        return;
+    };
+}
+
+pub fn ptraceDetach(self: *Coff, handle: std.ChildProcess.Id) void {
+    if (!is_hot_update_compatible) return;
+
+    log.debug("detaching from process with handle {*}", .{handle});
+    self.hot_state.loaded_base_address = null;
+}
+
 fn freeAtom(self: *Coff, atom_index: Atom.Index) void {
     log.debug("freeAtom {d}", .{atom_index});
 
@@ -1421,7 +1520,7 @@ pub fn flushModule(self: *Coff, comp: *Compilation, prog_node: *std.Progress.Nod
 
     for (self.relocs.keys(), self.relocs.values()) |atom_index, relocs| {
         const needs_update = for (relocs.items) |reloc| {
-            if (reloc.dirty) break true;
+            if (reloc.isResolvable(self)) break true;
         } else false;
 
         if (!needs_update) continue;
diff --git a/src/link/Coff/Relocation.zig b/src/link/Coff/Relocation.zig
index 37bd3e292f..2fafa0bbdc 100644
--- a/src/link/Coff/Relocation.zig
+++ b/src/link/Coff/Relocation.zig
@@ -72,14 +72,18 @@ pub fn getTargetAddress(self: Relocation, coff_file: *const Coff) ?u32 {
     }
 }
 
-/// Returns `false` if obtaining the target address has been deferred until `flushModule`.
-/// This can happen when trying to resolve address of an import table entry ahead of time.
-pub fn resolve(self: Relocation, atom_index: Atom.Index, code: []u8, coff_file: *Coff) bool {
+/// Returns true if and only if the reloc is dirty AND the target address is available.
+pub fn isResolvable(self: Relocation, coff_file: *Coff) bool {
+    _ = self.getTargetAddress(coff_file) orelse return false;
+    return self.dirty;
+}
+
+pub fn resolve(self: Relocation, atom_index: Atom.Index, code: []u8, image_base: u64, coff_file: *Coff) void {
     const atom = coff_file.getAtom(atom_index);
     const source_sym = atom.getSymbol(coff_file);
     const source_vaddr = source_sym.value + self.offset;
 
-    const target_vaddr = self.getTargetAddress(coff_file) orelse return false;
+    const target_vaddr = self.getTargetAddress(coff_file).?; // Oops, you didn't check if the relocation can be resolved with isResolvable().
     const target_vaddr_with_addend = target_vaddr + self.addend;
 
     log.debug("  ({x}: [() => 0x{x} ({s})) ({s}) ", .{
@@ -92,7 +96,7 @@ pub fn resolve(self: Relocation, atom_index: Atom.Index, code: []u8, coff_file:
     const ctx: Context = .{
         .source_vaddr = source_vaddr,
         .target_vaddr = target_vaddr_with_addend,
-        .image_base = coff_file.getImageBase(),
+        .image_base = image_base,
         .code = code,
         .ptr_width = coff_file.ptr_width,
     };
@@ -102,8 +106,6 @@ pub fn resolve(self: Relocation, atom_index: Atom.Index, code: []u8, coff_file:
         .x86, .x86_64 => self.resolveX86(ctx),
         else => unreachable, // unhandled target architecture
     }
-
-    return true;
 }
 
 const Context = struct {
diff --git a/src/main.zig b/src/main.zig
index c96fd25766..1a445107ba 100644
--- a/src/main.zig
+++ b/src/main.zig
@@ -3817,11 +3817,25 @@ fn runOrTestHotSwap(
     runtime_args_start: ?usize,
 ) !std.ChildProcess.Id {
     const exe_emit = comp.bin_file.options.emit.?;
-    // A naive `directory.join` here will indeed get the correct path to the binary,
-    // however, in the case of cwd, we actually want `./foo` so that the path can be executed.
-    const exe_path = try fs.path.join(gpa, &[_][]const u8{
-        exe_emit.directory.path orelse ".", exe_emit.sub_path,
-    });
+
+    const exe_path = switch (builtin.target.os.tag) {
+        // On Windows it seems impossible to perform an atomic rename of a file that is currently
+        // running in a process. Therefore, we do the opposite. We create a copy of the file in
+        // tmp zig-cache and use it to spawn the child process. This way we are free to update
+        // the binary with each requested hot update.
+        .windows => blk: {
+            try exe_emit.directory.handle.copyFile(exe_emit.sub_path, comp.local_cache_directory.handle, exe_emit.sub_path, .{});
+            break :blk try fs.path.join(gpa, &[_][]const u8{
+                comp.local_cache_directory.path orelse ".", exe_emit.sub_path,
+            });
+        },
+
+        // A naive `directory.join` here will indeed get the correct path to the binary,
+        // however, in the case of cwd, we actually want `./foo` so that the path can be executed.
+        else => try fs.path.join(gpa, &[_][]const u8{
+            exe_emit.directory.path orelse ".", exe_emit.sub_path,
+        }),
+    };
     defer gpa.free(exe_path);
 
     var argv = std.ArrayList([]const u8).init(gpa);