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

     1  package redis
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"strings"
     7  	"time"
     8  
     9  	"github.com/projecteru2/core/log"
    10  )
    11  
    12  type endpoints map[string]struct{}
    13  
    14  func (e *endpoints) Add(endpoint string) (changed bool) {
    15  	if _, ok := (*e)[endpoint]; !ok {
    16  		(*e)[endpoint] = struct{}{}
    17  		changed = true
    18  	}
    19  	return
    20  }
    21  
    22  func (e *endpoints) Remove(endpoint string) (changed bool) {
    23  	if _, ok := (*e)[endpoint]; ok {
    24  		delete(*e, endpoint)
    25  		changed = true
    26  	}
    27  	return
    28  }
    29  
    30  func (e endpoints) ToSlice() (eps []string) {
    31  	for ep := range e {
    32  		eps = append(eps, ep)
    33  	}
    34  	return
    35  }
    36  
    37  // ServiceStatusStream watches /services/ --prefix
    38  func (r *Rediaron) ServiceStatusStream(ctx context.Context) (chan []string, error) {
    39  	key := fmt.Sprintf(serviceStatusKey, "*")
    40  	ch := make(chan []string)
    41  	_ = r.pool.Invoke(func() {
    42  		defer close(ch)
    43  
    44  		watchC := r.KNotify(ctx, key)
    45  
    46  		data, err := r.getByKeyPattern(ctx, key, 0)
    47  		if err != nil {
    48  			log.WithFunc("store.redis.ServiceStatusStream").Error(ctx, err, "failed to get current services")
    49  			return
    50  		}
    51  		eps := endpoints{}
    52  		for k := range data {
    53  			eps.Add(parseServiceKey(k))
    54  		}
    55  		ch <- eps.ToSlice()
    56  
    57  		for message := range watchC {
    58  			changed := false
    59  			endpoint := parseServiceKey(message.Key)
    60  			switch message.Action {
    61  			case actionSet, actionExpire:
    62  				changed = eps.Add(endpoint)
    63  			case actionDel, actionExpired:
    64  				changed = eps.Remove(endpoint)
    65  			}
    66  			if changed {
    67  				ch <- eps.ToSlice()
    68  			}
    69  		}
    70  	})
    71  	return ch, nil
    72  }
    73  
    74  // RegisterService put /services/{address}
    75  func (r *Rediaron) RegisterService(ctx context.Context, serviceAddress string, expire time.Duration) (<-chan struct{}, func(), error) {
    76  	key := fmt.Sprintf(serviceStatusKey, serviceAddress)
    77  	return r.StartEphemeral(ctx, key, expire)
    78  }
    79  
    80  func parseServiceKey(key string) (endpoint string) {
    81  	parts := strings.Split(key, "/")
    82  	return parts[len(parts)-1]
    83  }