diff --git a/metric/common.go b/metric/common.go index 996b2d1..35f6792 100644 --- a/metric/common.go +++ b/metric/common.go @@ -29,6 +29,7 @@ type MetricDescription struct { type Exporter interface { RegisterMetric(*MetricDescription) UpdateMetric(string, float64) + Shutdown() } type MetricWriter interface { diff --git a/metric/prometheus.go b/metric/prometheus.go index 9a32756..1be333e 100644 --- a/metric/prometheus.go +++ b/metric/prometheus.go @@ -1,6 +1,7 @@ package metric import ( + "context" "math" "sync/atomic" @@ -49,6 +50,7 @@ type prometheusExporter struct { writerChan chan *writeRequest registerChan chan *prometheusMetricDesc namespace string + cancel context.CancelFunc } func (pe *prometheusExporter) RegisterMetric(nm *MetricDescription) { @@ -64,6 +66,12 @@ func (pe *prometheusExporter) UpdateMetric(key string, val float64) { pe.writerChan <- &writeRequest{key: key, val: val} } +func (pe *prometheusExporter) Shutdown() { + if pe.cancel != nil { + pe.cancel() + } +} + type prometheusCollector struct { metrics map[string]*prometheusMetricDesc } @@ -83,7 +91,7 @@ func (pc *prometheusCollector) Collect(ch chan<- prometheus.Metric) { } } -func (pe *prometheusExporter) loop() { +func (pe *prometheusExporter) loop(ctx context.Context) { defer func() { r := recover() if r != nil { @@ -92,9 +100,17 @@ func (pe *prometheusExporter) loop() { }() var collector *prometheusCollector + defer func() { + if collector != nil { + prometheus.Unregister(collector) + } + }() for { select { + case <-ctx.Done(): + return + case req := <-pe.writerChan: if m := collector.metrics[req.key]; m != nil { atomic.StoreUint64(m.valptr, math.Float64bits(req.val)) @@ -128,12 +144,14 @@ func (pe *prometheusExporter) loop() { } func NewPrometheusExport(namespace string) Exporter { + ctx, cancel := context.WithCancel(context.Background()) exp := &prometheusExporter{ registerChan: make(chan *prometheusMetricDesc, 10), writerChan: make(chan *writeRequest, 100), namespace: namespace, + cancel: cancel, } - go exp.loop() + go exp.loop(ctx) return exp }