THRIFT-5214: Use peek to implement socket connectivity check

Client: go

In previous implementation of socket connectivity check, we try to read
1 byte and put it into buffer when succeeded. The buffer complicates
the implementation of Read function. Change to use syscall.Recvfrom with
MSG_PEEK flag instead so that the buffer is no longer needed.

3 files changed, 8 insertions, 27 deletions

diff --git a/lib/go/thrift/socket_conn.go b/lib/go/thrift/socket_conn.go
index 40aea61b1..c1cc30c6c 100644
--- a/lib/go/thrift/socket_conn.go
+++ b/lib/go/thrift/socket_conn.go
@@ -20,8 +20,6 @@
 package thrift
 
 import (
-	"bytes"
-	"io"
 	"net"
 )
 
@@ -29,7 +27,6 @@ import (
 type socketConn struct {
 	net.Conn
 
-	buf    bytes.Buffer
 	buffer [1]byte
 }
 
@@ -101,16 +98,5 @@ func (sc *socketConn) Read(p []byte) (n int, err error) {
 		return 0, sc.read0()
 	}
 
-	n, err = sc.buf.Read(p)
-	if err != nil && err != io.EOF {
-		return
-	}
-	if n == len(p) {
-		return n, nil
-	}
-	// Continue reading from the wire.
-	var newRead int
-	newRead, err = sc.Conn.Read(p[n:])
-	n += newRead
-	return
+	return sc.Conn.Read(p)
 }
diff --git a/lib/go/thrift/socket_unix_conn.go b/lib/go/thrift/socket_unix_conn.go
index 4535e750d..98e5a04dd 100644
--- a/lib/go/thrift/socket_unix_conn.go
+++ b/lib/go/thrift/socket_unix_conn.go
@@ -56,22 +56,20 @@ func (sc *socketConn) checkConn() error {
 	var n int
 
 	if readErr := rc.Read(func(fd uintptr) bool {
-		n, err = syscall.Read(int(fd), sc.buffer[:])
+		n, _, err = syscall.Recvfrom(int(fd), sc.buffer[:], syscall.MSG_PEEK|syscall.MSG_DONTWAIT)
 		return true
 	}); readErr != nil {
 		return readErr
 	}
 
-	if err == syscall.EAGAIN || err == syscall.EWOULDBLOCK {
-		// This means the connection is still open but we don't have
-		// anything to read right now.
+	if n > 0 {
+		// We got something, which means we are good
 		return nil
 	}
 
-	if n > 0 {
-		// We got something,
-		// put it to sc's buf for the next real read to use.
-		sc.buf.Write(sc.buffer[:n])
+	if err == syscall.EAGAIN || err == syscall.EWOULDBLOCK {
+		// This means the connection is still open but we don't have
+		// anything to read right now.
 		return nil
 	}
 
diff --git a/lib/go/thrift/socket_unix_conn_test.go b/lib/go/thrift/socket_unix_conn_test.go
index 3563a259c..1d4c806cf 100644
--- a/lib/go/thrift/socket_unix_conn_test.go
+++ b/lib/go/thrift/socket_unix_conn_test.go
@@ -75,7 +75,7 @@ func TestSocketConnUnix(t *testing.T) {
 		t.Error("Expected sc to report open, got false")
 	}
 	// Do connection check again twice after server already wrote new data,
-	// make sure we correctly buffered the read bytes
+	// make sure we don't cause any data loss with the check.
 	time.Sleep(interval * 10)
 	if !sc.IsOpen() {
 		t.Error("Expected sc to report open, got false")
@@ -83,9 +83,6 @@ func TestSocketConnUnix(t *testing.T) {
 	if !sc.IsOpen() {
 		t.Error("Expected sc to report open, got false")
 	}
-	if sc.buf.Len() == 0 {
-		t.Error("Expected sc to buffer read bytes, got empty buffer")
-	}
 	n, err = sc.Read(buf)
 	if err != nil {
 		t.Fatal(err)