github.com/projecteru2/core@v0.0.0-20240321043226-06bcc1c23f58/discovery/helium/helium.go (about)

     1  package helium
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  	"time"
     7  
     8  	"github.com/cockroachdb/errors"
     9  	"github.com/projecteru2/core/log"
    10  	"github.com/projecteru2/core/store"
    11  	"github.com/projecteru2/core/types"
    12  
    13  	"github.com/alphadose/haxmap"
    14  	"github.com/google/uuid"
    15  )
    16  
    17  const interval = 15 * time.Second
    18  
    19  // Helium .
    20  type Helium struct {
    21  	sync.Once
    22  	store     store.Store
    23  	subs      *haxmap.Map[uint32, entry]
    24  	interval  time.Duration
    25  	unsubChan chan uint32
    26  }
    27  
    28  type entry struct {
    29  	ch     chan types.ServiceStatus
    30  	ctx    context.Context
    31  	cancel context.CancelFunc
    32  }
    33  
    34  // New .
    35  func New(ctx context.Context, config types.GRPCConfig, store store.Store) *Helium {
    36  	h := &Helium{
    37  		interval:  config.ServiceDiscoveryPushInterval,
    38  		store:     store,
    39  		subs:      haxmap.New[uint32, entry](),
    40  		unsubChan: make(chan uint32),
    41  	}
    42  	if h.interval < time.Second {
    43  		h.interval = interval
    44  	}
    45  	h.Do(func() {
    46  		h.start(ctx)
    47  	})
    48  	return h
    49  }
    50  
    51  // Subscribe .
    52  func (h *Helium) Subscribe(ctx context.Context) (uuid.UUID, <-chan types.ServiceStatus) {
    53  	ID := uuid.New()
    54  	key := ID.ID()
    55  	subCtx, cancel := context.WithCancel(ctx)
    56  	ch := make(chan types.ServiceStatus)
    57  	h.subs.Set(key, entry{
    58  		ch:     ch,
    59  		ctx:    subCtx,
    60  		cancel: cancel,
    61  	})
    62  	return ID, ch
    63  }
    64  
    65  // Unsubscribe .
    66  func (h *Helium) Unsubscribe(ID uuid.UUID) {
    67  	h.unsubChan <- ID.ID()
    68  }
    69  
    70  func (h *Helium) start(ctx context.Context) {
    71  	logger := log.WithFunc("helium.start")
    72  	ch, err := h.store.ServiceStatusStream(ctx)
    73  	if err != nil {
    74  		logger.Error(ctx, err, "failed to start watch")
    75  		return
    76  	}
    77  
    78  	go func() {
    79  		logger.Info(ctx, "service discovery start")
    80  		defer logger.Warn(ctx, "service discovery exited")
    81  		var latestStatus types.ServiceStatus
    82  		ticker := time.NewTicker(h.interval)
    83  		defer ticker.Stop()
    84  		for {
    85  			select {
    86  			case addresses, ok := <-ch:
    87  				if !ok {
    88  					logger.Warn(ctx, "watch channel closed")
    89  					return
    90  				}
    91  
    92  				latestStatus = types.ServiceStatus{
    93  					Addresses: addresses,
    94  					Interval:  h.interval * 2,
    95  				}
    96  
    97  			case ID := <-h.unsubChan:
    98  				if entry, ok := h.subs.Get(ID); ok {
    99  					entry.cancel()
   100  					h.subs.Del(ID)
   101  					close(entry.ch)
   102  				}
   103  
   104  			case <-ticker.C:
   105  			}
   106  
   107  			h.dispatch(ctx, latestStatus)
   108  		}
   109  	}()
   110  }
   111  
   112  func (h *Helium) dispatch(ctx context.Context, status types.ServiceStatus) {
   113  	f := func(key uint32, val entry) {
   114  		defer func() {
   115  			if err := recover(); err != nil {
   116  				log.WithFunc("helium.dispatch").Errorf(ctx, errors.Errorf("%+v", err), "dispatch %+v failed", key)
   117  			}
   118  		}()
   119  		select {
   120  		case val.ch <- status:
   121  			return
   122  		case <-val.ctx.Done():
   123  			return
   124  		}
   125  	}
   126  	h.subs.ForEach(func(k uint32, v entry) bool {
   127  		f(k, v)
   128  		return true
   129  	})
   130  }