Implement infinite lease renewal logic.

Signed-off-by: Antonio Navarro Perez <antnavper@gmail.com>

etcd/embed.go

@@ -201,10 +201,10 @@ func (e *Etcd) Start(ctx context.Context) (disco.InitialClusterState, error) {
 }
 
 func (e *Etcd) startHeartbeat(ctx context.Context) error {
-	key, value := heartbeatPrefix+e.e.Server.ID().String(), disco.NodeStateStarting
+	key := heartbeatPrefix + e.e.Server.ID().String()
 	e.heartBeatLeasedKV = newLeasedKV(e.cli, key, e.options.HeartbeatTTL)
 
-	if err := e.heartBeatLeasedKV.Start(ctx, string(value)); err != nil {
+	if err := e.heartBeatLeasedKV.Start(string(disco.NodeStateStarting)); err != nil {
 		return errors.Wrap(err, "startHeartbeat: starting a new heartbeat")
 	}
 
@@ -372,7 +372,7 @@ func (e *Etcd) Resize(ctx context.Context) (func([]byte) error, error) {
 		e.resizeLeasedKV = newLeasedKV(e.cli, key, e.options.HeartbeatTTL)
 	}
 
-	if err := e.resizeLeasedKV.Start(ctx, ""); err != nil {
+	if err := e.resizeLeasedKV.Start(""); err != nil {
 		return nil, errors.Wrap(err, "Resize: creates a new hearbeat")
 	}
 

etcd/leasedkv.go

@@ -16,7 +16,9 @@ package etcd
 
 import (
 	"context"
+	"log"
 	"sync"
+	"time"
 
 	"github.com/pilosa/pilosa/v2/disco"
 	"github.com/pkg/errors"
@@ -31,6 +33,7 @@ type leasedKV struct {
 	mu sync.Mutex
 
 	key, value string
+	stopped    bool
 	ttlSeconds int64
 }
 
@@ -42,17 +45,17 @@ func newLeasedKV(cli *clientv3.Client, key string, ttlSeconds int64) *leasedKV {
 	}
 }
 
-func (l *leasedKV) Start(ctx context.Context, initValue string) error {
+func (l *leasedKV) Start(initValue string) error {
 	l.mu.Lock()
 	defer l.mu.Unlock()
 
-	ctx, cancel := context.WithCancel(ctx)
-	l.cancel = cancel
-
-	return l.create(ctx, initValue)
+	return l.create(initValue)
 }
 
-func (l *leasedKV) create(ctx context.Context, initValue string) error {
+func (l *leasedKV) create(initValue string) error {
+	ctx, cancel := context.WithCancel(context.Background())
+	l.cancel = cancel
+
 	leaseResp, err := l.cli.Grant(ctx, l.ttlSeconds)
 	if err != nil {
 		return errors.Wrap(err, "creating a lease")
@@ -64,15 +67,44 @@ func (l *leasedKV) create(ctx context.Context, initValue string) error {
 		return errors.Wrapf(err, "creating key %s with value [%s]", l.key, initValue)
 	}
 
-	if _, err := l.cli.KeepAlive(ctx, leaseResp.ID); err != nil {
+	kaChann, err := l.cli.KeepAlive(ctx, leaseResp.ID)
+	if err != nil {
 		return errors.Wrapf(err, "keeping alive the lease for the key %s with value %s", l.key, l.value)
 	}
 
 	l.value = initValue
 
+	go l.consumeLease(kaChann)
+
 	return nil
 }
 
+func (l *leasedKV) consumeLease(ch <-chan *clientv3.LeaseKeepAliveResponse) {
+	for {
+		kresp, ok := <-ch
+		if !ok {
+			if l.stopped {
+				return
+			}
+
+			l.mu.Lock()
+			defer l.mu.Unlock()
+
+			retry(1*time.Second, func() error {
+				return l.create(l.value)
+			})
+
+			log.Println("lease recreated after a problem. Key:", l.key)
+
+			return
+		}
+
+		if kresp == nil {
+			continue
+		}
+	}
+}
+
 func (l *leasedKV) Stop() {
 	l.mu.Lock()
 	defer l.mu.Unlock()
@@ -80,6 +112,8 @@ func (l *leasedKV) Stop() {
 	if l.cancel != nil {
 		l.cancel()
 	}
+
+	l.stopped = true
 }
 
 func (l *leasedKV) Set(ctx context.Context, value string) error {
@@ -117,3 +151,16 @@ func (l *leasedKV) Get(ctx context.Context) (string, error) {
 
 	return l.value, nil
 }
+
+func retry(sleep time.Duration, f func() error) {
+	for {
+		err := f()
+		if err == nil {
+			return
+		}
+
+		time.Sleep(sleep)
+
+		log.Println("retrying after error:", err)
+	}
+}

etcd/leasedkv_test.go

@@ -0,0 +1,87 @@
+package etcd
+
+import (
+	"context"
+	"errors"
+	"io/ioutil"
+	"os"
+	"testing"
+	"time"
+
+	"github.com/pilosa/pilosa/v2/disco"
+	"go.etcd.io/etcd/embed"
+	"go.etcd.io/etcd/etcdserver/api/v3client"
+)
+
+const initVal = "test"
+const newVal = "newValue"
+
+func TestLeasedKv(t *testing.T) {
+	cfg := embed.NewConfig()
+
+	dir, err := ioutil.TempDir("", "leasedkv-*")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer os.RemoveAll(dir)
+
+	cfg.Dir = dir
+	etcd, err := embed.StartEtcd(cfg)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	cli := v3client.New(etcd.Server)
+
+	lkv := newLeasedKV(cli, "/test", 1)
+
+	ctx := context.Background()
+
+	err = lkv.Start(initVal)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	v, err := lkv.Get(ctx)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if v != initVal {
+		t.Fatal("obtained value is not the same as expected. Obtained:", v, "Expected:", initVal)
+	}
+
+	err = lkv.Set(ctx, "otherValue")
+	if err != nil {
+		t.Fatal(err)
+	}
+	err = lkv.Set(ctx, "3")
+	if err != nil {
+		t.Fatal(err)
+	}
+	err = lkv.Set(ctx, newVal)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	v, err = lkv.Get(ctx)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if v != newVal {
+		t.Fatal("obtained value is not the same as expected. Obtained:", v, "Expected:", initVal)
+	}
+
+	time.Sleep(5 * time.Second)
+
+	lkv.Stop()
+
+	// we need to wait to force the lease expiration
+	time.Sleep(5 * time.Second)
+
+	_, err = lkv.Get(ctx)
+	if err == nil || !errors.Is(err, disco.ErrNoResults) {
+		t.Fatal("expected error:", disco.ErrNoResults, "obtained:", err)
+	}
+}