Refactor rpc package

sanjayarun07 · Mar 29, 2019 · 219a94a · 219a94a
1 parent 519f53d
commit 219a94a
Show file tree

Hide file tree

Showing 29 changed files with 976 additions and 1,611 deletions.
diff --git a/' b/'
diff --git a/client/helper_test.go b/client/helper_test.go
@@ -34,6 +34,7 @@ import (
 	"github.com/CovenantSQL/CovenantSQL/crypto/kms"
 	"github.com/CovenantSQL/CovenantSQL/proto"
 	"github.com/CovenantSQL/CovenantSQL/route"
+	rrpc "github.com/CovenantSQL/CovenantSQL/rpc"
 	rpc "github.com/CovenantSQL/CovenantSQL/rpc/mux"
 	"github.com/CovenantSQL/CovenantSQL/types"
 	"github.com/CovenantSQL/CovenantSQL/utils"
@@ -233,7 +234,7 @@ func initNode() (cleanupFunc func(), tempDir string, server *rpc.Server, err err
 	}
 
 	// init rpc
-	if server, err = rpc.NewServerWithService(rpc.ServiceMap{route.DHTRPCName: dht}); err != nil {
+	if server, err = rpc.NewServerWithService(rrpc.ServiceMap{route.DHTRPCName: dht}); err != nil {
 		return
 	}
 

diff --git a/cmd/cqld/bench_test.go b/cmd/cqld/bench_test.go
@@ -33,6 +33,7 @@ import (
 	"github.com/CovenantSQL/CovenantSQL/crypto/kms"
 	"github.com/CovenantSQL/CovenantSQL/proto"
 	"github.com/CovenantSQL/CovenantSQL/route"
+	rrpc "github.com/CovenantSQL/CovenantSQL/rpc"
 	rpc "github.com/CovenantSQL/CovenantSQL/rpc/mux"
 	"github.com/CovenantSQL/CovenantSQL/utils"
 	"github.com/CovenantSQL/CovenantSQL/utils/log"
@@ -163,14 +164,11 @@ func TestStartBP_CallRPC(t *testing.T) {
 
 	leaderNodeID := kms.BP.NodeID
 	var conn net.Conn
-	var RPCClient *rpc.Client
 
-	if conn, err = rpc.DialToNodeWithPool(rpc.GetSessionPoolInstance(), leaderNodeID, false); err != nil {
-		t.Fatal(err)
-	}
-	if RPCClient, err = rpc.InitClientConn(conn); err != nil {
+	if conn, err = rrpc.DialToNodeWithPool(rpc.GetSessionPoolInstance(), leaderNodeID, false); err != nil {
 		t.Fatal(err)
 	}
+	RPCClient := rrpc.NewClientWithConn(conn)
 
 	nodePayload := proto.NewNode()
 	nodePayload.InitNodeCryptoInfo(100 * time.Millisecond)
@@ -294,14 +292,11 @@ func BenchmarkKVServer_GetAllNodeInfo(b *testing.B) {
 
 	leaderNodeID := kms.BP.NodeID
 	var conn net.Conn
-	var RPCClient *rpc.Client
 
-	if conn, err = rpc.DialToNodeWithPool(rpc.GetSessionPoolInstance(), leaderNodeID, false); err != nil {
-		return
-	}
-	if RPCClient, err = rpc.InitClientConn(conn); err != nil {
+	if conn, err = rrpc.DialToNodeWithPool(rpc.GetSessionPoolInstance(), leaderNodeID, false); err != nil {
 		return
 	}
+	RPCClient := rrpc.NewClientWithConn(conn)
 
 	var reqType = "FindNeighbor"
 	nodePayload := proto.NewNode()