github.com/grafana/pyroscope@v1.18.0/pkg/distributor/aggregator/multitenant_aggregator.go (about)

     1  package aggregator
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  	"time"
     7  
     8  	"github.com/grafana/dskit/services"
     9  	"github.com/prometheus/client_golang/prometheus"
    10  	"github.com/prometheus/common/model"
    11  	"go.uber.org/atomic"
    12  )
    13  
    14  type MultiTenantAggregator[T any] struct {
    15  	*services.BasicService
    16  
    17  	limits     Limits
    18  	registerer prometheus.Registerer
    19  
    20  	m       sync.RWMutex
    21  	tenants map[tenantKey]*tenantAggregator[T]
    22  
    23  	closeOnce sync.Once
    24  	stop      chan struct{}
    25  	done      chan struct{}
    26  }
    27  
    28  type Limits interface {
    29  	DistributorAggregationWindow(tenantID string) model.Duration
    30  	DistributorAggregationPeriod(tenantID string) model.Duration
    31  }
    32  
    33  func NewMultiTenantAggregator[T any](limits Limits, registerer prometheus.Registerer) *MultiTenantAggregator[T] {
    34  	m := MultiTenantAggregator[T]{
    35  		limits:     limits,
    36  		registerer: registerer,
    37  		tenants:    make(map[tenantKey]*tenantAggregator[T]),
    38  		stop:       make(chan struct{}),
    39  		done:       make(chan struct{}),
    40  	}
    41  	m.BasicService = services.NewBasicService(
    42  		m.starting,
    43  		m.running,
    44  		m.stopping,
    45  	)
    46  	return &m
    47  }
    48  
    49  type tenantAggregator[T any] struct {
    50  	lastSeen   atomic.Time
    51  	key        tenantKey
    52  	collector  prometheus.Collector
    53  	registerer prometheus.Registerer
    54  	aggregator *Aggregator[T]
    55  }
    56  
    57  type tenantKey struct {
    58  	tenantID string
    59  	window   model.Duration
    60  	period   model.Duration
    61  }
    62  
    63  func (m *MultiTenantAggregator[T]) AggregatorForTenant(tenantID string) (*Aggregator[T], bool) {
    64  	now := time.Now()
    65  	t, ok := m.aggregatorForTenant(tenantID)
    66  	if ok {
    67  		t.lastSeen.Store(now)
    68  		return t.aggregator, true
    69  	}
    70  	return nil, false
    71  }
    72  
    73  func (m *MultiTenantAggregator[T]) aggregatorForTenant(tenantID string) (*tenantAggregator[T], bool) {
    74  	window := m.limits.DistributorAggregationWindow(tenantID)
    75  	period := m.limits.DistributorAggregationPeriod(tenantID)
    76  	if window <= 0 || period <= 0 {
    77  		return nil, false
    78  	}
    79  	k := tenantKey{
    80  		tenantID: tenantID,
    81  		window:   window,
    82  		period:   period,
    83  	}
    84  	m.m.RLock()
    85  	t, ok := m.tenants[k]
    86  	m.m.RUnlock()
    87  	if ok {
    88  		return t, true
    89  	}
    90  
    91  	m.m.Lock()
    92  	defer m.m.Unlock()
    93  	if t, ok = m.tenants[k]; ok {
    94  		return t, true
    95  	}
    96  
    97  	labels := prometheus.Labels{
    98  		"tenant_id":          tenantID,
    99  		"aggregation_window": window.String(),
   100  		"aggregation_period": period.String(),
   101  	}
   102  
   103  	a := NewAggregator[T](time.Duration(window), time.Duration(period))
   104  	const metricNamePrefix = "pyroscope_distributor_aggregation_"
   105  	t = &tenantAggregator[T]{
   106  		key:        k,
   107  		registerer: prometheus.WrapRegistererWith(labels, m.registerer),
   108  		aggregator: a,
   109  		collector:  NewAggregatorCollector(a, metricNamePrefix),
   110  	}
   111  
   112  	t.registerer.MustRegister(t.collector)
   113  	go t.aggregator.Start()
   114  	m.tenants[k] = t
   115  	return t, true
   116  }
   117  
   118  func (m *MultiTenantAggregator[T]) starting(_ context.Context) error { return nil }
   119  
   120  func (m *MultiTenantAggregator[T]) stopping(_ error) error {
   121  	m.closeOnce.Do(func() {
   122  		close(m.stop)
   123  		<-m.done
   124  	})
   125  	return nil
   126  }
   127  
   128  const maxTenantAge = time.Second * 10
   129  
   130  func (m *MultiTenantAggregator[T]) running(ctx context.Context) error {
   131  	t := time.NewTicker(maxTenantAge)
   132  	defer func() {
   133  		t.Stop()
   134  		close(m.done)
   135  	}()
   136  	for {
   137  		select {
   138  		case <-t.C:
   139  			m.removeStaleTenants()
   140  		case <-m.stop:
   141  			return nil
   142  		case <-ctx.Done():
   143  			return nil
   144  		}
   145  	}
   146  }
   147  
   148  func (m *MultiTenantAggregator[T]) removeStaleTenants() {
   149  	stale := make([]*tenantAggregator[T], len(m.tenants)/8)
   150  	m.m.Lock()
   151  	for _, v := range m.tenants {
   152  		if time.Since(v.lastSeen.Load()) > maxTenantAge {
   153  			stale = append(stale, v)
   154  		}
   155  	}
   156  	m.m.Unlock()
   157  	for _, v := range stale {
   158  		v.aggregator.Stop()
   159  		v.registerer.Unregister(v.collector)
   160  		delete(m.tenants, v.key)
   161  	}
   162  }