github.com/projecteru2/core@v0.0.0-20240321043226-06bcc1c23f58/store/etcdv3/meta/ephemeral.go (about)

     1  package meta
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  	"time"
     7  
     8  	"github.com/projecteru2/core/log"
     9  	"github.com/projecteru2/core/types"
    10  
    11  	"github.com/cockroachdb/errors"
    12  	clientv3 "go.etcd.io/etcd/client/v3"
    13  )
    14  
    15  // StartEphemeral starts an empheral kv pair.
    16  func (e *ETCD) StartEphemeral(ctx context.Context, path string, heartbeat time.Duration) (<-chan struct{}, func(), error) {
    17  	lease, err := e.cliv3.Grant(ctx, int64(heartbeat/time.Second))
    18  	if err != nil {
    19  		return nil, nil, err
    20  	}
    21  
    22  	switch tx, err := e.cliv3.Txn(ctx).
    23  		If(clientv3.Compare(clientv3.Version(path), "=", 0)).
    24  		Then(clientv3.OpPut(path, "", clientv3.WithLease(lease.ID))).
    25  		Commit(); {
    26  	case err != nil:
    27  		return nil, nil, err
    28  	case !tx.Succeeded:
    29  		return nil, nil, errors.Wrap(types.ErrKeyExists, path)
    30  	}
    31  
    32  	ctx, cancel := context.WithCancel(ctx)
    33  	expiry := make(chan struct{})
    34  	logger := log.WithFunc("store.etcdv3.meta.StartEphemeral")
    35  
    36  	var wg sync.WaitGroup
    37  	wg.Add(1)
    38  	go func() {
    39  		defer wg.Done()
    40  		defer close(expiry)
    41  
    42  		tick := time.NewTicker(heartbeat / 3)
    43  		defer tick.Stop()
    44  
    45  		// Revokes the lease.
    46  		defer func() {
    47  			// It shouldn't be inheriting from the ctx.
    48  			ctx, cancel := context.WithTimeout(context.TODO(), time.Minute)
    49  			defer cancel()
    50  			if _, err := e.cliv3.Revoke(ctx, lease.ID); err != nil {
    51  				logger.Errorf(ctx, err, "revoke %d with %s failed", lease.ID, path)
    52  			}
    53  		}()
    54  
    55  		for {
    56  			select {
    57  			case <-tick.C:
    58  				if _, err := e.cliv3.KeepAliveOnce(ctx, lease.ID); err != nil {
    59  					logger.Errorf(ctx, err, "keepalive %d with %s failed", lease.ID, path)
    60  					return
    61  				}
    62  			case <-ctx.Done():
    63  				return
    64  			}
    65  		}
    66  	}()
    67  
    68  	return expiry, func() {
    69  		cancel()
    70  		wg.Wait()
    71  	}, nil
    72  }