github.com/yankunsam/loki/v2@v2.6.3-0.20220817130409-389df5235c27/pkg/util/dns_watcher.go (about)

     1  package util
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"time"
     7  
     8  	"github.com/grafana/dskit/grpcutil"
     9  	"github.com/grafana/dskit/services"
    10  	"github.com/pkg/errors"
    11  
    12  	util_log "github.com/grafana/loki/pkg/util/log"
    13  )
    14  
    15  // Notifications about address resolution. All notifications are sent on the same goroutine.
    16  type DNSNotifications interface {
    17  	// New address has been discovered by DNS watcher for supplied hostname.
    18  	AddressAdded(address string)
    19  
    20  	// Previously-discovered address is no longer resolved for the hostname.
    21  	AddressRemoved(address string)
    22  }
    23  
    24  type dnsWatcher struct {
    25  	watcher       grpcutil.Watcher
    26  	notifications DNSNotifications
    27  }
    28  
    29  // NewDNSWatcher creates a new DNS watcher and returns a service that is wrapping it.
    30  func NewDNSWatcher(address string, dnsLookupPeriod time.Duration, notifications DNSNotifications) (services.Service, error) {
    31  	resolver, err := grpcutil.NewDNSResolverWithFreq(dnsLookupPeriod, util_log.Logger)
    32  	if err != nil {
    33  		return nil, err
    34  	}
    35  
    36  	watcher, err := resolver.Resolve(address, "")
    37  	if err != nil {
    38  		return nil, err
    39  	}
    40  
    41  	w := &dnsWatcher{
    42  		watcher:       watcher,
    43  		notifications: notifications,
    44  	}
    45  	return services.NewBasicService(nil, w.watchDNSLoop, nil), nil
    46  }
    47  
    48  // watchDNSLoop watches for changes in DNS and sends notifications.
    49  func (w *dnsWatcher) watchDNSLoop(servCtx context.Context) error {
    50  	go func() {
    51  		// Close the watcher, when this service is asked to stop.
    52  		// Closing the watcher makes watchDNSLoop exit, since it only iterates on watcher updates, and has no other
    53  		// way to stop. We cannot close the watcher in `stopping` method, because it is only called *after*
    54  		// watchDNSLoop exits.
    55  		<-servCtx.Done()
    56  		w.watcher.Close()
    57  	}()
    58  
    59  	for {
    60  		updates, err := w.watcher.Next()
    61  		if err != nil {
    62  			// watcher.Next returns error when Close is called, but we call Close when our context is done.
    63  			// we don't want to report error in that case.
    64  			if servCtx.Err() != nil {
    65  				return nil
    66  			}
    67  			return errors.Wrapf(err, "error from DNS watcher")
    68  		}
    69  
    70  		for _, update := range updates {
    71  			switch update.Op {
    72  			case grpcutil.Add:
    73  				w.notifications.AddressAdded(update.Addr)
    74  
    75  			case grpcutil.Delete:
    76  				w.notifications.AddressRemoved(update.Addr)
    77  
    78  			default:
    79  				return fmt.Errorf("unknown op: %v", update.Op)
    80  			}
    81  		}
    82  	}
    83  }