RPCChainVM fail-fast health RPCs (#2123)

Signed-off-by: Sam Batschelet <[email protected]>
ava-labs · Oct 4, 2023 · 3b6514d · 3b6514d
1 parent 9aa3d25
commit 3b6514d
Show file tree

Hide file tree

Showing 4 changed files with 29 additions and 11 deletions.
diff --git a/vms/rpcchainvm/factory.go b/vms/rpcchainvm/factory.go
@@ -13,8 +13,6 @@ import (
 	"github.com/ava-labs/avalanchego/vms/rpcchainvm/grpcutils"
 	"github.com/ava-labs/avalanchego/vms/rpcchainvm/runtime"
 	"github.com/ava-labs/avalanchego/vms/rpcchainvm/runtime/subprocess"
-
-	vmpb "github.com/ava-labs/avalanchego/proto/pb/vm"
 )
 
 var _ vms.Factory = (*factory)(nil)
@@ -61,7 +59,7 @@ func (f *factory) New(log logging.Logger) (interface{}, error) {
 		return nil, err
 	}
 
-	vm := NewClient(vmpb.NewVMClient(clientConn))
+	vm := NewClient(clientConn)
 	vm.SetProcess(stopper, status.Pid, f.processTracker)
 
 	f.runtimeTracker.TrackRuntime(stopper)

diff --git a/vms/rpcchainvm/grpcutils/client_test.go b/vms/rpcchainvm/grpcutils/client_test.go
@@ -4,6 +4,8 @@
 package grpcutils
 
 import (
+	"context"
+	"fmt"
 	"testing"
 	"time"
 
@@ -55,12 +57,11 @@ func TestWaitForReady(t *testing.T) {
 		Serve(listener, server)
 	}()
 
-	// The default includes grpc.WaitForReady(true).
+	// The default is WaitForReady = true.
 	conn, err := Dial(listener.Addr().String())
 	require.NoError(err)
 
 	db := rpcdb.NewClient(pb.NewDatabaseClient(conn))
-
 	require.NoError(db.Put([]byte("foo"), []byte("bar")))
 
 	noWaitListener, err := NewListener()
@@ -84,3 +85,21 @@ func TestWaitForReady(t *testing.T) {
 	require.True(ok)
 	require.Equal(codes.Unavailable, status.Code())
 }
+
+func TestWaitForReadyCallOption(t *testing.T) {
+	require := require.New(t)
+
+	listener, err := NewListener()
+	require.NoError(err)
+	conn, err := Dial(listener.Addr().String())
+	require.NoError(err)
+	// close listener causes RPC to fail fast.
+	_ = listener.Close()
+
+	db := pb.NewDatabaseClient(conn)
+	_, err = db.Put(context.Background(), &pb.PutRequest{Key: []byte("foo"), Value: []byte("bar")}, grpc.WaitForReady(false))
+	s, ok := status.FromError(err)
+	fmt.Printf("status: %v\n", s)
+	require.True(ok)
+	require.Equal(codes.Unavailable, s.Code())
+}
diff --git a/vms/rpcchainvm/state_syncable_vm_test.go b/vms/rpcchainvm/state_syncable_vm_test.go
@@ -25,8 +25,6 @@ import (
 	"github.com/ava-labs/avalanchego/vms/rpcchainvm/grpcutils"
 	"github.com/ava-labs/avalanchego/vms/rpcchainvm/runtime"
 	"github.com/ava-labs/avalanchego/vms/rpcchainvm/runtime/subprocess"
-
-	vmpb "github.com/ava-labs/avalanchego/proto/pb/vm"
 )
 
 var (
@@ -295,7 +293,7 @@ func buildClientHelper(require *require.Assertions, testKey string) (*VMClient,
 	clientConn, err := grpcutils.Dial(status.Addr)
 	require.NoError(err)
 
-	return NewClient(vmpb.NewVMClient(clientConn)), stopper
+	return NewClient(clientConn), stopper
 }
 
 func TestStateSyncEnabled(t *testing.T) {

diff --git a/vms/rpcchainvm/vm_client.go b/vms/rpcchainvm/vm_client.go
@@ -110,9 +110,10 @@ type VMClient struct {
 }
 
 // NewClient returns a VM connected to a remote VM
-func NewClient(client vmpb.VMClient) *VMClient {
+func NewClient(clientConn *grpc.ClientConn) *VMClient {
 	return &VMClient{
-		client: client,
+		client: vmpb.NewVMClient(clientConn),
+		conns:  []*grpc.ClientConn{clientConn},
 	}
 }
 
@@ -527,7 +528,9 @@ func (vm *VMClient) SetPreference(ctx context.Context, blkID ids.ID) error {
 }
 
 func (vm *VMClient) HealthCheck(ctx context.Context) (interface{}, error) {
-	health, err := vm.client.Health(ctx, &emptypb.Empty{})
+	// HealthCheck is a special case, where we want to fail fast instead of block.
+	failFast := grpc.WaitForReady(false)
+	health, err := vm.client.Health(ctx, &emptypb.Empty{}, failFast)
 	if err != nil {
 		return nil, fmt.Errorf("health check failed: %w", err)
 	}