Allow setting leadership extra functions

distribworks · Jul 8, 2019 · c6da283 · c6da283
1 parent 0c38604
commit c6da283
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 0 deletions.
diff --git a/dkron/agent.go b/dkron/agent.go
@@ -86,6 +86,9 @@ type Agent struct {
 	peers      map[string][]*serverParts
 	localPeers map[raft.ServerAddress]*serverParts
 	peerLock   sync.RWMutex
+
+	establishLeadershipFuncs []func() error
+	revokeLeadershipFuncs    []func() error
 }
 
 // ProcessorFactory is a function type that creates a new instance

diff --git a/dkron/leader.go b/dkron/leader.go
@@ -188,6 +188,12 @@ func (a *Agent) establishLeadership(stopCh chan struct{}) error {
 	defer metrics.MeasureSince([]string{"dkron", "leader", "establish_leadership"}, time.Now())
 	a.schedule()
 
+	for _, f := range a.establishLeadershipFuncs {
+		if err := f(); err != nil {
+			return err
+		}
+	}
+
 	return nil
 }
 
@@ -197,6 +203,12 @@ func (a *Agent) revokeLeadership() error {
 	defer metrics.MeasureSince([]string{"dkron", "leader", "revoke_leadership"}, time.Now())
 	a.sched.Stop()
 
+	for _, f := range a.revokeLeadershipFuncs {
+		if err := f(); err != nil {
+			return err
+		}
+	}
+
 	return nil
 }