refactor balancer server

pysel · Nov 29, 2023 · deb6aa5 · deb6aa5
1 parent 3b0cd7f
commit deb6aa5
Showing 1 changed file with 9 additions and 7 deletions.
diff --git a/balancer/balancer-server.go b/balancer/balancer-server.go
@@ -14,18 +14,20 @@ type BalancerServer struct {
 	*Balancer
 }
 
-func RunBalancerServer(port int64, partitions int) net.Addr {
+func RegisterBalancerServer(b *Balancer) *grpc.Server {
+	s := grpc.NewServer()
+	reflection.Register(s)
+	pbbalancer.RegisterBalancerServiceServer(s, &BalancerServer{Balancer: b})
+
+	return s
+}
+
+func listenOnPort(s *grpc.Server, port int64) net.Addr {
 	lis, err := net.Listen("tcp", fmt.Sprintf(":%d", port))
 	if err != nil {
 		panic(err)
 	}
 
-	balancer := NewBalancer(partitions)
-
-	s := grpc.NewServer()
-	reflection.Register(s)
-	pbbalancer.RegisterBalancerServiceServer(s, &BalancerServer{Balancer: balancer})
 	go s.Serve(lis)
-
 	return lis.Addr()
 }