From 851c41cd45d5825d08dcb130e09f55a0e58e9245 Mon Sep 17 00:00:00 2001 From: Antonio Navarro Perez Date: Fri, 5 Mar 2021 13:49:18 +0100 Subject: [PATCH] Implement infinite lease renewal logic. Signed-off-by: Antonio Navarro Perez --- etcd/embed.go | 6 +-- etcd/leasedkv.go | 61 ++++++++++++++++++++++++++---- etcd/leasedkv_test.go | 87 +++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 144 insertions(+), 10 deletions(-) create mode 100644 etcd/leasedkv_test.go diff --git a/etcd/embed.go b/etcd/embed.go index 0ad0d7b6b..db8e5672f 100644 --- a/etcd/embed.go +++ b/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") } diff --git a/etcd/leasedkv.go b/etcd/leasedkv.go index 6f0082310..77df1c73c 100644 --- a/etcd/leasedkv.go +++ b/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) + } +} diff --git a/etcd/leasedkv_test.go b/etcd/leasedkv_test.go new file mode 100644 index 000000000..288a02052 --- /dev/null +++ b/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) + } +}