github.com/grafana/pyroscope@v1.18.0/pkg/metastore/raftnode/observer.go (about)

     1  package raftnode
     2  
     3  import (
     4  	"sync"
     5  
     6  	"github.com/go-kit/log"
     7  	"github.com/go-kit/log/level"
     8  	"github.com/hashicorp/raft"
     9  	"github.com/prometheus/client_golang/prometheus"
    10  )
    11  
    12  // StateHandler is called every time the
    13  // Raft state change is observed.
    14  type StateHandler interface {
    15  	Observe(raft.RaftState)
    16  }
    17  
    18  type Observer struct {
    19  	logger   log.Logger
    20  	raft     *raft.Raft
    21  	observer *raft.Observer
    22  	state    *prometheus.GaugeVec
    23  	mu       sync.Mutex
    24  	handlers []StateHandler
    25  	c        chan raft.Observation
    26  	stop     chan struct{}
    27  	done     chan struct{}
    28  }
    29  
    30  func NewRaftStateObserver(logger log.Logger, r *raft.Raft, state *prometheus.GaugeVec) *Observer {
    31  	o := &Observer{
    32  		logger: logger,
    33  		raft:   r,
    34  		c:      make(chan raft.Observation, 1),
    35  		stop:   make(chan struct{}),
    36  		done:   make(chan struct{}),
    37  		state:  state,
    38  	}
    39  	level.Debug(o.logger).Log("msg", "registering raft state observer")
    40  	o.observer = raft.NewObserver(o.c, true, func(o *raft.Observation) bool {
    41  		_, ok := o.Data.(raft.RaftState)
    42  		return ok
    43  	})
    44  	r.RegisterObserver(o.observer)
    45  	o.updateRaftState()
    46  	go o.run()
    47  	return o
    48  }
    49  
    50  func (o *Observer) RegisterHandler(h StateHandler) {
    51  	o.mu.Lock()
    52  	o.handlers = append(o.handlers, h)
    53  	o.mu.Unlock()
    54  	o.updateRaftState()
    55  }
    56  
    57  func (o *Observer) Deregister() {
    58  	level.Debug(o.logger).Log("msg", "deregistering raft observer")
    59  	o.raft.DeregisterObserver(o.observer)
    60  	close(o.stop)
    61  	<-o.done
    62  }
    63  
    64  func (o *Observer) run() {
    65  	defer func() {
    66  		close(o.done)
    67  	}()
    68  	for {
    69  		select {
    70  		case <-o.c:
    71  			o.updateRaftState()
    72  		case <-o.stop:
    73  			return
    74  		}
    75  	}
    76  }
    77  
    78  func (o *Observer) updateRaftState() {
    79  	state := o.raft.State()
    80  	level.Debug(o.logger).Log("msg", "raft state changed", "raft_state", state)
    81  	if o.state != nil {
    82  		o.state.Reset()
    83  		o.state.WithLabelValues(state.String()).Set(1)
    84  	}
    85  	o.mu.Lock()
    86  	handlers := o.handlers
    87  	o.mu.Unlock()
    88  	for _, h := range handlers {
    89  		h.Observe(state)
    90  	}
    91  }