From 77cc9e6a7095f7d8edc208fbb34bd34a5e859a46 Mon Sep 17 00:00:00 2001
From: Lars Kanis <lars@greiz-reinsdorf.de>
Date: Fri, 31 Mar 2023 22:40:18 +0200
Subject: [PATCH] Add PG::Connection#check_socket as a way to verify the socket
 is OK

... without data exchange with the server.

Related to #509
---
 ext/pg_connection.c        |  4 ++++
 lib/pg/connection.rb       | 21 ++++++++++++++++++++-
 spec/pg/connection_spec.rb | 34 ++++++++++++++++++++++++++++++++++
 3 files changed, 58 insertions(+), 1 deletion(-)

diff --git a/ext/pg_connection.c b/ext/pg_connection.c
index 6833a6b52..253c33aeb 100644
--- a/ext/pg_connection.c
+++ b/ext/pg_connection.c
@@ -763,6 +763,10 @@ pgconn_conninfo( VALUE self )
  *
  * ... and other constants of kind PG::Constants::CONNECTION_*
  *
+ * This method returns the status of the last command from memory.
+ * It doesn't do any socket access hence is not suitable to test the connectivity.
+ * See check_socket for a way to verify the socket state.
+ *
  * Example:
  *   PG.constants.grep(/CONNECTION_/).find{|c| PG.const_get(c) == conn.status} # => :CONNECTION_OK
  */
diff --git a/lib/pg/connection.rb b/lib/pg/connection.rb
index e2f6129bf..8b36afce1 100644
--- a/lib/pg/connection.rb
+++ b/lib/pg/connection.rb
@@ -343,6 +343,20 @@ def ssl_attributes
 		end
 	end
 
+	# Read all pending socket input to internal memory and raise an exception in case of errors.
+	#
+	# This verifies that the connection socket is in a usable state and not aborted in any way.
+	# No communication is done with the server.
+	# Only pending data is read from the socket - the method doesn't wait for any outstanding server answers.
+	#
+	# Raises a kind of PG::Error if there was an error reading the data or if the socket is in a failure state.
+	def check_socket
+		while socket_io.wait_readable(0)
+			consume_input
+		end
+		nil
+	end
+
 	# call-seq:
 	#    conn.get_result() -> PG::Result
 	#    conn.get_result() {|pg_result| block }
@@ -797,7 +811,10 @@ def new(*args)
 		#    PG::Connection.ping(connection_string)     -> Integer
 		#    PG::Connection.ping(host, port, options, tty, dbname, login, password) ->  Integer
 		#
-		# Check server status.
+		# PQpingParams reports the status of the server.
+		#
+		# It accepts connection parameters identical to those of PQ::Connection.new .
+		# It is not necessary to supply correct user name, password, or database name values to obtain the server status; however, if incorrect values are provided, the server will log a failed connection attempt.
 		#
 		# See PG::Connection.new for a description of the parameters.
 		#
@@ -810,6 +827,8 @@ def new(*args)
 		#   could not establish connection
 		# [+PQPING_NO_ATTEMPT+]
 		#   connection not attempted (bad params)
+		#
+		# See also check_socket for a way to check the connection without doing any server communication.
 		def ping(*args)
 			if Fiber.respond_to?(:scheduler) && Fiber.scheduler
 				# Run PQping in a second thread to avoid blocking of the scheduler.
diff --git a/spec/pg/connection_spec.rb b/spec/pg/connection_spec.rb
index 78dd611fd..ebf77e6f3 100644
--- a/spec/pg/connection_spec.rb
+++ b/spec/pg/connection_spec.rb
@@ -1678,6 +1678,40 @@
 		expect{ conn.consume_input }.to raise_error(PG::ConnectionBad, /can't get socket descriptor|connection not open/){|err| expect(err).to have_attributes(connection: conn) }
 	end
 
+	describe :check_socket do
+		it "does nothing if connection is OK" do
+			expect( @conn.check_socket ).to be_nil
+		end
+
+		def wait_check_socket(conn)
+			retries = 100
+			loop do
+				conn.check_socket
+				sleep 0.1
+				break if (retries-=1) < 0
+			end
+		end
+
+		it "raises error on broken connection" do
+			conn = PG.connect(@conninfo)
+			conn.send_query "SELECT pg_terminate_backend(pg_backend_pid())"
+			expect( conn.get_result.result_status ).to be( PG::PGRES_FATAL_ERROR )
+
+			expect do
+				wait_check_socket(conn)
+			end.to raise_error(PG::ConnectionBad, /SSL connection has been closed unexpectedly|server closed the connection unexpectedly/)
+		end
+
+		it "processes messages before connection error" do
+			conn = PG.connect(@conninfo)
+			conn.send_query "do $$ BEGIN RAISE NOTICE 'foo'; PERFORM pg_terminate_backend(pg_backend_pid()); END; $$ LANGUAGE plpgsql;"
+
+			expect do
+				wait_check_socket(conn)
+			end.to raise_error(PG::ConnectionBad, /SSL connection has been closed unexpectedly|server closed the connection unexpectedly/)
+		end
+	end
+
 	it "calls the block supplied to wait_for_notify with the notify payload if it accepts " +
 			"any number of arguments" do