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

     1  // This directory was copied and adapted from https://github.com/grafana/agent/tree/main/pkg/metrics.
     2  // We cannot vendor the agent in since the agent vendors loki in, which would cause a cyclic dependency.
     3  // NOTE: many changes have been made to the original code for our use-case.
     4  // This directory was copied and adapted from https://github.com/grafana/agent/tree/main/pkg/metrics.
     5  // We cannot vendor the agent in since the agent vendors loki in, which would cause a cyclic dependency.
     6  // NOTE: many changes have been made to the original code for our use-case.
     7  package util
     8  
     9  import "github.com/prometheus/client_golang/prometheus"
    10  
    11  // Unregisterer is a Prometheus Registerer that can unregister all collectors
    12  // passed to it.
    13  type Unregisterer struct {
    14  	wrap prometheus.Registerer
    15  	cs   map[prometheus.Collector]struct{}
    16  }
    17  
    18  // WrapWithUnregisterer wraps a prometheus Registerer with capabilities to
    19  // unregister all collectors.
    20  func WrapWithUnregisterer(reg prometheus.Registerer) *Unregisterer {
    21  	return &Unregisterer{
    22  		wrap: reg,
    23  		cs:   make(map[prometheus.Collector]struct{}),
    24  	}
    25  }
    26  
    27  // Register implements prometheus.Registerer.
    28  func (u *Unregisterer) Register(c prometheus.Collector) error {
    29  	if u.wrap == nil {
    30  		return nil
    31  	}
    32  
    33  	err := u.wrap.Register(c)
    34  	if err != nil {
    35  		return err
    36  	}
    37  	u.cs[c] = struct{}{}
    38  	return nil
    39  }
    40  
    41  // MustRegister implements prometheus.Registerer.
    42  func (u *Unregisterer) MustRegister(cs ...prometheus.Collector) {
    43  	for _, c := range cs {
    44  		if err := u.Register(c); err != nil {
    45  			panic(err)
    46  		}
    47  	}
    48  }
    49  
    50  // Unregister implements prometheus.Registerer.
    51  func (u *Unregisterer) Unregister(c prometheus.Collector) bool {
    52  	if u.wrap != nil && u.wrap.Unregister(c) {
    53  		delete(u.cs, c)
    54  		return true
    55  	}
    56  	return false
    57  }
    58  
    59  // UnregisterAll unregisters all collectors that were registered through the
    60  // Reigsterer.
    61  func (u *Unregisterer) UnregisterAll() bool {
    62  	success := true
    63  	for c := range u.cs {
    64  		if !u.Unregister(c) {
    65  			success = false
    66  		}
    67  	}
    68  	return success
    69  }