Make Execute() wait for initialization of status gRPC server before use

plugin/executor.go

@@ -37,27 +37,42 @@ func (p *ExecutorPlugin) GRPCClient(ctx context.Context, broker *plugin.GRPCBrok
 	return &ExecutorClient{client: types.NewExecutorClient(c), broker: broker}, nil
 }
 
+type Broker interface {
+	NextId() uint32
+	AcceptAndServe(id uint32, s func([]grpc.ServerOption) *grpc.Server)
+}
+
 // Here is the gRPC client that GRPCClient talks to.
 type ExecutorClient struct {
 	// This is the real implementation
 	client types.ExecutorClient
-	broker *plugin.GRPCBroker
+	broker Broker
 }
 
 func (m *ExecutorClient) Execute(args *types.ExecuteRequest, cb StatusHelper) (*types.ExecuteResponse, error) {
 	// This is where the magic conversion to Proto happens
 	statusHelperServer := &GRPCStatusHelperServer{Impl: cb}
 
+	initChan := make(chan bool, 1)
 	var s *grpc.Server
 	serverFunc := func(opts []grpc.ServerOption) *grpc.Server {
 		s = grpc.NewServer(opts...)
 		types.RegisterStatusHelperServer(s, statusHelperServer)
+		initChan <- true
 
 		return s
 	}
 
 	brokerID := m.broker.NextId()
-	go m.broker.AcceptAndServe(brokerID, serverFunc)
+	go func() {
+		m.broker.AcceptAndServe(brokerID, serverFunc)
+		// AcceptAndServe might terminate without calling serverFunc
+		// To prevent eternal blocking, send 'init done' signal
+		initChan <- true
+	}()
+
+	// Wait for s to be initialized in the goroutine
+	<-initChan
 
 	args.StatusServer = brokerID
 	r, err := m.client.Execute(context.Background(), args)

plugin/executor_test.go

@@ -0,0 +1,47 @@
+package plugin
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+
+	"github.com/distribworks/dkron/v3/plugin/types"
+	dktypes "github.com/distribworks/dkron/v3/plugin/types"
+	grpc "google.golang.org/grpc"
+)
+
+type MockedExecutor struct{}
+
+func (m *MockedExecutor) Execute(ctx context.Context, in *types.ExecuteRequest, opts ...grpc.CallOption) (*types.ExecuteResponse, error) {
+	resp := &dktypes.ExecuteResponse{}
+	return resp, nil
+}
+
+type MockedStatusHelper struct{}
+
+func (m MockedStatusHelper) Update([]byte, bool) (int64, error) {
+	return 0, nil
+}
+
+type MockedBroker struct{}
+
+func (m *MockedBroker) AcceptAndServe(id uint32, s func([]grpc.ServerOption) *grpc.Server) {
+}
+
+func (m *MockedBroker) NextId() uint32 {
+	return 0
+}
+
+func TestExecuteDoesNotPanicIfGRPCIsNotInitializedOnTime(t *testing.T) {
+	var brokerMock MockedBroker
+	var execMock MockedExecutor
+	execClient := ExecutorClient{
+		client: &execMock,
+		broker: &brokerMock,
+	}
+
+	var requestStub dktypes.ExecuteRequest
+	var statusHelperMock MockedStatusHelper
+	assert.NotPanics(t, func() { execClient.Execute(&requestStub, statusHelperMock) })
+}