From aa3f46c89786ab8d79f0213731cc980b1838a051 Mon Sep 17 00:00:00 2001
From: vyzo <vyzo@hackzen.org>
Date: Mon, 23 Apr 2018 14:50:46 +0300
Subject: [PATCH] user parameter for ttl in Register

---
 client.go | 18 +++++++++++-------
 1 file changed, 11 insertions(+), 7 deletions(-)

diff --git a/client.go b/client.go
index 8a42219..b20c31a 100644
--- a/client.go
+++ b/client.go
@@ -73,27 +73,31 @@ func (cli *client) Register(ctx context.Context, ns string, ttl int) error {
 	return nil
 }
 
-func Register(ctx context.Context, rz Rendezvous, ns string) error {
-	err := rz.Register(ctx, ns, DefaultTTL)
+func Register(ctx context.Context, rz Rendezvous, ns string, ttl int) error {
+	if ttl < 120 {
+		return fmt.Errorf("registration TTL is too short")
+	}
+
+	err := rz.Register(ctx, ns, ttl)
 	if err != nil {
 		return err
 	}
 
-	go registerRefresh(ctx, rz, ns)
+	go registerRefresh(ctx, rz, ns, ttl)
 	return nil
 }
 
-func registerRefresh(ctx context.Context, rz Rendezvous, ns string) {
-	const refresh = DefaultTTL - 30
+func registerRefresh(ctx context.Context, rz Rendezvous, ns string, ttl int) {
+	refresh := time.Duration(ttl-30) * time.Second
 
 	for {
 		select {
-		case <-time.After(refresh * time.Second):
+		case <-time.After(refresh):
 		case <-ctx.Done():
 			return
 		}
 
-		err := rz.Register(ctx, ns, DefaultTTL)
+		err := rz.Register(ctx, ns, ttl)
 		if err != nil {
 			log.Errorf("Error registering [%s]: %s", ns, err.Error())
 		}