github.com/grafana/pyroscope@v1.18.0/pkg/metastore/discovery/kuberesolver/builder.go (about)

     1  package kuberesolver
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"io"
     7  	"sync"
     8  	"time"
     9  
    10  	"github.com/go-kit/log"
    11  )
    12  
    13  const (
    14  	defaultFreq = time.Minute * 30
    15  )
    16  
    17  type TargetInfo struct {
    18  	ServiceName      string
    19  	ServiceNamespace string
    20  }
    21  
    22  func Build(l log.Logger, k8sClient K8sClient, upd ResolveUpdates, target TargetInfo) (*KResolver, error) {
    23  	if k8sClient == nil {
    24  		return nil, fmt.Errorf("k8sClient is nil")
    25  	}
    26  	ti := target
    27  	if ti.ServiceNamespace == "" {
    28  		ti.ServiceNamespace = getCurrentNamespaceOrDefault()
    29  	}
    30  	ctx, cancel := context.WithCancel(context.Background())
    31  	r := &KResolver{
    32  		target:    ti,
    33  		upd:       upd,
    34  		l:         l,
    35  		ctx:       ctx,
    36  		cancel:    cancel,
    37  		k8sClient: k8sClient,
    38  		t:         time.NewTimer(defaultFreq),
    39  		freq:      defaultFreq,
    40  	}
    41  	go until(func() {
    42  		r.wg.Add(1)
    43  		err := r.watch()
    44  		if err != nil && err != io.EOF {
    45  			l.Log("msg", "watching ended with error, will reconnect again", "err", err)
    46  		}
    47  	}, time.Second, time.Second*30, ctx.Done())
    48  	return r, nil
    49  }
    50  
    51  type ResolveUpdates interface {
    52  	Resolved(e Endpoints)
    53  }
    54  
    55  type ResolveUpdatesFunc func(e Endpoints)
    56  
    57  func (f ResolveUpdatesFunc) Resolved(e Endpoints) {
    58  	f(e)
    59  }
    60  
    61  type KResolver struct {
    62  	target    TargetInfo
    63  	ctx       context.Context
    64  	cancel    context.CancelFunc
    65  	k8sClient K8sClient
    66  	// wg is used to enforce Close() to return after the watcher() goroutine has finished.
    67  	wg   sync.WaitGroup
    68  	t    *time.Timer
    69  	freq time.Duration
    70  	upd  ResolveUpdates
    71  	l    log.Logger
    72  }
    73  
    74  // Close closes the resolver.
    75  func (k *KResolver) Close() {
    76  	k.cancel()
    77  	k.wg.Wait()
    78  }
    79  
    80  func (k *KResolver) handle(e Endpoints) {
    81  	k.upd.Resolved(e)
    82  }
    83  
    84  func (k *KResolver) resolve() {
    85  	e, err := getEndpoints(k.k8sClient, k.target.ServiceNamespace, k.target.ServiceName)
    86  	if err == nil {
    87  		k.handle(e)
    88  	} else {
    89  		k.l.Log("msg", "lookup endpoints failed", "err", err)
    90  	}
    91  	// Next lookup should happen after an interval defined by k.freq.
    92  	k.t.Reset(k.freq)
    93  }
    94  
    95  func (k *KResolver) watch() error {
    96  	defer k.wg.Done()
    97  	// watch endpoints lists existing endpoints at start
    98  	sw, err := watchEndpoints(k.ctx, k.k8sClient, k.target.ServiceNamespace, k.target.ServiceName)
    99  	if err != nil {
   100  		return err
   101  	}
   102  	for {
   103  		select {
   104  		case <-k.ctx.Done():
   105  			return nil
   106  		case <-k.t.C:
   107  			k.resolve()
   108  		case up, hasMore := <-sw.ResultChan():
   109  			if hasMore {
   110  				k.handle(up.Object)
   111  			} else {
   112  				return nil
   113  			}
   114  		}
   115  	}
   116  }