Prevent deadlock on connection close and expose connection to the stream context (#52)

During an RPC, the generated method blocks until a response is received
or the passed context is cancelled, but if the connection is closed
while waiting for either event, the method remains blocked.

The new generated code also listens for a connection close event,
unblocking the method.

Signed-off-by: Luiz Aoqui <[email protected]>

Author: lgfa29

Makefile

@@ -0,0 +1,11 @@
+.PHONY: proto
+proto: install
+	protoc --go-frpc_out=./pkg/generator ./examples/test/test.proto
+
+.PHONY: install
+install:
+	go install ./protoc-gen-go-frpc
+
+.PHONY: test
+test: proto
+	go test -v ./...

go.mod

@@ -3,17 +3,17 @@ module github.com/loopholelabs/frpc-go
 go 1.22
 
 require (
-	github.com/loopholelabs/frisbee-go v0.10.0
+	github.com/loopholelabs/frisbee-go v0.10.2
 	github.com/loopholelabs/logging v0.3.1
-	github.com/loopholelabs/polyglot/v2 v2.0.2
+	github.com/loopholelabs/polyglot/v2 v2.0.3
 	github.com/loopholelabs/testing v0.2.3
 	github.com/stretchr/testify v1.9.0
 	google.golang.org/protobuf v1.35.1
 )
 
 require (
 	github.com/davecgh/go-spew v1.1.1 // indirect
-	github.com/loopholelabs/common v0.4.9 // indirect
+	github.com/loopholelabs/common v0.4.10 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 )

go.sum

@@ -2,18 +2,18 @@ github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/google/go-cmp v0.5.5 h1:Khx7svrCpmxxtHBq5j2mp/xVjsi8hQMfNLvJFAlrGgU=
 github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
-github.com/kr/pretty v0.1.0 h1:L/CwN0zerZDmRFUapSPitk6f+Q3+0za1rQkzVuMiMFI=
-github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORNo=
+github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE=
+github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk=
 github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
 github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE=
-github.com/loopholelabs/common v0.4.9 h1:9MPUYlZZ/qx3Kt8LXgXxcSXthrM91od8026c4DlGpAU=
-github.com/loopholelabs/common v0.4.9/go.mod h1:Wop5srN1wYT+mdQ9gZ+kn2I9qKAyVd0FB48pThwIa9M=
-github.com/loopholelabs/frisbee-go v0.10.0 h1:eqdDqm44V23GMxhjDL9OBz2Fxecsu42M0KHM8kQObBQ=
-github.com/loopholelabs/frisbee-go v0.10.0/go.mod h1:RwKglItbNcQq9UW6Vm2aOwOWXR6wTxUNjvuXcCIyJkU=
+github.com/loopholelabs/common v0.4.10 h1:BMJSMwH0PiVtdpOlXNPlW827B9WPJ/Gkb/q20NLeOjw=
+github.com/loopholelabs/common v0.4.10/go.mod h1:wc17hLpzZaDbndb7Fh3MXQDnhf4Cmf/JKC+LmXaD6II=
+github.com/loopholelabs/frisbee-go v0.10.2 h1:TvfONSpoCrrvG7/8QzBbHou2gPqOwN5iFL+4fp3bu4s=
+github.com/loopholelabs/frisbee-go v0.10.2/go.mod h1:JJhInJ5zjxpwOLAqEviOjTo1k8E3NNEF/fQTrC4lZ/4=
 github.com/loopholelabs/logging v0.3.1 h1:VA9DF3WrbmvJC1uQJ/XcWgz8KWXydWwe3BdDiMbN2FY=
 github.com/loopholelabs/logging v0.3.1/go.mod h1:uRDUydiqPqKbZkb0WoQ3dfyAcJ2iOMhxdEafZssLVv0=
-github.com/loopholelabs/polyglot/v2 v2.0.2 h1:v308fg2ZKSvkKDnWgBnDvvmiu4YypCxcDe5Ih5GUVnY=
-github.com/loopholelabs/polyglot/v2 v2.0.2/go.mod h1:kFoSKvnKAWmV0ICfbaCHDv/+cz5LSuA+xXG4WtYV/z4=
+github.com/loopholelabs/polyglot/v2 v2.0.3 h1:CpH2az5shkOgOBASnzjc1XC5SVzQMbWyHt4R7ds/FFc=
+github.com/loopholelabs/polyglot/v2 v2.0.3/go.mod h1:yodgE9ile0RS/npD0WnHfFpMLvL5FlC9n3qZ1tTkB9g=
 github.com/loopholelabs/testing v0.2.3 h1:4nVuK5ctaE6ua5Z0dYk2l7xTFmcpCYLUeGjRBp8keOA=
 github.com/loopholelabs/testing v0.2.3/go.mod h1:gqtGY91soYD1fQoKQt/6kP14OYpS7gcbcIgq5mc9m8Q=
 github.com/mattn/go-colorable v0.1.13 h1:fFA4WZxdEF4tXPZVKMLwD8oUnCTTo08duU7wxecdEvA=
@@ -22,6 +22,8 @@ github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWE
 github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/rogpeppe/go-internal v1.13.1 h1:KvO1DLK/DRN07sQ1LQKScxyZJuNnedQ5/wKSR38lUII=
+github.com/rogpeppe/go-internal v1.13.1/go.mod h1:uMEvuHeurkdAXX61udpOXGD/AzZDWNMNyH2VO9fmH0o=
 github.com/rs/zerolog v1.33.0 h1:1cU2KZkvPxNyfgEmhHAz/1A9Bz+llsdYzklWFzgp0r8=
 github.com/rs/zerolog v1.33.0/go.mod h1:/7mN4D5sKwJLZQ2b/znpjC3/GQWY/xaDXUM0kKWRHss=
 github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg=
@@ -35,7 +37,7 @@ golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8T
 google.golang.org/protobuf v1.35.1 h1:m3LfL6/Ca+fqnjnlqQXNpFPABW1UD7mjh8KO2mKFytA=
 google.golang.org/protobuf v1.35.1/go.mod h1:9fA7Ob0pmnwhb644+1+CVWFRbNajQ6iRojtC/QF5bRE=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
-gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127 h1:qIbj1fsPNlZgppZ+VLlY7N33q108Sa+fhmuc+sWQYwY=
-gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk=
+gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q=
 gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
 gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

pkg/generator/test/generator_test.go

@@ -4,11 +4,18 @@ package test
 
 import (
 	"context"
+	"fmt"
 	"io"
+	"net/http"
+	"net/http/httptest"
+	"os"
 	"testing"
+	"time"
 
+	"github.com/loopholelabs/polyglot/v2"
 	"github.com/loopholelabs/testing/conn/pair"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 func TestRPC(t *testing.T) {
@@ -130,3 +137,51 @@ func testClientStreaming(client *Client, t *testing.T) {
 	assert.NoError(t, err)
 	assert.Equal(t, "Hello World", res.Message)
 }
+
+func TestRPCInvalidConnection(t *testing.T) {
+	// Create non-Frisbee server the client can connect to but not exchange
+	// messages, so the connection will be broken soon after connect.
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		fmt.Fprintln(w, "test")
+	}))
+	t.Cleanup(ts.Close)
+
+	// Create a client and connect to test server.
+	client, err := NewClient(nil, nil)
+	require.NoError(t, err)
+
+	err = client.Connect(ts.Listener.Addr().String())
+	require.NoError(t, err)
+
+	// Make RPC request with a 3s timeout.
+	timeout := 3 * time.Second
+	if os.Getenv("CI") != "" {
+		timeout = 10 * time.Second
+	}
+	ctx, cancel := context.WithTimeout(context.Background(), timeout)
+	t.Cleanup(cancel)
+
+	req := &Request{
+		Message: "Hello World",
+		Corpus:  RequestUNIVERSAL,
+	}
+	response, err := client.EchoService.Echo(ctx, req)
+
+	// Verify request doesn't block forever.
+	require.NoError(t, ctx.Err())
+
+	// Verify request fails.
+	require.Error(t, err)
+	require.Nil(t, response)
+}
+
+func TestEncodeDecodePreservesNilFields(t *testing.T) {
+	r := &Response{Message: "test", Test: nil}
+	b := polyglot.NewBuffer()
+	r.Encode(b)
+
+	got := &Response{}
+	err := got.Decode(b.Bytes())
+	require.NoError(t, err)
+	require.Equal(t, r, got)
+}