github.com/projecteru2/core@v0.0.0-20240321043226-06bcc1c23f58/cluster/calcium/service.go (about)

     1  package calcium
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  	"time"
     7  
     8  	"github.com/projecteru2/core/log"
     9  	"github.com/projecteru2/core/types"
    10  	"github.com/projecteru2/core/utils"
    11  
    12  	"github.com/cockroachdb/errors"
    13  )
    14  
    15  // WatchServiceStatus returns chan of available service address
    16  func (c *Calcium) WatchServiceStatus(ctx context.Context) (<-chan types.ServiceStatus, error) {
    17  	id, ch := c.watcher.Subscribe(ctx)
    18  	_ = c.pool.Invoke(func() {
    19  		<-ctx.Done()
    20  		c.watcher.Unsubscribe(id)
    21  	})
    22  	return ch, nil
    23  }
    24  
    25  // RegisterService writes self service address in store
    26  func (c *Calcium) RegisterService(ctx context.Context) (unregister func(), err error) {
    27  	serviceAddress, err := utils.GetOutboundAddress(c.config.Bind, c.config.ProbeTarget)
    28  	logger := log.WithFunc("calcium.RegisterService")
    29  	if err != nil {
    30  		logger.Error(ctx, err, "failed to get outbound address")
    31  		return nil, err
    32  	}
    33  
    34  	var (
    35  		expiry            <-chan struct{}
    36  		unregisterService func()
    37  	)
    38  	for {
    39  		if expiry, unregisterService, err = c.registerService(ctx, serviceAddress); err == nil {
    40  			break
    41  		}
    42  		if errors.Is(err, types.ErrKeyExists) {
    43  			logger.Debugf(ctx, "service key exists: %+v", err)
    44  			time.Sleep(time.Second)
    45  			continue
    46  		}
    47  		logger.Error(ctx, err, "failed to first register service")
    48  		return nil, err
    49  	}
    50  
    51  	wg := &sync.WaitGroup{}
    52  	wg.Add(1)
    53  	ctx, cancel := context.WithCancel(ctx)
    54  	_ = c.pool.Invoke(func() {
    55  		defer func() {
    56  			unregisterService()
    57  			wg.Done()
    58  		}()
    59  
    60  		for {
    61  			select {
    62  			case <-expiry:
    63  				// The original one had been expired, we're going to register again.
    64  				if ne, us, err := c.registerService(ctx, serviceAddress); err != nil {
    65  					logger.Error(ctx, err, "failed to re-register service")
    66  					time.Sleep(c.config.GRPCConfig.ServiceHeartbeatInterval)
    67  				} else {
    68  					expiry = ne
    69  					unregisterService = us
    70  				}
    71  
    72  			case <-ctx.Done():
    73  				logger.Infof(ctx, "heartbeat done: %+v", ctx.Err())
    74  				return
    75  			}
    76  		}
    77  	})
    78  	return func() {
    79  		cancel()
    80  		wg.Wait()
    81  	}, nil
    82  }
    83  
    84  func (c *Calcium) registerService(ctx context.Context, addr string) (<-chan struct{}, func(), error) {
    85  	return c.store.RegisterService(ctx, addr, c.config.GRPCConfig.ServiceHeartbeatInterval)
    86  }