added grpc interceptor and singleton changes

pkg/client/client.go

@@ -4,6 +4,7 @@
 package client
 
 import (
+	"context"
 	"fmt"
 	"os"
 	"strings"
@@ -15,7 +16,9 @@ import (
 	"google.golang.org/grpc/keepalive"
 	log "k8s.io/klog"
 
+	"github.com/microsoft/moc-sdk-for-go/pkg/constant"
 	"github.com/microsoft/moc/pkg/auth"
+	"github.com/microsoft/moc/pkg/errors"
 )
 
 const (
@@ -29,6 +32,20 @@ var (
 	connectionCache map[string]*grpc.ClientConn
 )
 
+func clientConnOptionsInterceptor() grpc.UnaryClientInterceptor {
+	return func(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
+		err := invoker(ctx, method, req, reply, cc, opts...)
+		if err != nil {
+			grpcConnFailure := errors.IsGRPCUnavailable(err) || errors.IsGRPCDeadlineExceeded(err)
+			exitProcess := !constant.GetClientOpts().NoExitOnConnFailure
+			if grpcConnFailure && exitProcess {
+				log.Fatalf("Communication with cloud agent failed. Exiting Process with error: %+v\n", err)
+			}
+		}
+		return err
+	}
+}
+
 func init() {
 	connectionCache = map[string]*grpc.ClientConn{}
 }
@@ -107,6 +124,7 @@ func getClientConnection(serverAddress *string, authorizer auth.Authorizer) (*gr
 	}
 
 	opts := getDefaultDialOption(authorizer)
+	opts = append(opts, grpc.WithUnaryInterceptor(clientConnOptionsInterceptor()))
 	conn, err := grpc.Dial(endpoint, opts...)
 	if err != nil {
 		log.Fatalf("Failed to dial: %v", err)

pkg/constant/constant.go

@@ -4,6 +4,7 @@
 package constant
 
 import (
+	"sync"
 	"time"
 )
 
@@ -12,3 +13,24 @@ const (
 	CertificateValidityThreshold float64 = (30.0 / 100.0)
 	RenewalBackoff               float64 = (2.0 / 100.0)
 )
+
+type ClientOpts struct {
+	NoExitOnConnFailure bool
+}
+
+var clientOptsSingleton *ClientOpts
+var once sync.Once
+
+//If GetClientOpts called before SetClientOpts, it will return a default ClientOpts
+func GetClientOpts() *ClientOpts {
+	once.Do(func() {
+		clientOptsSingleton = &ClientOpts{}
+	})
+	return clientOptsSingleton
+}
+
+func SetClientOpts(clientOpts *ClientOpts) {
+	once.Do(func() {
+		clientOptsSingleton = clientOpts
+	})
+}