github.com/projecteru2/core@v0.0.0-20240321043226-06bcc1c23f58/client/servicediscovery/eru_service_discovery.go (about)

     1  package servicediscovery
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"math"
     7  	"time"
     8  
     9  	"github.com/projecteru2/core/auth"
    10  	"github.com/projecteru2/core/client/interceptor"
    11  	"github.com/projecteru2/core/client/utils"
    12  	"github.com/projecteru2/core/log"
    13  	pb "github.com/projecteru2/core/rpc/gen"
    14  	"github.com/projecteru2/core/types"
    15  
    16  	"google.golang.org/grpc"
    17  	"google.golang.org/grpc/credentials/insecure"
    18  )
    19  
    20  // EruServiceDiscovery watches eru service status
    21  type EruServiceDiscovery struct {
    22  	endpoint   string
    23  	authConfig types.AuthConfig
    24  }
    25  
    26  // New EruServiceDiscovery
    27  func New(endpoint string, authConfig types.AuthConfig) *EruServiceDiscovery {
    28  	return &EruServiceDiscovery{
    29  		endpoint:   endpoint,
    30  		authConfig: authConfig,
    31  	}
    32  }
    33  
    34  // Watch .
    35  func (w *EruServiceDiscovery) Watch(ctx context.Context) (_ <-chan []string, err error) {
    36  	cc, err := w.dial(ctx, w.endpoint, w.authConfig)
    37  	logger := log.WithFunc("servicediscovery.Watch").WithField("endpoint", w.endpoint)
    38  	if err != nil {
    39  		logger.Error(ctx, err, "dial failed")
    40  		return nil, err
    41  	}
    42  	client := pb.NewCoreRPCClient(cc)
    43  	ch := make(chan []string)
    44  	epPusher := utils.NewEndpointPusher()
    45  	epPusher.Register(ch)
    46  	epPusher.Register(lbResolverBuilder.updateCh)
    47  	go func() {
    48  		defer close(ch)
    49  		for {
    50  			watchCtx, cancelWatch := context.WithCancel(ctx)
    51  			stream, err := client.WatchServiceStatus(watchCtx, &pb.Empty{})
    52  			if err != nil {
    53  				logger.Error(ctx, err, "watch failed, try later")
    54  				time.Sleep(10 * time.Second) // TODO can config
    55  				continue
    56  			}
    57  			expectedInterval := time.Duration(math.MaxInt64) / time.Second
    58  
    59  			for {
    60  				cancelTimer := make(chan struct{})
    61  				go func(expectedInterval time.Duration) {
    62  					timer := time.NewTimer(expectedInterval * time.Second)
    63  					defer timer.Stop()
    64  					select {
    65  					case <-timer.C:
    66  						cancelWatch()
    67  					case <-cancelTimer:
    68  						return
    69  					}
    70  				}(expectedInterval)
    71  				status, err := stream.Recv()
    72  				close(cancelTimer)
    73  				if err != nil {
    74  					logger.Error(ctx, err, "recv failed")
    75  					break
    76  				}
    77  				expectedInterval = time.Duration(status.GetIntervalInSecond())
    78  
    79  				epPusher.Push(ctx, status.GetAddresses())
    80  			}
    81  		}
    82  	}()
    83  
    84  	return ch, nil
    85  }
    86  
    87  func (w *EruServiceDiscovery) dial(ctx context.Context, addr string, authConfig types.AuthConfig) (*grpc.ClientConn, error) {
    88  	opts := []grpc.DialOption{
    89  		grpc.WithTransportCredentials(insecure.NewCredentials()),
    90  		grpc.WithStreamInterceptor(interceptor.NewStreamRetry(interceptor.RetryOptions{Max: 1})),
    91  	}
    92  
    93  	if authConfig.Username != "" {
    94  		opts = append(opts, grpc.WithPerRPCCredentials(auth.NewCredential(authConfig)))
    95  	}
    96  
    97  	target := makeServiceDiscoveryTarget(addr)
    98  	return grpc.DialContext(ctx, target, opts...)
    99  }
   100  
   101  func makeServiceDiscoveryTarget(addr string) string {
   102  	return fmt.Sprintf("lb://_/%s", addr)
   103  }