github.com/grafana/pyroscope@v1.18.0/pkg/util/servicediscovery/dns.go (about)

     1  // SPDX-License-Identifier: AGPL-3.0-only
     2  // Provenance-includes-location: https://github.com/cortexproject/cortex/blob/master/pkg/util/dns_watcher.go
     3  // Provenance-includes-license: Apache-2.0
     4  // Provenance-includes-copyright: The Cortex Authors.
     5  
     6  package servicediscovery
     7  
     8  import (
     9  	"context"
    10  	"fmt"
    11  	"time"
    12  
    13  	"github.com/grafana/dskit/grpcutil"
    14  	"github.com/grafana/dskit/services"
    15  	"github.com/pkg/errors"
    16  
    17  	util_log "github.com/grafana/pyroscope/pkg/util"
    18  )
    19  
    20  // Instance notified by the service discovery.
    21  type Instance struct {
    22  	Address string
    23  
    24  	// InUse is true if this instance should be actively used. For example, if a service discovery
    25  	// implementation enforced a max number of instances to be used, this flag will be set to true
    26  	// only on a number of instances up to the configured max.
    27  	InUse bool
    28  }
    29  
    30  func (i Instance) Equal(other Instance) bool {
    31  	return i.Address == other.Address && i.InUse == other.InUse
    32  }
    33  
    34  // Notifications about address resolution. All notifications are sent on the same goroutine.
    35  type Notifications interface {
    36  	// InstanceAdded is called each time a new instance has been discovered.
    37  	InstanceAdded(instance Instance)
    38  
    39  	// InstanceRemoved is called each time an instance that was previously notified by AddressAdded()
    40  	// is no longer available.
    41  	InstanceRemoved(instance Instance)
    42  
    43  	// InstanceChanged is called each time an instance that was previously notified by AddressAdded()
    44  	// has changed its InUse value.
    45  	InstanceChanged(instance Instance)
    46  }
    47  
    48  type dnsServiceDiscovery struct {
    49  	watcher  grpcutil.Watcher
    50  	receiver Notifications
    51  }
    52  
    53  // NewDNS creates a new DNS-based service discovery.
    54  func NewDNS(address string, dnsLookupPeriod time.Duration, notifications Notifications) (services.Service, error) {
    55  	resolver, err := grpcutil.NewDNSResolverWithFreq(dnsLookupPeriod, util_log.Logger)
    56  	if err != nil {
    57  		return nil, err
    58  	}
    59  
    60  	// Pass empty string for service argument, since we don't intend to lookup any SRV record
    61  	watcher, err := resolver.Resolve(address, "")
    62  	if err != nil {
    63  		return nil, err
    64  	}
    65  
    66  	w := &dnsServiceDiscovery{
    67  		watcher:  watcher,
    68  		receiver: notifications,
    69  	}
    70  	return services.NewBasicService(nil, w.watchDNSLoop, nil), nil
    71  }
    72  
    73  // watchDNSLoop watches for changes in DNS and sends notifications.
    74  func (w *dnsServiceDiscovery) watchDNSLoop(servCtx context.Context) error {
    75  	go func() {
    76  		// Close the watcher, when this service is asked to stop.
    77  		// Closing the watcher makes watchDNSLoop exit, since it only iterates on watcher updates, and has no other
    78  		// way to stop. We cannot close the watcher in `stopping` method, because it is only called *after*
    79  		// watchDNSLoop exits.
    80  		<-servCtx.Done()
    81  		w.watcher.Close()
    82  	}()
    83  
    84  	for {
    85  		updates, err := w.watcher.Next()
    86  		if err != nil {
    87  			// watcher.Next returns error when Close is called, but we call Close when our context is done.
    88  			// we don't want to report error in that case.
    89  			if servCtx.Err() != nil {
    90  				return nil
    91  			}
    92  			return errors.Wrapf(err, "error from DNS service discovery")
    93  		}
    94  
    95  		for _, update := range updates {
    96  			switch update.Op {
    97  			case grpcutil.Add:
    98  				w.receiver.InstanceAdded(Instance{Address: update.Addr, InUse: true})
    99  
   100  			case grpcutil.Delete:
   101  				w.receiver.InstanceRemoved(Instance{Address: update.Addr, InUse: true})
   102  
   103  			default:
   104  				return fmt.Errorf("unknown op: %v", update.Op)
   105  			}
   106  		}
   107  	}
   108  }