github.com/mailgun/holster/v4@v4.20.0/etcdutil/session.go (about)

     1  package etcdutil
     2  
     3  import (
     4  	"context"
     5  	"sync/atomic"
     6  	"time"
     7  
     8  	"github.com/mailgun/holster/v4/errors"
     9  	"github.com/mailgun/holster/v4/setter"
    10  	"github.com/mailgun/holster/v4/syncutil"
    11  	etcd "go.etcd.io/etcd/client/v3"
    12  )
    13  
    14  const NoLease = etcd.LeaseID(-1)
    15  
    16  type SessionObserver func(etcd.LeaseID, error)
    17  
    18  type Session struct {
    19  	keepAlive     <-chan *etcd.LeaseKeepAliveResponse
    20  	lease         *etcd.LeaseGrantResponse
    21  	backOff       *backOffCounter
    22  	wg            syncutil.WaitGroup
    23  	ctx           context.Context
    24  	cancel        context.CancelFunc
    25  	observer      SessionObserver
    26  	client        *etcd.Client
    27  	ttl           time.Duration
    28  	lastKeepAlive time.Time
    29  	isRunning     int32
    30  }
    31  
    32  type SessionConfig struct {
    33  	TTL      int64
    34  	Observer SessionObserver
    35  }
    36  
    37  // NewSession creates a lease and monitors lease keep alive's for connectivity.
    38  // Once a lease ID is granted SessionConfig.Observer is called with the granted lease.
    39  // If connectivity is lost with etcd SessionConfig.Observer is called again with -1 (NoLease)
    40  // as the lease ID. The Session will continue to try to gain another lease, once a new lease
    41  // is gained SessionConfig.Observer is called again with the new lease id.
    42  func NewSession(c *etcd.Client, conf SessionConfig) (*Session, error) {
    43  	setter.SetDefault(&conf.TTL, int64(30))
    44  
    45  	if conf.Observer == nil {
    46  		return nil, errors.New("provided observer function cannot be nil")
    47  	}
    48  
    49  	if c == nil {
    50  		return nil, errors.New("provided etcd client cannot be nil")
    51  	}
    52  
    53  	ttlDuration := time.Second * time.Duration(conf.TTL)
    54  	s := Session{
    55  		observer: conf.Observer,
    56  		ttl:      ttlDuration,
    57  		backOff:  newBackOffCounter(time.Millisecond*500, ttlDuration, 2),
    58  		client:   c,
    59  	}
    60  
    61  	s.start()
    62  	return &s, nil
    63  }
    64  
    65  func (s *Session) start() {
    66  	s.ctx, s.cancel = context.WithCancel(context.Background())
    67  	ticker := time.NewTicker(s.ttl)
    68  	s.lastKeepAlive = time.Now()
    69  	atomic.StoreInt32(&s.isRunning, 1)
    70  
    71  	s.wg.Until(func(done chan struct{}) bool {
    72  		// If we have lost our keep alive, attempt to regain it
    73  		if s.keepAlive == nil {
    74  			if err := s.gainLease(s.ctx); err != nil {
    75  				s.observer(NoLease, errors.Wrap(err, "while attempting to gain new lease"))
    76  				select {
    77  				case <-time.After(s.backOff.Next()):
    78  					return true
    79  				case <-s.ctx.Done():
    80  					atomic.StoreInt32(&s.isRunning, 0)
    81  					return false
    82  				}
    83  				// TODO: Fix this in the library. Unreachable code
    84  				// return true
    85  			}
    86  		}
    87  		s.backOff.Reset()
    88  
    89  		select {
    90  		case _, ok := <-s.keepAlive:
    91  			if !ok {
    92  				// heartbeat lost
    93  				s.keepAlive = nil
    94  			} else {
    95  				// heartbeat received
    96  				s.lastKeepAlive = time.Now()
    97  			}
    98  		case <-ticker.C:
    99  			// Ensure we are getting heartbeats regularly
   100  			if time.Since(s.lastKeepAlive) > s.ttl {
   101  				// too long between heartbeats
   102  				s.keepAlive = nil
   103  			}
   104  		case <-done:
   105  			s.keepAlive = nil
   106  			if s.lease != nil {
   107  				ctx, cancel := context.WithTimeout(context.Background(), s.ttl)
   108  				if _, err := s.client.Revoke(ctx, s.lease.ID); err != nil {
   109  					s.observer(NoLease, errors.Wrap(err, "while revoking our lease during shutdown"))
   110  				}
   111  				cancel()
   112  			}
   113  			atomic.StoreInt32(&s.isRunning, 0)
   114  			return false
   115  		}
   116  
   117  		if s.keepAlive == nil {
   118  			s.observer(NoLease, nil)
   119  		}
   120  		return true
   121  	})
   122  }
   123  
   124  func (s *Session) Reset() {
   125  	if atomic.LoadInt32(&s.isRunning) != 1 {
   126  		return
   127  	}
   128  	s.Close()
   129  	s.start()
   130  }
   131  
   132  // Close terminates the session shutting down all network operations,
   133  // then SessionConfig.Observer is called with -1 (NoLease), only returns
   134  // once the session has closed successfully.
   135  func (s *Session) Close() {
   136  	if atomic.LoadInt32(&s.isRunning) != 1 {
   137  		return
   138  	}
   139  
   140  	s.cancel()
   141  	s.wg.Stop()
   142  	s.observer(NoLease, nil)
   143  }
   144  
   145  func (s *Session) gainLease(ctx context.Context) error {
   146  	var err error
   147  	s.lease, err = s.client.Grant(ctx, int64(s.ttl/time.Second))
   148  	if err != nil {
   149  		return errors.Wrapf(err, "during grant lease")
   150  	}
   151  
   152  	s.keepAlive, err = s.client.KeepAlive(s.ctx, s.lease.ID)
   153  	if err != nil {
   154  		return err
   155  	}
   156  	s.observer(s.lease.ID, nil)
   157  	return nil
   158  }