Making an RPC call would hang forever if the connection to the GRPC server dies

(it just retries to connect, and the RPC call would just block until it does), so change up
how this works by timing out the health check.

Signed-off-by: Ying Li <ying.li@docker.com>

cmd/notary-server/main.go

@@ -114,10 +114,10 @@ func main() {
 			viper.GetString("trust_service.tls_ca_file"),
 		)
 		health.RegisterPeriodicFunc(
-			"Key manager operational", trust.(*signer.NotarySigner).KMHealth,
-			time.Second*60)
-		health.RegisterPeriodicFunc(
-			"Signer operational", trust.(*signer.NotarySigner).SHealth,
+			"Trust operational",
+			func() error {
+				return trust.(*signer.NotarySigner).CheckHealth(60)
+			},
 			time.Second*60)
 	} else {
 		logrus.Info("Using local signing service")

signer/signer_trust.go

@@ -2,7 +2,9 @@ package signer
 
 import (
 	"errors"
+	"fmt"
 	"net"
+	"time"
 
 	"github.com/Sirupsen/logrus"
 	pb "github.com/docker/notary/proto"
@@ -12,10 +14,15 @@ import (
 	"google.golang.org/grpc/credentials"
 )
 
+type checkableConnectionState interface {
+	State() grpc.ConnectivityState
+}
+
 // NotarySigner implements a RPC based Trust service that calls the Notary-signer Service
 type NotarySigner struct {
-	kmClient pb.KeyManagementClient
-	sClient  pb.SignerClient
+	kmClient   pb.KeyManagementClient
+	sClient    pb.SignerClient
+	clientConn checkableConnectionState
 }
 
 // NewNotarySigner is a convinience method that returns NotarySigner
@@ -35,8 +42,9 @@ func NewNotarySigner(hostname string, port string, tlscafile string) *NotarySign
 	kmClient := pb.NewKeyManagementClient(conn)
 	sClient := pb.NewSignerClient(conn)
 	return &NotarySigner{
-		kmClient: kmClient,
-		sClient:  sClient,
+		kmClient:   kmClient,
+		sClient:    sClient,
+		clientConn: conn,
 	}
 }
 
@@ -87,26 +95,55 @@ func (trust *NotarySigner) GetKey(keyid string) data.PublicKey {
 	return data.NewPublicKey(data.KeyAlgorithm(publicKey.KeyInfo.Algorithm.Algorithm), publicKey.PublicKey)
 }
 
-// KMHealth returns the health of the KeyManagement RPC service.
-func (trust *NotarySigner) KMHealth() error {
-	status, err := trust.kmClient.CheckHealth(context.Background(), &pb.Void{})
-	if err != nil {
-		return err
+// CheckHealth returns true if trust is healthy - that is if it can connect
+// to the trust server, and both the key management and signer services are
+// healthy
+func (trust *NotarySigner) CheckHealth(timeout int) error {
+	if e := trust.checkServiceHealth("Key manager", trust.kmClient,
+		timeout); e != nil {
+		return e
 	}
-	if len(status.Status) > 0 {
-		return errors.New("KeyManagement not healthy")
-	}
-	return nil
+	return trust.checkServiceHealth("Signer", trust.sClient, timeout)
 }
 
-// SHealth returns the health of the Signer RPC service.
-func (trust *NotarySigner) SHealth() error {
-	status, err := trust.sClient.CheckHealth(context.Background(), &pb.Void{})
-	if err != nil {
+// Generalized function that can check the health of both the signer client
+// and the key management client.
+func (trust *NotarySigner) checkServiceHealth(
+	serviceName string,
+	client interface {
+		CheckHealth(context.Context, *pb.Void, ...grpc.CallOption) (*pb.HealthStatus, error)
+	},
+	timeout int) error {
+
+	// Do not bother starting goroutine if the connection is broken.
+	if trust.clientConn.State() != grpc.Idle &&
+		trust.clientConn.State() != grpc.Ready {
+		return errors.New("Not currently connected to trust server.")
+	}
+
+	// We still want to time out getting health, because the connection could
+	// have disconnected sometime between when we checked the connection and
+	// when we try to make an RPC call.
+	channel := make(chan error)
+	go func() {
+		status, err := client.CheckHealth(context.Background(), &pb.Void{})
+		// if this function gets timed out, it might panic when writing to a
+		// closed channel
+		defer func() { recover() }()
+		if err != nil {
+			channel <- err
+		}
+		if len(status.Status) > 0 {
+			channel <- fmt.Errorf("%s not healthy", serviceName)
+		}
+		channel <- nil
+	}()
+	select {
+	case err := <-channel:
+		close(channel)
 		return err
+	case <-time.After(time.Second * time.Duration(timeout)):
+		close(channel)
+		return fmt.Errorf("Timed out connecting to %s after 60 seconds", serviceName)
 	}
-	if len(status.Status) > 0 {
-		return errors.New("Signer not healthy")
-	}
-	return nil
 }

signer/signer_trust_test.go

@@ -3,6 +3,7 @@ package signer
 import (
 	"errors"
 	"testing"
+	"time"
 
 	"google.golang.org/grpc"
 
@@ -13,110 +14,116 @@ import (
 
 type StubKeyManagementClient struct {
 	pb.KeyManagementClient
-	status map[string]string
-	err    error
+	healthCheck func() (map[string]string, error)
 }
 
 func (c StubKeyManagementClient) CheckHealth(x context.Context, v *pb.Void, o ...grpc.CallOption) (*pb.HealthStatus, error) {
-	if c.err != nil {
-		return nil, c.err
+	status, err := c.healthCheck()
+	if err != nil {
+		return nil, err
 	}
-	return &pb.HealthStatus{c.status}, nil
+	return &pb.HealthStatus{status}, nil
 }
 
 type StubSignerClient struct {
 	pb.SignerClient
-	status map[string]string
-	err    error
+	healthCheck func() (map[string]string, error)
 }
 
 func (c StubSignerClient) CheckHealth(x context.Context, v *pb.Void, o ...grpc.CallOption) (*pb.HealthStatus, error) {
-	if c.err != nil {
-		return nil, c.err
+	status, err := c.healthCheck()
+	if err != nil {
+		return nil, err
 	}
-	return &pb.HealthStatus{c.status}, nil
+	return &pb.HealthStatus{status}, nil
 }
 
-// KMHealth only succeeds if the KeyManagement service is healthy.
-func TestKMHealthFailure(t *testing.T) {
-	errs := []error{errors.New("Connection failure!"), nil}
-	status := map[string]string{"db": "bad"}
-
-	for _, err := range errs {
-		signer := NotarySigner{
-			StubKeyManagementClient{
-				pb.NewKeyManagementClient(nil),
-				status,
-				err,
-			},
-			StubSignerClient{
-				pb.NewSignerClient(nil),
-				make(map[string]string),
-				nil,
-			},
-		}
-		err := signer.KMHealth()
-		assert.Error(t, err)
-	}
+type StubGRPCConnection struct {
+	fakeConnStatus grpc.ConnectivityState
 }
 
-// SHealth only succeeds if the Signer service is healthy.
-func TestSignerHealthFailure(t *testing.T) {
-	errs := []error{errors.New("Connection failure!"), nil}
-	status := map[string]string{"db": "bad"}
-
-	for _, err := range errs {
-		signer := NotarySigner{
-			StubKeyManagementClient{
-				pb.NewKeyManagementClient(nil),
-				make(map[string]string),
-				nil,
-			},
-			StubSignerClient{
-				pb.NewSignerClient(nil),
-				status,
-				err,
-			},
-		}
-		err := signer.SHealth()
-		assert.Error(t, err)
-	}
+func (c StubGRPCConnection) State() grpc.ConnectivityState {
+	return c.fakeConnStatus
 }
 
-// SHealth succeeds if the signer service is healthy, regardless of the
-// key management service.
-func TestSignerHealthGood(t *testing.T) {
-	signer := NotarySigner{
+type healthSideEffect func() (map[string]string, error)
+
+func healthOk() (map[string]string, error) {
+	return make(map[string]string), nil
+}
+
+func healthBad() (map[string]string, error) {
+	return map[string]string{"health": "not good"}, nil
+}
+
+func healthError() (map[string]string, error) {
+	return nil, errors.New("Something's wrong")
+}
+
+func healthTimeout() (map[string]string, error) {
+	time.Sleep(time.Second * 10)
+	return healthOk()
+}
+
+func makeSigner(kmFunc healthSideEffect, sFunc healthSideEffect, conn StubGRPCConnection) NotarySigner {
+	return NotarySigner{
 		StubKeyManagementClient{
 			pb.NewKeyManagementClient(nil),
-			nil,
-			errors.New("Connection failure!"),
+			kmFunc,
 		},
 		StubSignerClient{
 			pb.NewSignerClient(nil),
-			make(map[string]string),
-			nil,
+			sFunc,
 		},
+		conn,
 	}
-	err := signer.SHealth()
-	assert.NoError(t, err)
 }
 
-// KMHealth succeeds if the key management service is healthy, regardless of
-// the signer service.
-func TestKMHealthGood(t *testing.T) {
-	signer := NotarySigner{
-		StubKeyManagementClient{
-			pb.NewKeyManagementClient(nil),
-			make(map[string]string),
-			nil,
-		},
-		StubSignerClient{
-			pb.NewSignerClient(nil),
-			nil,
-			errors.New("Connection failure!"),
-		},
-	}
-	err := signer.KMHealth()
-	assert.NoError(t, err)
+// CheckHealth does not succeed if the KM server is unhealthy
+func TestHealthCheckKMUnhealthy(t *testing.T) {
+	signer := makeSigner(healthBad, healthOk, StubGRPCConnection{})
+	assert.Error(t, signer.CheckHealth(1))
+}
+
+// CheckHealth does not succeed if the health check to the KM server errors
+func TestHealthCheckKMError(t *testing.T) {
+	signer := makeSigner(healthBad, healthOk, StubGRPCConnection{})
+	assert.Error(t, signer.CheckHealth(1))
+}
+
+// CheckHealth does not succeed if the health check to the KM server times out
+func TestHealthCheckKMTimeout(t *testing.T) {
+	signer := makeSigner(healthTimeout, healthOk, StubGRPCConnection{})
+	assert.Error(t, signer.CheckHealth(1))
+}
+
+// CheckHealth does not succeed if the signer is unhealthy
+func TestHealthCheckSignerUnhealthy(t *testing.T) {
+	signer := makeSigner(healthOk, healthBad, StubGRPCConnection{})
+	assert.Error(t, signer.CheckHealth(1))
+}
+
+// CheckHealth does not succeed if the health check to the signer errors
+func TestHealthCheckSignerError(t *testing.T) {
+	signer := makeSigner(healthOk, healthBad, StubGRPCConnection{})
+	assert.Error(t, signer.CheckHealth(1))
+}
+
+// CheckHealth does not succeed if the health check to the signer times out
+func TestHealthCheckSignerTimeout(t *testing.T) {
+	signer := makeSigner(healthOk, healthTimeout, StubGRPCConnection{})
+	assert.Error(t, signer.CheckHealth(1))
+}
+
+// CheckHealth succeeds if both services are healthy and reachable
+func TestHealthCheckBothHealthy(t *testing.T) {
+	signer := makeSigner(healthOk, healthOk, StubGRPCConnection{})
+	assert.NoError(t, signer.CheckHealth(1))
+}
+
+// CheckHealth fails immediately if not connected to the server.
+func TestHealthCheckConnectionDied(t *testing.T) {
+	signer := makeSigner(healthTimeout, healthTimeout,
+		StubGRPCConnection{grpc.Connecting})
+	assert.Error(t, signer.CheckHealth(30))
 }