From 095407df5828ba8da963ad52e960f440ec1bf1e1 Mon Sep 17 00:00:00 2001
From: Gyu-Ho Lee <gyuhox@gmail.com>
Date: Mon, 27 Feb 2017 17:09:36 -0800
Subject: [PATCH] etcdmain: add register,resolver flags

Signed-off-by: Gyu-Ho Lee <gyuhox@gmail.com>
---
 etcdmain/grpc_proxy.go | 39 ++++++++++++++++++++++++++++++---------
 1 file changed, 30 insertions(+), 9 deletions(-)

diff --git a/etcdmain/grpc_proxy.go b/etcdmain/grpc_proxy.go
index b92825693..08312e86b 100644
--- a/etcdmain/grpc_proxy.go
+++ b/etcdmain/grpc_proxy.go
@@ -27,20 +27,22 @@ import (
 	"github.com/coreos/etcd/pkg/transport"
 	"github.com/coreos/etcd/proxy/grpcproxy"
 
-	"github.com/spf13/cobra"
-	"google.golang.org/grpc"
-
 	"github.com/cockroachdb/cmux"
 	grpc_prometheus "github.com/grpc-ecosystem/go-grpc-prometheus"
 	"github.com/prometheus/client_golang/prometheus"
+	"github.com/spf13/cobra"
+	"google.golang.org/grpc"
 )
 
 var (
-	grpcProxyListenAddr string
-	grpcProxyEndpoints  []string
-	grpcProxyCert       string
-	grpcProxyKey        string
-	grpcProxyCA         string
+	grpcProxyListenAddr         string
+	grpcProxyEndpoints          []string
+	grpcProxyCert               string
+	grpcProxyKey                string
+	grpcProxyCA                 string
+	grpcProxyAdvertiseClientURL string
+	grpcProxyResolverPrefix     string
+	grpcProxyResolverTTL        int
 )
 
 func init() {
@@ -70,11 +72,27 @@ func newGRPCProxyStartCommand() *cobra.Command {
 	cmd.Flags().StringVar(&grpcProxyCert, "cert", "", "identify secure connections with etcd servers using this TLS certificate file")
 	cmd.Flags().StringVar(&grpcProxyKey, "key", "", "identify secure connections with etcd servers using this TLS key file")
 	cmd.Flags().StringVar(&grpcProxyCA, "cacert", "", "verify certificates of TLS-enabled secure etcd servers using this CA bundle")
+	cmd.Flags().StringVar(&grpcProxyAdvertiseClientURL, "advertise-client-url", "127.0.0.1:23790", "advertise address to register (must be reachable by client)")
+	cmd.Flags().StringVar(&grpcProxyResolverPrefix, "resolver-prefix", "", "prefix to use for registering proxy (must be shared with other grpc-proxy members)")
+	cmd.Flags().IntVar(&grpcProxyResolverTTL, "resolver-ttl", 0, "specify TTL, in seconds, when registering proxy endpoints")
 
 	return &cmd
 }
 
 func startGRPCProxy(cmd *cobra.Command, args []string) {
+	if grpcProxyResolverPrefix != "" && grpcProxyResolverTTL < 1 {
+		fmt.Fprintln(os.Stderr, fmt.Errorf("invalid resolver-ttl %d", grpcProxyResolverTTL))
+		os.Exit(1)
+	}
+	if grpcProxyResolverPrefix == "" && grpcProxyResolverTTL > 0 {
+		fmt.Fprintln(os.Stderr, fmt.Errorf("invalid resolver-prefix %q", grpcProxyResolverPrefix))
+		os.Exit(1)
+	}
+	if grpcProxyResolverPrefix != "" && grpcProxyResolverTTL > 0 && grpcProxyAdvertiseClientURL == "" {
+		fmt.Fprintln(os.Stderr, fmt.Errorf("invalid advertise-client-url %q", grpcProxyAdvertiseClientURL))
+		os.Exit(1)
+	}
+
 	l, err := net.Listen("tcp", grpcProxyListenAddr)
 	if err != nil {
 		fmt.Fprintln(os.Stderr, err)
@@ -105,7 +123,10 @@ func startGRPCProxy(cmd *cobra.Command, args []string) {
 
 	kvp, _ := grpcproxy.NewKvProxy(client)
 	watchp, _ := grpcproxy.NewWatchProxy(client)
-	clusterp := grpcproxy.NewClusterProxy(client)
+	if grpcProxyResolverPrefix != "" {
+		grpcproxy.Register(client, grpcProxyResolverPrefix, grpcProxyAdvertiseClientURL, grpcProxyResolverTTL)
+	}
+	clusterp, _ := grpcproxy.NewClusterProxy(client, grpcProxyAdvertiseClientURL, grpcProxyResolverPrefix)
 	leasep, _ := grpcproxy.NewLeaseProxy(client)
 	mainp := grpcproxy.NewMaintenanceProxy(client)
 	authp := grpcproxy.NewAuthProxy(client)