github.com/grafana/pyroscope@v1.18.0/pkg/util/extprom/tx_gauge.go (about)

     1  // SPDX-License-Identifier: AGPL-3.0-only
     2  // Provenance-includes-location: https://github.com/thanos-io/thanos/blob/main/pkg/extprom/tx_gauge.go
     3  // Provenance-includes-license: Apache-2.0
     4  // Provenance-includes-copyright: The Thanos Authors.
     5  
     6  package extprom
     7  
     8  import (
     9  	"sync"
    10  
    11  	"github.com/prometheus/client_golang/prometheus"
    12  	"github.com/prometheus/client_golang/prometheus/promauto"
    13  )
    14  
    15  type TxGaugeVec struct {
    16  	current      *prometheus.GaugeVec
    17  	mtx          sync.Mutex
    18  	newMetricVal func() *prometheus.GaugeVec
    19  
    20  	tx *prometheus.GaugeVec
    21  }
    22  
    23  // NewTxGaugeVec is a prometheus.GaugeVec that allows to start atomic metric value transaction.
    24  // It might be useful if long process that wants to update a GaugeVec but wants to build/accumulate those metrics
    25  // in a concurrent way without exposing partial state to Prometheus.
    26  // Caller can also use this as normal GaugeVec.
    27  //
    28  // Additionally it allows to init LabelValues on each transaction.
    29  // NOTE: This is quite naive implementation creating new prometheus.GaugeVec on each `ResetTx`, use wisely.
    30  func NewTxGaugeVec(reg prometheus.Registerer, opts prometheus.GaugeOpts, labelNames []string, initLabelValues ...[]string) *TxGaugeVec {
    31  	// Nil as we will register it on our own later.
    32  	f := func() *prometheus.GaugeVec {
    33  		g := promauto.With(nil).NewGaugeVec(opts, labelNames)
    34  		for _, vals := range initLabelValues {
    35  			g.WithLabelValues(vals...)
    36  		}
    37  		return g
    38  	}
    39  	tx := &TxGaugeVec{
    40  		current:      f(),
    41  		newMetricVal: f,
    42  	}
    43  	if reg != nil {
    44  		reg.MustRegister(tx)
    45  	}
    46  	return tx
    47  }
    48  
    49  // ResetTx starts new transaction. Not goroutine-safe.
    50  func (tx *TxGaugeVec) ResetTx() {
    51  	tx.tx = tx.newMetricVal()
    52  }
    53  
    54  // Submit atomically and fully applies new values from existing transaction GaugeVec. Not goroutine-safe.
    55  func (tx *TxGaugeVec) Submit() {
    56  	if tx.tx == nil {
    57  		return
    58  	}
    59  
    60  	tx.mtx.Lock()
    61  	tx.current = tx.tx
    62  	tx.mtx.Unlock()
    63  }
    64  
    65  // Describe is used in Register.
    66  func (tx *TxGaugeVec) Describe(ch chan<- *prometheus.Desc) {
    67  	tx.mtx.Lock()
    68  	defer tx.mtx.Unlock()
    69  
    70  	tx.current.Describe(ch)
    71  }
    72  
    73  // Collect is used by Registered.
    74  func (tx *TxGaugeVec) Collect(ch chan<- prometheus.Metric) {
    75  	tx.mtx.Lock()
    76  	defer tx.mtx.Unlock()
    77  
    78  	tx.current.Collect(ch)
    79  }
    80  
    81  // With works as GetMetricWith, but panics where GetMetricWithLabels would have
    82  // returned an error. Not returning an error allows shortcuts like
    83  //
    84  //	myVec.With(prometheus.Labels{"code": "404", "method": "GET"}).Add(42)
    85  func (tx *TxGaugeVec) With(labels prometheus.Labels) prometheus.Gauge {
    86  	if tx.tx == nil {
    87  		tx.ResetTx()
    88  	}
    89  	return tx.tx.With(labels)
    90  }
    91  
    92  // WithLabelValues works as GetMetricWithLabelValues, but panics where
    93  // GetMetricWithLabelValues would have returned an error. Not returning an
    94  // error allows shortcuts like
    95  //
    96  //	myVec.WithLabelValues("404", "GET").Add(42)
    97  func (tx *TxGaugeVec) WithLabelValues(lvs ...string) prometheus.Gauge {
    98  	if tx.tx == nil {
    99  		tx.ResetTx()
   100  	}
   101  	return tx.tx.WithLabelValues(lvs...)
   102  }