refactoring

1 files changed, 37 insertions, 0 deletions

diff --git a/weed/util/grpc_client_server.go b/weed/util/grpc_client_server.go
index 8dbb4c0cd..18d5c02c9 100644
--- a/weed/util/grpc_client_server.go
+++ b/weed/util/grpc_client_server.go
@@ -1,12 +1,20 @@
 package util
 
 import (
+	"fmt"
+	"sync"
 	"time"
 
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/keepalive"
 )
 
+var (
+	// cache grpc connections
+	grpcClients     = make(map[string]*grpc.ClientConn)
+	grpcClientsLock sync.Mutex
+)
+
 func NewGrpcServer() *grpc.Server {
 	return grpc.NewServer(grpc.KeepaliveParams(keepalive.ServerParameters{
 		Time:    10 * time.Second, // wait time before ping if no activity
@@ -26,3 +34,32 @@ func GrpcDial(address string, opts ...grpc.DialOption) (*grpc.ClientConn, error)
 
 	return grpc.Dial(address, opts...)
 }
+
+func WithCachedGrpcClient(fn func(*grpc.ClientConn) error, address string, opts ...grpc.DialOption) error {
+
+	grpcClientsLock.Lock()
+
+	existingConnection, found := grpcClients[address]
+	if found {
+		grpcClientsLock.Unlock()
+		return fn(existingConnection)
+	}
+
+	grpcConnection, err := GrpcDial(address, opts...)
+	if err != nil {
+		grpcClientsLock.Unlock()
+		return fmt.Errorf("fail to dial %s: %v", address, err)
+	}
+
+	grpcClients[address] = grpcConnection
+	grpcClientsLock.Unlock()
+
+	err = fn(grpcConnection)
+	if err != nil {
+		grpcClientsLock.Lock()
+		delete(grpcClients, address)
+		grpcClientsLock.Unlock()
+	}
+
+	return err
+}