From aad56d6eb69581813ed19cd7ca7c8a83089a6809 Mon Sep 17 00:00:00 2001
From: Grant Griffiths <grant@portworx.com>
Date: Tue, 30 Apr 2019 21:27:59 -0700
Subject: [PATCH] Add support for SDK gRPC server extensions (#1094)

Signed-off-by: Grant Griffiths <grant@portworx.com>
---
 api/server/sdk/server.go | 21 +++++++++++++++++++++
 1 file changed, 21 insertions(+)

diff --git a/api/server/sdk/server.go b/api/server/sdk/server.go
index 6405bbc5c..e351269e6 100644
--- a/api/server/sdk/server.go
+++ b/api/server/sdk/server.go
@@ -107,6 +107,18 @@ type ServerConfig struct {
 	StoragePolicy policy.PolicyManager
 	// Security configuration
 	Security *SecurityConfig
+	// ServerExtensions allows you to extend the SDK gRPC server
+	// with callback functions that are sequentially executed
+	// at the end of Server.Start()
+	//
+	// To add your own service to the SDK gRPC server,
+	// just append a function callback that registers it:
+	//
+	// s.config.ServerExtensions = append(s.config.ServerExtensions,
+	// 		func(gs *grpc.Server) {
+	//			api.RegisterCustomService(gs, customHandler)
+	//		})
+	ServerExtensions []func(s *grpc.Server)
 }
 
 // Server is an implementation of the gRPC SDK interface
@@ -451,6 +463,9 @@ func (s *sdkGrpcServer) Start() error {
 		if s.config.Security.Role != nil {
 			api.RegisterOpenStorageRoleServer(grpcServer, s.roleServer)
 		}
+
+		s.RegisterServerExtensions(grpcServer)
+
 		return grpcServer
 	})
 	if err != nil {
@@ -460,6 +475,12 @@ func (s *sdkGrpcServer) Start() error {
 	return nil
 }
 
+func (s *sdkGrpcServer) RegisterServerExtensions(grpcServer *grpc.Server) {
+	for _, ext := range s.config.ServerExtensions {
+		ext(grpcServer)
+	}
+}
+
 func (s *sdkGrpcServer) useCluster(c cluster.Cluster) {
 	s.lock.Lock()
 	defer s.lock.Unlock()