From c207e96a51416638d01e800d02265f7ec5ff1cd7 Mon Sep 17 00:00:00 2001
From: miraclesu <suchuangji@gmail.com>
Date: Mon, 20 Mar 2017 12:03:54 +0800
Subject: [PATCH] node: keepAlive with timeout

---
 models/client.go | 12 ++++++++++++
 models/proc.go   |  5 ++---
 node/node.go     |  6 ++----
 3 files changed, 16 insertions(+), 7 deletions(-)

diff --git a/models/client.go b/models/client.go
index 9c3048c..dc282ab 100644
--- a/models/client.go
+++ b/models/client.go
@@ -80,6 +80,18 @@ func (c *Client) Watch(key string, opts ...client.OpOption) client.WatchChan {
 	return c.Client.Watch(context.Background(), key, opts...)
 }
 
+func (c *Client) Grant(ttl int64) (*client.LeaseGrantResponse, error) {
+	ctx, cancel := context.WithTimeout(context.Background(), c.reqTimeout)
+	defer cancel()
+	return c.Client.Grant(ctx, ttl)
+}
+
+func (c *Client) KeepAliveOnce(id client.LeaseID) (*client.LeaseKeepAliveResponse, error) {
+	ctx, cancel := context.WithTimeout(context.Background(), c.reqTimeout)
+	defer cancel()
+	return c.Client.KeepAliveOnce(ctx, id)
+}
+
 func IsValidAsKeyPath(s string) bool {
 	return strings.IndexByte(s, '/') == -1
 }
diff --git a/models/proc.go b/models/proc.go
index cc798ef..968da64 100644
--- a/models/proc.go
+++ b/models/proc.go
@@ -1,7 +1,6 @@
 package models
 
 import (
-	"context"
 	"fmt"
 	"sync"
 	"time"
@@ -78,7 +77,7 @@ func (l *leaseID) get() client.LeaseID {
 
 func (l *leaseID) set() error {
 	id := client.LeaseID(-1)
-	resp, err := DefalutClient.Grant(context.TODO(), l.ttl+2)
+	resp, err := DefalutClient.Grant(l.ttl + 2)
 	if err == nil {
 		id = resp.ID
 	}
@@ -103,7 +102,7 @@ func (l *leaseID) keepAlive() {
 
 			id := l.get()
 			if id > 0 {
-				_, err := DefalutClient.KeepAliveOnce(context.TODO(), l.ID)
+				_, err := DefalutClient.KeepAliveOnce(l.ID)
 				if err == nil {
 					timer.Reset(duration)
 					continue
diff --git a/node/node.go b/node/node.go
index 88781ea..2e138dd 100644
--- a/node/node.go
+++ b/node/node.go
@@ -6,8 +6,6 @@ import (
 	"strconv"
 	"time"
 
-	"golang.org/x/net/context"
-
 	client "github.com/coreos/etcd/clientv3"
 
 	"sunteng/commons/log"
@@ -77,7 +75,7 @@ func (n *Node) Register() (err error) {
 }
 
 func (n *Node) set() error {
-	resp, err := n.Client.Grant(context.TODO(), n.ttl+2)
+	resp, err := n.Client.Grant(n.ttl + 2)
 	if err != nil {
 		return err
 	}
@@ -100,7 +98,7 @@ func (n *Node) keepAlive() {
 			return
 		case <-timer.C:
 			if n.lID > 0 {
-				_, err := n.Client.KeepAliveOnce(context.TODO(), n.lID)
+				_, err := n.Client.KeepAliveOnce(n.lID)
 				if err == nil {
 					timer.Reset(duration)
 					continue