From 401866b68a056d1464bb718851ab2cb960b3fedc Mon Sep 17 00:00:00 2001
From: Earlopain <14981592+Earlopain@users.noreply.github.com>
Date: Tue, 7 May 2024 21:54:25 +0200
Subject: [PATCH] [Prod] Simplify implementation of worker killer

Credits to byroot for pointing me in this direction
---
 config/pitchfork.rb | 15 +++++++++------
 1 file changed, 9 insertions(+), 6 deletions(-)

diff --git a/config/pitchfork.rb b/config/pitchfork.rb
index ee6e36741..d9c7734ec 100644
--- a/config/pitchfork.rb
+++ b/config/pitchfork.rb
@@ -11,25 +11,28 @@ timeout 180
 listen ENV.fetch("PITCHFORK_LISTEN_ADDRESS"), tcp_nopush: true, backlog: 2048
 worker_processes ENV.fetch("PITCHFORK_WORKER_COUNT").to_i
 
+# Each worker will have its own copy of this data structure.
+WorkerData = Data.define(:max_requests, :max_mem)
+worker_data = nil
+
 after_worker_ready do |server, worker|
   max_requests = Random.rand(5_000..10_000)
-  worker.instance_variable_set(:@_max_requests, max_requests)
   max_mem = Random.rand((386 * (1024**2))..(768 * (1024**2)))
-  worker.instance_variable_set(:@_max_mem, max_mem)
+  worker_data = WorkerData.new(max_requests: max_requests, max_mem: max_mem)
 
   server.logger.info("worker=#{worker.nr} gen=#{worker.generation} ready, serving #{max_requests} requests, #{max_mem} bytes")
 end
 
 after_request_complete do |server, worker, _env|
-  if worker.requests_count > worker.instance_variable_get(:@_max_requests)
-    server.logger.info("worker=#{worker.nr} gen=#{worker.generation}) exit: request limit (#{worker.instance_variable_get(:@_max_requests)})")
+  if worker.requests_count > worker_data.max_requests
+    server.logger.info("worker=#{worker.nr} gen=#{worker.generation}) exit: request limit (#{worker_data.max_requests})")
     exit # rubocop:disable Rails/Exit
   end
 
   if worker.requests_count % 16 == 0
     mem_info = Pitchfork::MemInfo.new(worker.pid)
-    if mem_info.pss > worker.instance_variable_get(:@_max_mem)
-      server.logger.info("worker=#{worker.nr} gen=#{worker.generation}) exit: memory limit (#{mem_info.pss} bytes > #{worker.instance_variable_get(:@_max_mem)} bytes)")
+    if mem_info.pss > worker_data.max_mem
+      server.logger.info("worker=#{worker.nr} gen=#{worker.generation}) exit: memory limit (#{mem_info.pss} bytes > #{worker_data.max_mem} bytes)")
       exit # rubocop:disable Rails/Exit
     end
   end