Even simpler - cancel the GRPC call using the context object passed

to the GRPC clients - thanks @endophage!

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

cmd/notary-server/main.go

@@ -113,21 +113,20 @@ func main() {
 			viper.GetString("trust_service.port"),
 			viper.GetString("trust_service.tls_ca_file"),
 		)
+		minute := 1 * time.Minute
 		health.RegisterPeriodicFunc(
 			"Trust operational",
 			// If the trust service fails, the server is degraded but not
 			// exactly unheatlthy, so always return healthy and just log an
 			// error.
 			func() error {
-				logrus.Info("Getting health")
-				err := trust.(*signer.NotarySigner).CheckHealth(5)
-				logrus.Info("Got health")
+				err := trust.(*signer.NotarySigner).CheckHealth(minute)
 				if err != nil {
 					logrus.Error("Trust not fully operational: ", err.Error())
 				}
 				return nil
 			},
-			time.Second*5)
+			minute)
 	} else {
 		logrus.Info("Using local signing service")
 		trust = signed.NewEd25519()

signer/signer_trust.go

@@ -11,6 +11,7 @@ import (
 	"github.com/endophage/gotuf/data"
 	"golang.org/x/net/context"
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/credentials"
 )
 
@@ -99,9 +100,9 @@ func (trust *NotarySigner) GetKey(keyid string) data.PublicKey {
 // 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 {
+func (trust *NotarySigner) CheckHealth(timeout time.Duration) error {
 	if e := trust.checkServiceHealth(
-		"Key manager", trust.kmClient.CheckHealth, timeout); e != nil {
+		"Key Manager", trust.kmClient.CheckHealth, timeout); e != nil {
 		return e
 	}
 	return trust.checkServiceHealth(
@@ -114,31 +115,22 @@ type rpcHealthCheck func(
 // Generalized function that can check the health of both the signer client
 // and the key management client.
 func (trust *NotarySigner) checkServiceHealth(
-	serviceName string, check rpcHealthCheck, timeout int) error {
+	serviceName string, check rpcHealthCheck, timeout time.Duration) error {
 
-	// Do not bother starting goroutine if the connection is broken.
+	// Do not bother starting checking at all 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, 1)
-	go func() {
-		status, err := check(context.Background(), &pb.Void{})
-		if len(status.Status) > 0 {
-			err = fmt.Errorf("%s not healthy", serviceName)
+	ctx, cancel := context.WithTimeout(context.Background(), timeout)
+	status, err := check(ctx, &pb.Void{})
+	defer cancel()
+	if err == nil && len(status.Status) > 0 {
+		return fmt.Errorf("%s not healthy", serviceName)
+	} else if err != nil && grpc.Code(err) == codes.DeadlineExceeded {
+		return fmt.Errorf("Timed out reaching %s after %s.", serviceName,
+			timeout)
 	}
-		channel <- err
-	}()
-	var returnErr error
-	select {
-	case err := <-channel:
-		returnErr = err
-	case <-time.After(time.Second * time.Duration(timeout)):
-		returnErr = fmt.Errorf("Timed out connecting to %s after %s", serviceName, timeout)
-	}
-	return returnErr
+	return err
 }

signer/signer_trust_test.go

@@ -2,10 +2,12 @@ package signer
 
 import (
 	"errors"
+	"strings"
 	"testing"
 	"time"
 
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
 
 	pb "github.com/docker/notary/proto"
 	"github.com/stretchr/testify/assert"
@@ -14,30 +16,22 @@ import (
 
 type StubKeyManagementClient struct {
 	pb.KeyManagementClient
-	healthCheck func() (map[string]string, error)
+	healthCheck rpcHealthCheck
 }
 
 func (c StubKeyManagementClient) CheckHealth(x context.Context,
 	v *pb.Void, o ...grpc.CallOption) (*pb.HealthStatus, error) {
-	status, err := c.healthCheck()
-	if err != nil {
-		return nil, err
-	}
-	return &pb.HealthStatus{Status: status}, nil
+	return c.healthCheck(x, v, o...)
 }
 
 type StubSignerClient struct {
 	pb.SignerClient
-	healthCheck func() (map[string]string, error)
+	healthCheck rpcHealthCheck
 }
 
 func (c StubSignerClient) CheckHealth(x context.Context, v *pb.Void,
 	o ...grpc.CallOption) (*pb.HealthStatus, error) {
-	status, err := c.healthCheck()
-	if err != nil {
-		return nil, err
-	}
-	return &pb.HealthStatus{Status: status}, nil
+	return c.healthCheck(x, v, o...)
 }
 
 type StubGRPCConnection struct {
@@ -48,26 +42,33 @@ func (c StubGRPCConnection) State() grpc.ConnectivityState {
 	return c.fakeConnStatus
 }
 
-type healthSideEffect func() (map[string]string, error)
+func stubHealthFunction(t *testing.T, status map[string]string, err error) rpcHealthCheck {
+	return func(ctx context.Context, v *pb.Void, o ...grpc.CallOption) (*pb.HealthStatus, error) {
+		_, withDeadline := ctx.Deadline()
+		assert.True(t, withDeadline)
 
-func healthOk() (map[string]string, error) {
-	return make(map[string]string), nil
+		return &pb.HealthStatus{status}, err
+	}
 }
 
-func healthBad() (map[string]string, error) {
-	return map[string]string{"health": "not good"}, nil
+func healthOk(t *testing.T) rpcHealthCheck {
+	return stubHealthFunction(t, make(map[string]string), nil)
 }
 
-func healthError() (map[string]string, error) {
-	return nil, errors.New("Something's wrong")
+func healthBad(t *testing.T) rpcHealthCheck {
+	return stubHealthFunction(t, map[string]string{"health": "not good"}, nil)
 }
 
-func healthTimeout() (map[string]string, error) {
-	time.Sleep(time.Second * 10)
-	return healthOk()
+func healthError(t *testing.T) rpcHealthCheck {
+	return stubHealthFunction(t, nil, errors.New("Something's wrong"))
 }
 
-func makeSigner(kmFunc healthSideEffect, sFunc healthSideEffect, conn StubGRPCConnection) NotarySigner {
+func healthTimeout(t *testing.T) rpcHealthCheck {
+	return stubHealthFunction(
+		t, nil, grpc.Errorf(codes.DeadlineExceeded, ""))
+}
+
+func makeSigner(kmFunc rpcHealthCheck, sFunc rpcHealthCheck, conn StubGRPCConnection) NotarySigner {
 	return NotarySigner{
 		StubKeyManagementClient{
 			pb.NewKeyManagementClient(nil),
@@ -83,49 +84,53 @@ func makeSigner(kmFunc healthSideEffect, sFunc healthSideEffect, conn StubGRPCCo
 
 // 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))
+	signer := makeSigner(healthBad(t), healthOk(t), StubGRPCConnection{})
+	assert.Error(t, signer.CheckHealth(1*time.Second))
 }
 
 // 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))
+	signer := makeSigner(healthBad(t), healthOk(t), StubGRPCConnection{})
+	assert.Error(t, signer.CheckHealth(1*time.Second))
 }
 
 // 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))
+	signer := makeSigner(healthTimeout(t), healthOk(t), StubGRPCConnection{})
+	err := signer.CheckHealth(1 * time.Second)
+	assert.Error(t, err)
+	assert.True(t, strings.Contains(err.Error(), "Timed out"))
 }
 
 // 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))
+	signer := makeSigner(healthOk(t), healthBad(t), StubGRPCConnection{})
+	assert.Error(t, signer.CheckHealth(1*time.Second))
 }
 
 // CheckHealth does not succeed if the health check to the signer errors
 func TestHealthCheckSignerError(t *testing.T) {
-	signer := makeSigner(healthOk, healthBad, StubGRPCConnection{})
+	signer := makeSigner(healthOk(t), healthBad(t), 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))
+	signer := makeSigner(healthOk(t), healthTimeout(t), StubGRPCConnection{})
+	err := signer.CheckHealth(1 * time.Second)
+	assert.Error(t, err)
+	assert.True(t, strings.Contains(err.Error(), "Timed out"))
 }
 
 // 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))
+	signer := makeSigner(healthOk(t), healthOk(t), StubGRPCConnection{})
+	assert.NoError(t, signer.CheckHealth(1*time.Second))
 }
 
 // CheckHealth fails immediately if not connected to the server.
 func TestHealthCheckConnectionDied(t *testing.T) {
-	signer := makeSigner(healthTimeout, healthTimeout,
+	signer := makeSigner(healthOk(t), healthOk(t),
 		StubGRPCConnection{grpc.Connecting})
-	assert.Error(t, signer.CheckHealth(30))
+	assert.Error(t, signer.CheckHealth(1*time.Second))
 }