Feature: On Broker Init callback

bit-wiz · Dec 1, 2023 · 188d8c1 · 188d8c1
1 parent fdb858a
commit 188d8c1
Show file tree

Hide file tree

Showing 4 changed files with 52 additions and 14 deletions.
diff --git a/engine/broker.go b/engine/broker.go
@@ -12,6 +12,11 @@ func (e *Engine) initBroker() error {
 	if err != nil {
 		return err
 	}
+	for _, cb := range e.onBrokerInit {
+		if err := cb(broker); err != nil {
+			return err
+		}
+	}
 	e.broker = broker
 	return nil
 }

diff --git a/engine/default.go b/engine/default.go
@@ -57,6 +57,10 @@ func SubmitJob(ctx context.Context, ij *input.Job, listeners ...web.JobListener)
 	return defaultEngine.SubmitJob(ctx, ij, listeners...)
 }
 
+func OnBrokerInit(fn func(b mq.Broker) error) {
+	defaultEngine.OnBrokerInit(fn)
+}
+
 func Start() error {
 	return defaultEngine.Start()
 }

diff --git a/engine/engine.go b/engine/engine.go
@@ -41,20 +41,21 @@ const (
 )
 
 type Engine struct {
-	quit        chan os.Signal
-	terminate   chan any
-	terminated  chan any
-	cfg         Config
-	state       string
-	mu          sync.Mutex
-	broker      mq.Broker
-	ds          datastore.Datastore
-	mounters    map[string]*runtime.MultiMounter
-	runtime     runtime.Runtime
-	coordinator *coordinator.Coordinator
-	worker      *worker.Worker
-	dsProviders map[string]datastore.Provider
-	mqProviders map[string]mq.Provider
+	quit         chan os.Signal
+	terminate    chan any
+	terminated   chan any
+	cfg          Config
+	state        string
+	mu           sync.Mutex
+	broker       mq.Broker
+	ds           datastore.Datastore
+	mounters     map[string]*runtime.MultiMounter
+	runtime      runtime.Runtime
+	coordinator  *coordinator.Coordinator
+	worker       *worker.Worker
+	dsProviders  map[string]datastore.Provider
+	mqProviders  map[string]mq.Provider
+	onBrokerInit []func(b mq.Broker) error
 }
 
 type Config struct {
@@ -337,6 +338,13 @@ func (e *Engine) SubmitJob(ctx context.Context, ij *input.Job, listeners ...web.
 	return job.Clone(), nil
 }
 
+func (e *Engine) OnBrokerInit(fn func(b mq.Broker) error) {
+	e.mu.Lock()
+	defer e.mu.Unlock()
+	e.mustState(StateIdle)
+	e.onBrokerInit = append(e.onBrokerInit, fn)
+}
+
 func (e *Engine) awaitTerm() {
 	signal.Notify(e.quit, os.Interrupt, syscall.SIGINT, syscall.SIGTERM)
 	select {

diff --git a/engine/engine_test.go b/engine/engine_test.go
@@ -255,3 +255,24 @@ func TestRegisterDatastoreProvider(t *testing.T) {
 	err = eng.Terminate()
 	assert.NoError(t, err)
 }
+
+func TestOnBrokerInit(t *testing.T) {
+	eng := New(Config{Mode: ModeStandalone})
+	assert.Equal(t, StateIdle, eng.state)
+
+	c := make(chan any)
+	eng.OnBrokerInit(func(b mq.Broker) error {
+		assert.NotNil(t, b)
+		close(c)
+		return nil
+	})
+
+	err := eng.Start()
+	assert.NoError(t, err)
+	assert.Equal(t, StateRunning, eng.state)
+
+	<-c
+
+	err = eng.Terminate()
+	assert.NoError(t, err)
+}