refactor: Allow pro appliers

dkron/agent.go

@@ -70,6 +70,7 @@ type Agent struct {
 	// Pro features
 	GlobalLock         bool
 	MemberEventHandler func(serf.Event)
+	ProAppliers        LogAppliers
 
 	serf        *serf.Serf
 	config      *Config
@@ -295,7 +296,7 @@ func (a *Agent) setupRaft() error {
 			if err != nil {
 				return fmt.Errorf("recovery failed to parse peers.json: %v", err)
 			}
-			tmpFsm := newFSM(nil)
+			tmpFsm := newFSM(nil, nil)
 			if err := raft.RecoverCluster(config, tmpFsm,
 				logStore, stableStore, snapshots, transport, configuration); err != nil {
 				return fmt.Errorf("recovery failed: %v", err)
@@ -331,13 +332,14 @@ func (a *Agent) setupRaft() error {
 
 	// Instantiate the Raft systems. The second parameter is a finite state machine
 	// which stores the actual kv pairs and is operated upon through Apply().
-	fsm := newFSM(a.Store)
+	fsm := newFSM(a.Store, a.ProAppliers)
 	rft, err := raft.NewRaft(config, fsm, logStore, stableStore, snapshots, transport)
 	if err != nil {
 		return fmt.Errorf("new raft: %s", err)
 	}
 	a.leaderCh = rft.LeaderCh()
 	a.raft = rft
+
 	return nil
 }
 
@@ -451,18 +453,7 @@ func (a *Agent) StartServer() {
 	}
 
 	if a.Store == nil {
-		dirExists, err := exists(a.config.DataDir)
-		if err != nil {
-			log.WithError(err).WithField("dir", a.config.DataDir).Fatal("Invalid Dir")
-		}
-		if !dirExists {
-			// Try to create the directory
-			err := os.Mkdir(a.config.DataDir, 0700)
-			if err != nil {
-				log.WithError(err).WithField("dir", a.config.DataDir).Fatal("Error Creating Dir")
-			}
-		}
-		s, err := NewStore(a, filepath.Join(a.config.DataDir, a.config.NodeName))
+		s, err := NewStore(filepath.Join(a.config.DataDir, a.config.NodeName))
 		if err != nil {
 			log.WithError(err).Fatal("dkron: Error initializing store")
 		}
@@ -572,6 +563,11 @@ func (a *Agent) LocalMember() serf.Member {
 	return a.serf.LocalMember()
 }
 
+// Leader is used to return the Raft leader
+func (a *Agent) Leader() raft.ServerAddress {
+	return a.raft.Leader()
+}
+
 // Servers returns a list of known server
 func (a *Agent) Servers() (members []*ServerParts) {
 	for _, member := range a.serf.Members() {
@@ -880,3 +876,8 @@ func (a *Agent) applySetJob(job *proto.Job) error {
 
 	return nil
 }
+
+// RaftApply applies a command to the Raft log
+func (a *Agent) RaftApply(cmd []byte) raft.ApplyFuture {
+	return a.raft.Apply(cmd, raftTimeout)
+}

dkron/api.go

@@ -52,7 +52,7 @@ func (h *HTTPTransport) ServeHTTP() {
 }
 
 // APIRoutes registers the api routes on the gin RouterGroup.
-func (h *HTTPTransport) APIRoutes(r *gin.RouterGroup) {
+func (h *HTTPTransport) APIRoutes(r *gin.RouterGroup, middleware ...gin.HandlerFunc) {
 	r.GET("/debug/vars", expvar.Handler())
 
 	h.Engine.GET("/health", func(c *gin.Context) {
@@ -63,6 +63,7 @@ func (h *HTTPTransport) APIRoutes(r *gin.RouterGroup) {
 
 	r.GET("/v1", h.indexHandler)
 	v1 := r.Group("/v1")
+	v1.Use(middleware...)
 	v1.GET("/", h.indexHandler)
 	v1.GET("/members", h.membersHandler)
 	v1.GET("/leader", h.leaderHandler)
@@ -117,7 +118,11 @@ func (h *HTTPTransport) indexHandler(c *gin.Context) {
 func (h *HTTPTransport) jobsHandler(c *gin.Context) {
 	metadata := c.QueryMap("metadata")
 
-	jobs, err := h.agent.Store.GetJobs(&JobOptions{Metadata: metadata})
+	jobs, err := h.agent.Store.GetJobs(
+		&JobOptions{
+			Metadata: metadata,
+		},
+	)
 	if err != nil {
 		log.WithError(err).Error("api: Unable to get jobs, store not reachable.")
 		return

dkron/fsm.go

@@ -26,15 +26,26 @@ const (
 	ExecutionDoneType
 )
 
+// LogApplier is the definition of a function that can apply a Raft log
+type LogApplier func(buf []byte, index uint64) interface{}
+
+// LogAppliers is a mapping of the Raft MessageType to the appropriate log
+// applier
+type LogAppliers map[MessageType]LogApplier
+
 type dkronFSM struct {
 	store Storage
 	mu    sync.Mutex
+
+	// proAppliers holds the set of pro only LogAppliers
+	proAppliers LogAppliers
 }
 
 // NewFSM is used to construct a new FSM with a blank state
-func newFSM(store Storage) *dkronFSM {
+func newFSM(store Storage, logAppliers LogAppliers) *dkronFSM {
 	return &dkronFSM{
-		store: store,
+		store:       store,
+		proAppliers: logAppliers,
 	}
 }
 
@@ -56,6 +67,11 @@ func (d *dkronFSM) Apply(l *raft.Log) interface{} {
 		return d.applySetExecution(buf[1:])
 	}
 
+	// Check enterprise only message types.
+	if applier, ok := d.proAppliers[msgType]; ok {
+		return applier(buf[1:], l.Index)
+	}
+
 	return nil
 }
 

dkron/job.go

@@ -279,12 +279,7 @@ func (j *Job) GetStatus() string {
 }
 
 // GetParent returns the parent job of a job
-func (j *Job) GetParent() (*Job, error) {
-	// Maybe we are testing
-	if j.Agent == nil {
-		return nil, ErrNoAgent
-	}
-
+func (j *Job) GetParent(store *Store) (*Job, error) {
 	if j.Name == j.ParentJob {
 		return nil, ErrSameParent
 	}
@@ -293,7 +288,7 @@ func (j *Job) GetParent() (*Job, error) {
 		return nil, ErrNoParent
 	}
 
-	parentJob, err := j.Agent.Store.GetJob(j.ParentJob, nil)
+	parentJob, err := store.GetJob(j.ParentJob, nil)
 	if err != nil {
 		if err == badger.ErrKeyNotFound {
 			return nil, ErrParentJobNotFound

dkron/job_test.go

@@ -15,11 +15,9 @@ func TestJobGetParent(t *testing.T) {
 	require.NoError(t, err)
 	defer os.RemoveAll(dir)
 
-	a := &Agent{}
-	s, err := NewStore(a, dir)
+	s, err := NewStore(dir)
 	defer s.Shutdown()
 	require.NoError(t, err)
-	a.Store = s
 
 	parentTestJob := &Job{
 		Name:           "parent_test",
@@ -42,11 +40,11 @@ func TestJobGetParent(t *testing.T) {
 	err = s.SetJob(dependentTestJob, true)
 	assert.NoError(t, err)
 
-	parentTestJob, err = dependentTestJob.GetParent()
+	parentTestJob, err = dependentTestJob.GetParent(s)
 	assert.NoError(t, err)
 	assert.Equal(t, []string{dependentTestJob.Name}, parentTestJob.DependentJobs)
 
-	ptj, err := dependentTestJob.GetParent()
+	ptj, err := dependentTestJob.GetParent(s)
 	assert.NoError(t, err)
 	assert.Equal(t, parentTestJob.Name, ptj.Name)
 
@@ -60,7 +58,7 @@ func TestJobGetParent(t *testing.T) {
 	assert.NoError(t, err)
 	assert.Equal(t, "", dtj.ParentJob)
 
-	ptj, err = dtj.GetParent()
+	ptj, err = dtj.GetParent(s)
 	assert.EqualError(t, ErrNoParent, err.Error())
 
 	ptj, err = s.GetJob(parentTestJob.Name, nil)
@@ -110,9 +108,10 @@ func Test_isRunnable(t *testing.T) {
 	require.NoError(t, err)
 	defer os.RemoveAll(dir)
 
-	a := &Agent{}
-	s, err := NewStore(a, dir)
-	a.Store = s
+	s, err := NewStore(dir)
+	a := &Agent{
+		Store: s,
+	}
 	defer s.Shutdown()
 	require.NoError(t, err)
 

dkron/leader.go

@@ -192,7 +192,7 @@ func (a *Agent) establishLeadership(stopCh chan struct{}) error {
 	if err != nil {
 		log.Fatal(err)
 	}
-	a.sched.Start(jobs)
+	a.sched.Start(jobs, a)
 
 	return nil
 }

dkron/options.go

@@ -18,3 +18,10 @@ func WithTransportCredentials(tls *tls.Config) AgentOption {
 		agent.TLSConfig = tls
 	}
 }
+
+// WithStore set store in the agent
+func WithStore(store Storage) AgentOption {
+	return func(agent *Agent) {
+		agent.Store = store
+	}
+}

dkron/queries_test.go

@@ -31,7 +31,8 @@ func TestRunQuery(t *testing.T) {
 	c.BootstrapExpect = 1
 
 	a := NewAgent(c)
-	a.Start()
+	err = a.Start()
+	require.NoError(t, err)
 	time.Sleep(2 * time.Second)
 
 	// Test error with no job
@@ -45,8 +46,10 @@ func TestRunQuery(t *testing.T) {
 	err = a.Store.SetJob(j1, false)
 	require.NoError(t, err)
 
-	a.sched.Start([]*Job{j1})
+	a.sched.Start([]*Job{j1}, a)
 
 	_, err = a.RunQuery("test_job", &Execution{})
 	assert.NoError(t, err)
+
+	a.Stop()
 }

dkron/scheduler.go

@@ -6,8 +6,8 @@ import (
 	"strings"
 
 	"github.com/armon/go-metrics"
-	"github.com/robfig/cron/v3"
 	"github.com/distribworks/dkron/v2/extcron"
+	"github.com/robfig/cron/v3"
 	"github.com/sirupsen/logrus"
 )
 
@@ -39,11 +39,12 @@ func NewScheduler() *Scheduler {
 
 // Start the cron scheduler, adding its corresponding jobs and
 // executing them on time.
-func (s *Scheduler) Start(jobs []*Job) error {
+func (s *Scheduler) Start(jobs []*Job, agent *Agent) error {
 	s.Cron = cron.New(cron.WithParser(extcron.NewParser()))
 
 	metrics.IncrCounter([]string{"scheduler", "start"}, 1)
 	for _, job := range jobs {
+		job.Agent = agent
 		s.AddJob(job)
 	}
 	s.Cron.Start()
@@ -70,9 +71,9 @@ func (s *Scheduler) Stop() {
 }
 
 // Restart the scheduler
-func (s *Scheduler) Restart(jobs []*Job) {
+func (s *Scheduler) Restart(jobs []*Job, agent *Agent) {
 	s.Stop()
-	s.Start(jobs)
+	s.Start(jobs, agent)
 }
 
 // GetEntry returns a scheduler entry from a snapshot in

dkron/scheduler_test.go

@@ -20,7 +20,7 @@ func TestSchedule(t *testing.T) {
 		Owner:          "John Dough",
 		OwnerEmail:     "foo@bar.com",
 	}
-	sched.Start([]*Job{testJob1})
+	sched.Start([]*Job{testJob1}, &Agent{})
 
 	assert.True(t, sched.Started)
 	now := time.Now().Truncate(time.Second)
@@ -36,10 +36,11 @@ func TestSchedule(t *testing.T) {
 		Owner:          "John Dough",
 		OwnerEmail:     "foo@bar.com",
 	}
-	sched.Restart([]*Job{testJob2})
+	sched.Restart([]*Job{testJob2}, &Agent{})
 
 	assert.True(t, sched.Started)
 	assert.Len(t, sched.Cron.Entries(), 1)
+	sched.Stop()
 }
 
 func TestTimezoneAwareJob(t *testing.T) {
@@ -52,8 +53,9 @@ func TestTimezoneAwareJob(t *testing.T) {
 		Executor:       "shell",
 		ExecutorConfig: map[string]string{"command": "echo 'test1'", "shell": "true"},
 	}
-	sched.Start([]*Job{tzJob})
+	sched.Start([]*Job{tzJob}, &Agent{})
 
 	assert.True(t, sched.Started)
 	assert.Len(t, sched.Cron.Entries(), 1)
+	sched.Stop()
 }