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

     1  package etcdv3
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"strings"
     7  	"time"
     8  
     9  	"github.com/projecteru2/core/log"
    10  
    11  	"go.etcd.io/etcd/api/v3/mvccpb"
    12  	clientv3 "go.etcd.io/etcd/client/v3"
    13  )
    14  
    15  type endpoints map[string]struct{}
    16  
    17  func (e *endpoints) Add(endpoint string) (changed bool) {
    18  	if _, ok := (*e)[endpoint]; !ok {
    19  		(*e)[endpoint] = struct{}{}
    20  		changed = true
    21  	}
    22  	return
    23  }
    24  
    25  func (e *endpoints) Remove(endpoint string) (changed bool) {
    26  	if _, ok := (*e)[endpoint]; ok {
    27  		delete(*e, endpoint)
    28  		changed = true
    29  	}
    30  	return
    31  }
    32  
    33  func (e endpoints) ToSlice() (eps []string) {
    34  	for ep := range e {
    35  		eps = append(eps, ep)
    36  	}
    37  	return
    38  }
    39  
    40  // ServiceStatusStream watches /services/ --prefix
    41  func (m *Mercury) ServiceStatusStream(ctx context.Context) (chan []string, error) {
    42  	ch := make(chan []string)
    43  	logger := log.WithFunc("store.etcdv3.ServiceStatusStream")
    44  	_ = m.pool.Invoke(func() {
    45  		defer close(ch)
    46  
    47  		// must watch prior to get
    48  		watchChan := m.Watch(ctx, fmt.Sprintf(serviceStatusKey, ""), clientv3.WithPrefix())
    49  
    50  		resp, err := m.Get(ctx, fmt.Sprintf(serviceStatusKey, ""), clientv3.WithPrefix())
    51  		if err != nil {
    52  			logger.Error(ctx, err, "failed to get current services")
    53  			return
    54  		}
    55  		eps := endpoints{}
    56  		for _, ev := range resp.Kvs {
    57  			eps.Add(parseServiceKey(ev.Key))
    58  		}
    59  		ch <- eps.ToSlice()
    60  
    61  		for resp := range watchChan {
    62  			if resp.Err() != nil {
    63  				if !resp.Canceled {
    64  					logger.Error(ctx, err, "watch failed")
    65  				}
    66  				return
    67  			}
    68  
    69  			changed := false
    70  			for _, ev := range resp.Events {
    71  				endpoint := parseServiceKey(ev.Kv.Key)
    72  				c := false
    73  				switch ev.Type {
    74  				case mvccpb.PUT:
    75  					c = eps.Add(endpoint)
    76  				case mvccpb.DELETE:
    77  					c = eps.Remove(endpoint)
    78  				}
    79  				if c {
    80  					changed = true
    81  				}
    82  			}
    83  			if changed {
    84  				ch <- eps.ToSlice()
    85  			}
    86  		}
    87  	})
    88  	return ch, nil
    89  }
    90  
    91  // RegisterService put /services/{address}
    92  func (m *Mercury) RegisterService(ctx context.Context, serviceAddress string, expire time.Duration) (<-chan struct{}, func(), error) {
    93  	key := fmt.Sprintf(serviceStatusKey, serviceAddress)
    94  	return m.StartEphemeral(ctx, key, expire)
    95  }
    96  
    97  func parseServiceKey(key []byte) (endpoint string) {
    98  	parts := strings.Split(string(key), "/")
    99  	return parts[len(parts)-1]
   100  }