From 81380e0862dbdd1001231cd0d08e04edd79bb6be Mon Sep 17 00:00:00 2001 From: Ying Li Date: Fri, 16 Oct 2015 09:46:08 -0700 Subject: [PATCH] Even simpler - cancel the GRPC call using the context object passed to the GRPC clients - thanks @endophage! Signed-off-by: Ying Li --- cmd/notary-server/main.go | 7 ++-- signer/signer_trust.go | 36 +++++++---------- signer/signer_trust_test.go | 81 ++++++++++++++++++++----------------- 3 files changed, 60 insertions(+), 64 deletions(-) diff --git a/cmd/notary-server/main.go b/cmd/notary-server/main.go index bae7cd8d68..d66f67bf29 100644 --- a/cmd/notary-server/main.go +++ b/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() diff --git a/signer/signer_trust.go b/signer/signer_trust.go index c7f8509eac..89b67fa404 100644 --- a/signer/signer_trust.go +++ b/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) - } - 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) + 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) } - return returnErr + return err } diff --git a/signer/signer_trust_test.go b/signer/signer_trust_test.go index 3776de91e4..c3e43b3763 100644 --- a/signer/signer_trust_test.go +++ b/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)) }