github.com/pyroscope-io/pyroscope@v0.37.3-0.20230725203016-5f6947968bd0/pkg/scrape/discovery/refresh/refresh.go (about)

     1  // Copyright 2013 The Prometheus Authors
     2  // Copyright 2021 The Pyroscope Authors
     3  //
     4  // Licensed under the Apache License, Version 2.0 (the "License");
     5  // you may not use this file except in compliance with the License.
     6  // You may obtain a copy of the License at
     7  //
     8  // http://www.apache.org/licenses/LICENSE-2.0
     9  //
    10  // Unless required by applicable law or agreed to in writing, software
    11  // distributed under the License is distributed on an "AS IS" BASIS,
    12  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  // See the License for the specific language governing permissions and
    14  // limitations under the License.
    15  
    16  package refresh
    17  
    18  import (
    19  	"context"
    20  	"time"
    21  
    22  	"github.com/prometheus/client_golang/prometheus"
    23  	"github.com/sirupsen/logrus"
    24  
    25  	"github.com/pyroscope-io/pyroscope/pkg/scrape/discovery/targetgroup"
    26  )
    27  
    28  var (
    29  	failuresCount = prometheus.NewCounterVec(
    30  		prometheus.CounterOpts{
    31  			Name: "pyroscope_sd_refresh_failures_total",
    32  			Help: "Number of refresh failures for the given SD mechanism.",
    33  		},
    34  		[]string{"mechanism"},
    35  	)
    36  	duration = prometheus.NewSummaryVec(
    37  		prometheus.SummaryOpts{
    38  			Name:       "pyroscope_sd_refresh_duration_seconds",
    39  			Help:       "The duration of a refresh in seconds for the given SD mechanism.",
    40  			Objectives: map[float64]float64{0.5: 0.05, 0.9: 0.01, 0.99: 0.001},
    41  		},
    42  		[]string{"mechanism"},
    43  	)
    44  )
    45  
    46  func init() {
    47  	prometheus.MustRegister(duration, failuresCount)
    48  }
    49  
    50  // Discovery implements the Discoverer interface.
    51  type Discovery struct {
    52  	logger   logrus.FieldLogger
    53  	interval time.Duration
    54  	refreshf func(ctx context.Context) ([]*targetgroup.Group, error)
    55  
    56  	failures prometheus.Counter
    57  	duration prometheus.Observer
    58  }
    59  
    60  // NewDiscovery returns a Discoverer function that calls a refresh() function at every interval.
    61  func NewDiscovery(l logrus.FieldLogger, mech string, interval time.Duration, refreshf func(ctx context.Context) ([]*targetgroup.Group, error)) *Discovery {
    62  	return &Discovery{
    63  		logger:   l,
    64  		interval: interval,
    65  		refreshf: refreshf,
    66  		failures: failuresCount.WithLabelValues(mech),
    67  		duration: duration.WithLabelValues(mech),
    68  	}
    69  }
    70  
    71  // Run implements the Discoverer interface.
    72  func (d *Discovery) Run(ctx context.Context, ch chan<- []*targetgroup.Group) {
    73  	// Get an initial set right away.
    74  	tgs, err := d.refresh(ctx)
    75  	if err != nil {
    76  		if ctx.Err() != context.Canceled {
    77  			d.logger.WithError(err).Errorf("unable to refresh target groups")
    78  		}
    79  	} else {
    80  		select {
    81  		case ch <- tgs:
    82  		case <-ctx.Done():
    83  			return
    84  		}
    85  	}
    86  
    87  	ticker := time.NewTicker(d.interval)
    88  	defer ticker.Stop()
    89  
    90  	for {
    91  		select {
    92  		case <-ticker.C:
    93  			tgs, err = d.refresh(ctx)
    94  			if err != nil {
    95  				if ctx.Err() != context.Canceled {
    96  					d.logger.WithError(err).Errorf("unable to refresh target groups")
    97  				}
    98  				continue
    99  			}
   100  
   101  			select {
   102  			case ch <- tgs:
   103  			case <-ctx.Done():
   104  				return
   105  			}
   106  		case <-ctx.Done():
   107  			return
   108  		}
   109  	}
   110  }
   111  
   112  func (d *Discovery) refresh(ctx context.Context) ([]*targetgroup.Group, error) {
   113  	now := time.Now()
   114  	defer d.duration.Observe(time.Since(now).Seconds())
   115  	tgs, err := d.refreshf(ctx)
   116  	if err != nil {
   117  		d.failures.Inc()
   118  	}
   119  	return tgs, err
   120  }