Merge pull request #253 from sgotti/renew_lease_if_expired

server: create new lease if expired

errors.go

@@ -62,4 +62,7 @@ var (
 
 	// ErrLocked is returned if the resource is locked.
 	ErrLocked = errors.New("torus: locked")
+
+	// ErrLeaseNotFound is returned if the lease cannot be found.
+	ErrLeaseNotFound = errors.New("torus: lease not found")
 )

heartbeat.go

@@ -52,7 +52,7 @@ func (s *Server) BeginHeartbeat(addr *url.URL) error {
 		s.peerInfo.Address = advertiseURI.String()
 	}
 	var err error
-	s.lease, err = s.MDS.GetLease()
+	err = s.createOrRenewLease(context.Background())
 	if err != nil {
 		return err
 	}
@@ -89,11 +89,15 @@ func (s *Server) oneHeartbeat() {
 	s.peerInfo.UsedBlocks = s.Blocks.UsedBlocks()
 	s.mut.Unlock()
 
-	s.infoMut.Lock()
-	defer s.infoMut.Unlock()
 	ctx, cancel := context.WithTimeout(context.Background(), heartbeatTimeout)
 	defer cancel()
-	err := s.MDS.WithContext(ctx).RegisterPeer(s.lease, s.peerInfo)
+	err := s.createOrRenewLease(ctx)
+	if err != nil {
+		clog.Warningf("failed to create or renew lease: %s", err)
+	}
+	s.infoMut.Lock()
+	defer s.infoMut.Unlock()
+	err = s.MDS.WithContext(ctx).RegisterPeer(s.lease, s.peerInfo)
 	if err != nil {
 		clog.Warningf("couldn't register heartbeat: %s", err)
 	}

metadata.go

@@ -41,6 +41,7 @@ type MetadataService interface {
 	WithContext(ctx context.Context) MetadataService
 
 	GetLease() (int64, error)
+	RenewLease(int64) error
 	RegisterPeer(lease int64, pi *models.PeerInfo) error
 	GetPeers() (PeerInfoList, error)
 

metadata/etcd/etcd.go

@@ -12,6 +12,7 @@ import (
 	"github.com/coreos/torus/ring"
 
 	etcdv3 "github.com/coreos/etcd/clientv3"
+	"github.com/coreos/etcd/etcdserver/api/v3rpc/rpctypes"
 	"github.com/coreos/pkg/capnslog"
 	"github.com/prometheus/client_golang/prometheus"
 	"golang.org/x/net/context"
@@ -190,11 +191,6 @@ func (c *etcdCtx) RegisterPeer(lease int64, p *models.PeerInfo) error {
 	}
 
 	lid := etcdv3.LeaseID(lease)
-	resp, err := c.etcd.Client.KeepAliveOnce(c.getContext(), lid)
-	if err != nil {
-		return err
-	}
-	clog.Tracef("updated lease for %d, TTL %d", resp.ID, resp.TTL)
 	_, err = c.etcd.Client.Put(
 		c.getContext(), MkKey("nodes", p.UUID), string(data), etcdv3.WithLease(lid))
 	return err
@@ -345,6 +341,18 @@ func (c *etcdCtx) GetLease() (int64, error) {
 	return int64(resp.ID), nil
 }
 
+func (c *etcdCtx) RenewLease(lease int64) error {
+	lid := etcdv3.LeaseID(lease)
+	resp, err := c.etcd.Client.KeepAliveOnce(c.getContext(), lid)
+	if err != nil {
+		if err == rpctypes.ErrLeaseNotFound {
+			return torus.ErrLeaseNotFound
+		}
+		return err
+	}
+	clog.Tracef("updated lease for %d, TTL %d", resp.ID, resp.TTL)
+	return nil
+}
 func (c *etcdCtx) GetRing() (torus.Ring, error) {
 	r, _, err := c.getRing()
 	return r, err

metadata/temp/temp.go

@@ -90,7 +90,9 @@ func (t *Client) UUID() string {
 	return t.uuid
 }
 
-func (t *Client) GetLease() (int64, error) { return 1, nil }
+func (t *Client) GetLease() (int64, error)     { return 1, nil }
+func (t *Client) RenewLease(lease int64) error { return nil }
+
 func (t *Client) GetPeers() (torus.PeerInfoList, error) {
 	t.srv.mut.Lock()
 	defer t.srv.mut.Unlock()

server.go

@@ -40,13 +40,31 @@ type Server struct {
 	peerInfo      *models.PeerInfo
 	ctx           context.Context
 
-	lease            int64
+	lease    int64
+	leaseMut sync.Mutex
+
 	heartbeating     bool
 	ReplicationOpen  bool
 	timeoutCallbacks []func(string)
 }
 
+func (s *Server) createOrRenewLease(ctx context.Context) error {
+	s.infoMut.Lock()
+	defer s.infoMut.Unlock()
+	if s.lease != 0 {
+		err := s.MDS.WithContext(ctx).RenewLease(s.lease)
+		if err == nil {
+			return nil
+		}
+	}
+	var err error
+	s.lease, err = s.MDS.WithContext(ctx).GetLease()
+	return err
+}
+
 func (s *Server) Lease() int64 {
+	s.infoMut.Lock()
+	defer s.infoMut.Unlock()
 	return s.lease
 }