From b077335a305278475e3e6386e78afc4d013be9c8 Mon Sep 17 00:00:00 2001
From: Sam <sam.saffron@gmail.com>
Date: Fri, 21 Apr 2017 17:12:29 -0400
Subject: [PATCH] make stat socket much more robust

---
 lib/stats_socket.rb                  | 23 +++++++++++++++++++----
 spec/components/stats_socket_spec.rb |  8 ++++++++
 2 files changed, 27 insertions(+), 4 deletions(-)

diff --git a/lib/stats_socket.rb b/lib/stats_socket.rb
index 23f6ebc2e0b..f4367c12c9d 100644
--- a/lib/stats_socket.rb
+++ b/lib/stats_socket.rb
@@ -38,16 +38,31 @@ class StatsSocket
       return false
     end
 
-    if IO.select(nil, [socket], nil, 10)
-      line = socket.read_nonblock(1000)
+    start = Time.now
+    line = ""
+
+    while Time.now - start < 10
+      if IO.select(nil, [socket], nil, 10)
+        begin
+          line << socket.read_nonblock(1000)
+        rescue IO::WaitReadable
+          sleep 0.001
+        end
+      end
+      break if line.include?("\n")
+    end
+
+    if line.include?("\n")
       socket.write get_response(line.strip)
     end
-    socket.close
+
     true
-  rescue IOError
+  rescue IOError => e
     # nothing to do here, case its normal on shutdown
   rescue => e
     Rails.logger.warn("Failed to handle connection in stats socket #{e}")
+  ensure
+    socket&.close rescue nil
   end
 
   def get_response(command)
diff --git a/spec/components/stats_socket_spec.rb b/spec/components/stats_socket_spec.rb
index 9352a6da531..cac144439ee 100644
--- a/spec/components/stats_socket_spec.rb
+++ b/spec/components/stats_socket_spec.rb
@@ -29,6 +29,14 @@ describe StatsSocket do
       socket.close
     end
 
+    socket = UNIXSocket.new(socket_path)
+    socket.send "gc_st", 0
+    socket.flush
+    sleep 0.001
+    socket.send "at\n", 0
+    line = socket.readline
+    socket.close
+
     parsed = JSON.parse(line)
 
     expect(parsed.keys.sort).to eq(GC.stat.keys.map(&:to_s).sort)