Allow Shutdown and HealthCheck RPCs to fail-fast

Signed-off-by: Sam Batschelet <[email protected]>

vms/rpcchainvm/factory.go

@@ -62,6 +62,7 @@ func (f *factory) New(log logging.Logger) (interface{}, error) {
 	}
 
 	vm := NewClient(vmpb.NewVMClient(clientConn))
+	vm.conns = append(vm.conns, clientConn)
 	vm.SetProcess(stopper, status.Pid, f.processTracker)
 
 	f.runtimeTracker.TrackRuntime(stopper)

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())
+}

vms/rpcchainvm/vm_client.go

@@ -353,7 +353,10 @@ func (vm *VMClient) SetState(ctx context.Context, state snow.State) error {
 
 func (vm *VMClient) Shutdown(ctx context.Context) error {
 	errs := wrappers.Errs{}
-	_, err := vm.client.Shutdown(ctx, &emptypb.Empty{})
+	// Shutdown is a special case, where we want to fail fast instead of block
+	// here.
+	failFast := grpc.WaitForReady(false)
+	_, err := vm.client.Shutdown(ctx, &emptypb.Empty{}, failFast)
 	errs.Add(err)
 
 	vm.serverCloser.Stop()
@@ -527,7 +530,10 @@ 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
+	// here.
+	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)
 	}