From f33b3e2e9f75d37978f6765f0e0f4012667b3ec7 Mon Sep 17 00:00:00 2001
From: Soulou <leo@unbekandt.eu>
Date: Fri, 6 Mar 2020 19:55:45 +0100
Subject: [PATCH] Fix lease time duration, it is now set according to the
 keepalive duration of the IP if set

---
 config/config.go      | 10 ++++++++--
 ip/event_manager.go   |  4 ++--
 locker/etcd_locker.go | 10 +++++++---
 3 files changed, 17 insertions(+), 7 deletions(-)

diff --git a/config/config.go b/config/config.go
index 140a4a6f..af45c8d0 100644
--- a/config/config.go
+++ b/config/config.go
@@ -28,8 +28,14 @@ type Config struct {
 	FailCountBeforeFailover int `envconfig:"FAIL_COUNT_BEFORE_FAILOVER" default:"3"`
 }
 
-func (c Config) LeaseTime() time.Duration {
-	return 5 * c.KeepAliveInterval
+// LeaseTime is either 5* the global keepalive interval, or 5 times the one
+// which is given as argument which can be speicific to the current IP.
+func (c Config) LeaseTime(ipKeepalive int) time.Duration {
+	duration := time.Duration(ipKeepalive) * time.Second
+	if duration == 0 {
+		duration = c.KeepAliveInterval
+	}
+	return 5 * duration
 }
 
 // RandomDurationAround returns a duration (1.0 - percent) * duration < n < (1.0 + percent) * duration
diff --git a/ip/event_manager.go b/ip/event_manager.go
index d61d95a8..14b940a6 100644
--- a/ip/event_manager.go
+++ b/ip/event_manager.go
@@ -88,7 +88,7 @@ func (m *manager) singleEventRun(ctx context.Context) bool {
 		// Sleeping twice the lease time will ensure that we've lost our lease and another node was elected MASTER.
 		// So after this sleep, we can safely remove our IP.
 
-		log.Infof("Stop order received, waiting %s to remove IP", (2 * m.config.LeaseTime()).String())
+		log.Infof("Stop order received, waiting %s to remove IP", (2 * m.config.LeaseTime(m.ip.KeepaliveInterval)).String())
 		m.waitTwiceLeaseTimeOrReallocation(ctx)
 		if m.stopOrder(ctx) {
 			log.Infof("Removing IP %s", m.ip.IP)
@@ -110,7 +110,7 @@ func (m *manager) singleEventRun(ctx context.Context) bool {
 
 func (m *manager) waitTwiceLeaseTimeOrReallocation(ctx context.Context) {
 	log := logger.Get(ctx)
-	timer := time.NewTimer(2 * m.config.LeaseTime())
+	timer := time.NewTimer(2 * m.config.LeaseTime(m.ip.KeepaliveInterval))
 	defer timer.Stop()
 
 	ticker := time.NewTicker(500 * time.Millisecond)
diff --git a/locker/etcd_locker.go b/locker/etcd_locker.go
index 43725f16..8dadc568 100644
--- a/locker/etcd_locker.go
+++ b/locker/etcd_locker.go
@@ -43,7 +43,7 @@ func (l *etcdLocker) Refresh(ctx context.Context) error {
 	log := logger.Get(ctx)
 
 	if l.leaseID == 0 {
-		grant, err := l.leaseEtcd.Grant(ctx, int64(l.config.LeaseTime().Seconds()))
+		grant, err := l.leaseEtcd.Grant(ctx, int64(l.config.LeaseTime(l.ip.KeepaliveInterval).Seconds()))
 		if err != nil {
 			return errors.Wrap(err, "fail to generate grant")
 		}
@@ -56,7 +56,11 @@ func (l *etcdLocker) Refresh(ctx context.Context) error {
 	// The goal of this transaction is to create the key with our leaseID only if this key does not exist
 	// We use a transaction to make sure that concurrent tries wont interfere with each others.
 
-	transactionCtx, cancel := context.WithTimeout(ctx, l.config.KeepAliveInterval)
+	transactionTimeout := time.Duration(l.ip.KeepaliveInterval) * time.Second
+	if transactionTimeout != 0 {
+		transactionTimeout = l.config.KeepAliveInterval
+	}
+	transactionCtx, cancel := context.WithTimeout(ctx, transactionTimeout)
 	defer cancel()
 
 	_, err := l.kvEtcd.Txn(transactionCtx).
@@ -96,7 +100,7 @@ func (l *etcdLocker) Refresh(ctx context.Context) error {
 }
 
 func (l *etcdLocker) leaseExpired() bool {
-	return l.lastLeaseRefresh.IsZero() || time.Now().After(l.lastLeaseRefresh.Add(l.config.LeaseTime()))
+	return l.lastLeaseRefresh.IsZero() || time.Now().After(l.lastLeaseRefresh.Add(l.config.LeaseTime(l.ip.KeepaliveInterval)))
 }
 
 func (l *etcdLocker) Unlock(ctx context.Context) error {