github.com/nycdavid/zeus@v0.0.0-20201208104106-9ba439429e03/rubygem/lib/zeus.rb (about) 1 # encoding: utf-8 2 require 'socket' 3 4 # load exact json version from Gemfile.lock to avoid conflicts 5 gemfile = "#{ENV["BUNDLE_GEMFILE"] || "Gemfile"}.lock" 6 if File.exist?(gemfile) && version = File.read(gemfile)[/^ json \((.*)\)/, 1] 7 gem 'json', version 8 end 9 require 'json' 10 require 'pty' 11 require 'set' 12 13 require 'zeus/load_tracking' 14 require 'zeus/plan' 15 require 'zeus/version' 16 require 'zeus/pry' 17 18 module Zeus 19 class << self 20 attr_accessor :plan, :dummy_tty, :master_socket 21 22 # this is totally asinine, but readline gets super confused when it's 23 # required at a time when stdin or stdout is not connected to a TTY, 24 # no matter what we do to tell it otherwise later. So we create a dummy 25 # TTY in case readline is required. 26 # 27 # Yup. 28 def setup_dummy_tty! 29 return if self.dummy_tty 30 master, self.dummy_tty = PTY.send(:open) 31 Thread.new { 32 loop { master.read(1024) } 33 } 34 STDIN.reopen(dummy_tty) 35 STDOUT.reopen(dummy_tty) 36 end 37 38 def setup_master_socket! 39 return master_socket if master_socket 40 41 fd = ENV['ZEUS_MASTER_FD'].to_i 42 self.master_socket = UNIXSocket.for_fd(fd) 43 end 44 45 def go(identifier=:boot) 46 # Thanks to the magic of fork, this following line will return 47 # many times: Every time the parent step receives a request to 48 # run a command. 49 if run_command = boot_steps(identifier) 50 ident, local = run_command 51 return command(ident, local) 52 end 53 end 54 55 def boot_steps(identifier) 56 while true 57 boot_step = catch(:boot_step) do 58 $0 = "zeus slave: #{identifier}" 59 60 setup_dummy_tty! 61 master = setup_master_socket! 62 feature_pipe_r, feature_pipe_w = IO.pipe 63 64 # I need to give the master a way to talk to me exclusively 65 local, remote = UNIXSocket.pair(Socket::SOCK_STREAM) 66 master.send_io(remote) 67 68 # Now I need to tell the master about my PID and ID 69 local.write "P:#{Process.pid}:#{@parent_pid || 0}:#{identifier}\0" 70 local.send_io(feature_pipe_r) 71 feature_pipe_r.close 72 73 Zeus::LoadTracking.set_feature_pipe(feature_pipe_w) 74 75 run_action(local, identifier) 76 77 # We are now 'connected'. From this point, we may receive requests to fork. 78 children = Set.new 79 while true 80 messages = local.recv(2**16) 81 82 # Reap any child runners or slaves that might have exited in 83 # the meantime. Note that reaping them like this can leave <=1 84 # zombie process per slave around while the slave waits for a 85 # new command. 86 children.each do |pid| 87 children.delete(pid) if Process.waitpid(pid, Process::WNOHANG) 88 end 89 90 messages.split("\0").each do |new_identifier| 91 new_identifier =~ /^(.):(.*)/ 92 code, ident = $1, $2 93 94 forked_from = Process.pid 95 96 pid = fork 97 if pid 98 # We're in the parent. Record the child: 99 children << pid 100 elsif code == "S" 101 # Child, supposed to start another step: 102 @parent_pid = forked_from 103 104 Zeus::LoadTracking.clear_feature_pipe 105 106 throw(:boot_step, ident.to_sym) 107 else 108 # Child, supposed to run a command: 109 @parent_pid = forked_from 110 111 Zeus::LoadTracking.clear_feature_pipe 112 113 return [ident.to_sym, local] 114 end 115 end 116 end 117 end 118 identifier = boot_step 119 end 120 end 121 122 private 123 124 def command(identifier, sock) 125 $0 = "zeus runner: #{identifier}" 126 Process.setsid 127 128 local, remote = UNIXSocket.pair(:DGRAM) 129 sock.send_io(remote) 130 remote.close 131 sock.close 132 133 pid_and_argument_count = local.recv(2**16) 134 pid_and_argument_count.chomp("\0") =~ /(.*?):(.*)/ 135 client_pid, argument_count = $1.to_i, $2.to_i 136 arg_io = local.recv_io 137 arguments = arg_io.read.chomp("\0").split("\0") 138 139 if arguments.length != argument_count 140 raise "Argument count mismatch: Expected #{argument_count}, got #{arguments.length}" 141 end 142 143 pid = fork { 144 $0 = "zeus command: #{identifier}" 145 146 plan.after_fork 147 remote_stdin_stdout = local.recv_io 148 remote_stderr = local.recv_io 149 local.write "P:#{Process.pid}:#{@parent_pid}:\0" 150 local.close 151 152 $stdin.reopen(remote_stdin_stdout) 153 $stdout.reopen(remote_stdin_stdout) 154 $stderr.reopen(remote_stderr) 155 ARGV.replace(arguments) 156 157 plan.send(identifier) 158 } 159 160 kill_command_if_client_quits!(pid, client_pid) 161 162 Process.wait(pid) 163 code = $?.exitstatus || 0 164 165 local.write "#{code}\0" 166 167 local.close 168 rescue Exception 169 # If anything at all went wrong, kill the client - if anything 170 # went wrong before the runner can clean up, it might hang 171 # around forever. 172 Process.kill(:TERM, client_pid) 173 end 174 175 def kill_command_if_client_quits!(command_pid, client_pid) 176 Thread.new { 177 loop { 178 begin 179 Process.kill(0, client_pid) 180 rescue Errno::ESRCH 181 Process.kill(9, command_pid) 182 exit 0 183 end 184 sleep 1 185 } 186 } 187 end 188 189 def report_error_to_master(local, error) 190 str = "R:" 191 str << "#{error.backtrace[0]}: #{error.message} (#{error.class})\n" 192 error.backtrace[1..-1].each do |line| 193 str << "\tfrom #{line}\n" 194 end 195 str << "\0" 196 local.write str 197 end 198 199 def run_action(socket, identifier) 200 # Now we run the action and report its success/fail status to the master. 201 begin 202 Zeus::LoadTracking.track_features_loaded_by do 203 plan.after_fork unless identifier == :boot 204 plan.send(identifier) 205 end 206 207 socket.write "R:OK\0" 208 rescue => err 209 report_error_to_master(socket, err) 210 raise 211 end 212 end 213 end 214 end