From 3f74e9db0dac0b113557fa3f9a5b0ee5cf8711c5 Mon Sep 17 00:00:00 2001
From: Xiang Li <xiangli.cs@gmail.com>
Date: Sat, 16 Jul 2016 10:00:14 -0700
Subject: [PATCH] *: add lease proxy

---
 etcdmain/grpc_proxy.go   |  2 +
 proxy/grpcproxy/lease.go | 82 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 84 insertions(+)
 create mode 100644 proxy/grpcproxy/lease.go

diff --git a/etcdmain/grpc_proxy.go b/etcdmain/grpc_proxy.go
index 279bf0003..43c487707 100644
--- a/etcdmain/grpc_proxy.go
+++ b/etcdmain/grpc_proxy.go
@@ -90,11 +90,13 @@ func startGRPCProxy(cmd *cobra.Command, args []string) {
 	kvp := grpcproxy.NewKvProxy(client)
 	watchp := grpcproxy.NewWatchProxy(client)
 	clusterp := grpcproxy.NewClusterProxy(client)
+	leasep := grpcproxy.NewLeaseProxy(client)
 
 	server := grpc.NewServer()
 	pb.RegisterKVServer(server, kvp)
 	pb.RegisterWatchServer(server, watchp)
 	pb.RegisterClusterServer(server, clusterp)
+	pb.RegisterLeaseServer(server, leasep)
 
 	server.Serve(l)
 }
diff --git a/proxy/grpcproxy/lease.go b/proxy/grpcproxy/lease.go
new file mode 100644
index 000000000..bb4a9a1a3
--- /dev/null
+++ b/proxy/grpcproxy/lease.go
@@ -0,0 +1,82 @@
+// Copyright 2016 The etcd Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package grpcproxy
+
+import (
+	"golang.org/x/net/context"
+
+	"github.com/coreos/etcd/clientv3"
+	pb "github.com/coreos/etcd/etcdserver/etcdserverpb"
+)
+
+type leaseProxy struct {
+	client *clientv3.Client
+}
+
+func NewLeaseProxy(c *clientv3.Client) *leaseProxy {
+	return &leaseProxy{
+		client: c,
+	}
+}
+
+func (lp *leaseProxy) LeaseGrant(ctx context.Context, cr *pb.LeaseGrantRequest) (*pb.LeaseGrantResponse, error) {
+	conn := lp.client.ActiveConnection()
+	return pb.NewLeaseClient(conn).LeaseGrant(ctx, cr)
+}
+
+func (lp *leaseProxy) LeaseRevoke(ctx context.Context, rr *pb.LeaseRevokeRequest) (*pb.LeaseRevokeResponse, error) {
+	conn := lp.client.ActiveConnection()
+	return pb.NewLeaseClient(conn).LeaseRevoke(ctx, rr)
+}
+
+func (lp *leaseProxy) LeaseKeepAlive(stream pb.Lease_LeaseKeepAliveServer) error {
+	conn := lp.client.ActiveConnection()
+	ctx, cancel := context.WithCancel(stream.Context())
+	lc, err := pb.NewLeaseClient(conn).LeaseKeepAlive(ctx)
+	if err != nil {
+		cancel()
+		return err
+	}
+
+	go func() {
+		// Cancel the context attached to lc to unblock lc.Recv when
+		// this routine returns on error.
+		defer cancel()
+
+		for {
+			// stream.Recv will be unblock when the loop in the parent routine
+			// returns on error.
+			rr, err := stream.Recv()
+			if err != nil {
+				return
+			}
+			err = lc.Send(rr)
+			if err != nil {
+				return
+			}
+		}
+	}()
+
+	for {
+		rr, err := lc.Recv()
+		if err != nil {
+			return err
+		}
+		err = stream.Send(rr)
+		if err != nil {
+			return err
+		}
+	}
+}