Use A Publisher Everwhere

config/armada/config.yaml

@@ -80,9 +80,9 @@ pulsar:
   maxConnectionsPerBroker: 1
   compressionType: zlib
   compressionLevel: faster
-  eventsPrinter: false
   eventsPrinterSubscription: "EventsPrinter"
-  maxAllowedMessageSize: 4194304 # 4MB
+  maxAllowedMessageSize: 4194304 # 1MB
+  sendTimeout: 5s
   receiverQueueSize: 100
 postgres:
   connection:

config/eventingester/config.yaml

@@ -7,7 +7,6 @@ redis:
 pulsar:
   URL: pulsar://pulsar:6650
   jobsetEventsTopic: events
-  receiveTimeout: 5s
   backoffTime: 1s
   receiverQueueSize: 100
 subscriptionName: "events-ingester"

config/lookoutingesterv2/config.yaml

@@ -10,7 +10,6 @@ metricsPort: 9002
 pulsar:
   URL: "pulsar://pulsar:6650"
   jobsetEventsTopic: "events"
-  receiveTimeout: 5s
   backoffTime: 1s
   receiverQueueSize: 100
 subscriptionName: "lookout-ingester-v2"

config/scheduler/config.yaml

@@ -3,7 +3,6 @@ schedulePeriod: 10s
 maxSchedulingDuration: 5s
 executorTimeout: 1h
 databaseFetchSize: 1000
-pulsarSendTimeout: 5s
 internedStringsCacheSize: 100000
 queueRefreshPeriod: 1s
 disableSubmitCheck: false
@@ -38,6 +37,7 @@ pulsar:
   compressionType: zlib
   compressionLevel: faster
   maxAllowedMessageSize: 4194304 #4Mi
+  sendTimeout: 5s
 redis:
   addrs:
     - redis:6379

config/scheduleringester/config.yaml

@@ -11,7 +11,6 @@ metrics:
 pulsar:
   URL: "pulsar://localhost:6650"
   jobsetEventsTopic: "events"
-  receiveTimeout: 5s
   backoffTime: 1s
   receiverQueueSize: 100
 subscriptionName: "scheduler-ingester"

internal/armada/configuration/types.go

@@ -72,8 +72,8 @@ type PulsarConfig struct {
 	DedupTable string
 	// Maximum allowed message size in bytes
 	MaxAllowedMessageSize uint
-	// Timeout when polling pulsar for messages
-	ReceiveTimeout time.Duration
+	// Timeout when sending messages asynchronously
+	SendTimeout time.Duration
 	// Backoff from polling when Pulsar returns an error
 	BackoffTime time.Duration
 	// Number of pulsar messages that will be queued by the pulsar consumer.

internal/armada/server.go

@@ -149,7 +149,7 @@ func Serve(ctx *armadacontext.Context, config *configuration.ArmadaConfig, healt
 		CompressionLevel: config.Pulsar.CompressionLevel,
 		BatchingMaxSize:  config.Pulsar.MaxAllowedMessageSize,
 		Topic:            config.Pulsar.JobsetEventsTopic,
-	}, config.Pulsar.MaxAllowedMessageSize)
+	}, config.Pulsar.SendTimeout, config.Pulsar.MaxAllowedMessageSize)
 	if err != nil {
 		return errors.Wrapf(err, "error creating pulsar producer")
 	}

internal/armada/submit/submit_test.go

@@ -16,14 +16,15 @@ import (
 	"github.com/armadaproject/armada/internal/armada/submit/testfixtures"
 	"github.com/armadaproject/armada/internal/common/armadacontext"
 	"github.com/armadaproject/armada/internal/common/auth/permission"
+	commonmocks "github.com/armadaproject/armada/internal/common/mocks"
 	"github.com/armadaproject/armada/internal/common/util"
 	"github.com/armadaproject/armada/pkg/api"
 	"github.com/armadaproject/armada/pkg/armadaevents"
 	"github.com/armadaproject/armada/pkg/client/queue"
 )
 
 type mockObjects struct {
-	publisher    *mocks.MockPublisher
+	publisher    *commonmocks.MockPublisher
 	queueRepo    *mocks.MockQueueRepository
 	deduplicator *mocks.MockDeduplicator
 	authorizer   *mocks.MockActionAuthorizer
@@ -32,7 +33,7 @@ type mockObjects struct {
 func createMocks(t *testing.T) *mockObjects {
 	ctrl := gomock.NewController(t)
 	return &mockObjects{
-		publisher:    mocks.NewMockPublisher(ctrl),
+		publisher:    commonmocks.NewMockPublisher(ctrl),
 		queueRepo:    mocks.NewMockQueueRepository(ctrl),
 		deduplicator: mocks.NewMockDeduplicator(ctrl),
 		authorizer:   mocks.NewMockActionAuthorizer(ctrl),

internal/common/ingest/ingestion_pipeline_test.go

@@ -281,8 +281,7 @@ func TestRun_HappyPath_MultipleMessages(t *testing.T) {
 func testPipeline(consumer pulsar.Consumer, converter InstructionConverter[*simpleMessages], sink Sink[*simpleMessages]) *IngestionPipeline[*simpleMessages] {
 	return &IngestionPipeline[*simpleMessages]{
 		pulsarConfig: configuration.PulsarConfig{
-			ReceiveTimeout: 10 * time.Second,
-			BackoffTime:    time.Second,
+			BackoffTime: time.Second,
 		},
 		pulsarSubscriptionName: "subscription",
 		pulsarBatchDuration:    batchDuration,

internal/common/mocks/generate.go

@@ -3,3 +3,4 @@ package mocks
 // Mock implementations used by tests
 //go:generate mockgen -destination=./mock_pulsar.go -package=mocks "github.com/apache/pulsar-client-go/pulsar" Client,Producer,Message
 //go:generate mockgen -destination=./mock_executorapi.go -package=mocks "github.com/armadaproject/armada/pkg/executorapi" ExecutorApiClient,ExecutorApi_LeaseJobRunsClient
+//go:generate mockgen -destination=./mock_publisher.go -package=mocks "github.com/armadaproject/armada/internal/common/pulsarutils" Publisher

internal/common/pulsarutils/publisher.go

@@ -1,23 +1,31 @@
 package pulsarutils
 
 import (
+	"sync"
+	"time"
+
 	"github.com/apache/pulsar-client-go/pulsar"
+	"github.com/gogo/protobuf/proto"
 	"github.com/pkg/errors"
 
 	"github.com/armadaproject/armada/internal/common/armadacontext"
+	"github.com/armadaproject/armada/internal/common/eventutil"
+	"github.com/armadaproject/armada/internal/common/logging"
 	"github.com/armadaproject/armada/pkg/armadaevents"
 )
 
 // Publisher is an interface to be implemented by structs that handle publishing messages to pulsar
 type Publisher interface {
-	PublishMessages(ctx *armadacontext.Context, es *armadaevents.EventSequence) error
+	PublishMessages(ctx *armadacontext.Context, events ...*armadaevents.EventSequence) error
 	Close()
 }
 
 // PulsarPublisher is the default implementation of Publisher
 type PulsarPublisher struct {
 	// Used to send messages to pulsar
 	producer pulsar.Producer
+	// Timeout after which async messages sends will be considered failed
+	sendTimeout time.Duration
 	// Maximum size (in bytes) of produced pulsar messages.
 	// This must be below 4MB which is the pulsar message size limit
 	maxAllowedMessageSize uint
@@ -26,6 +34,7 @@ type PulsarPublisher struct {
 func NewPulsarPublisher(
 	pulsarClient pulsar.Client,
 	producerOptions pulsar.ProducerOptions,
+	sendTimeout time.Duration,
 	maxAllowedMessageSize uint,
 ) (*PulsarPublisher, error) {
 	producer, err := pulsarClient.CreateProducer(producerOptions)
@@ -34,18 +43,54 @@ func NewPulsarPublisher(
 	}
 	return &PulsarPublisher{
 		producer:              producer,
+		sendTimeout:           sendTimeout,
 		maxAllowedMessageSize: maxAllowedMessageSize,
 	}, nil
 }
 
 // PublishMessages publishes all event sequences to pulsar. Event sequences for a given jobset will be combined into
 // single event sequences up to maxMessageBatchSize.
-func (p *PulsarPublisher) PublishMessages(ctx *armadacontext.Context, es *armadaevents.EventSequence) error {
-	return CompactAndPublishSequences(
-		ctx,
-		[]*armadaevents.EventSequence{es},
-		p.producer,
-		p.maxAllowedMessageSize)
+func (p *PulsarPublisher) PublishMessages(ctx *armadacontext.Context, events ...*armadaevents.EventSequence) error {
+	sequences := eventutil.CompactEventSequences(events)
+	sequences, err := eventutil.LimitSequencesByteSize(sequences, p.maxAllowedMessageSize, true)
+	if err != nil {
+		return err
+	}
+	msgs := make([]*pulsar.ProducerMessage, len(sequences))
+	for i, sequence := range sequences {
+		bytes, err := proto.Marshal(sequence)
+		if err != nil {
+			return err
+		}
+		msgs[i] = &pulsar.ProducerMessage{
+			Payload: bytes,
+			Key:     sequences[i].JobSetName,
+		}
+	}
+
+	wg := sync.WaitGroup{}
+	wg.Add(len(msgs))
+
+	// Send messages
+	sendCtx, cancel := armadacontext.WithTimeout(ctx, p.sendTimeout)
+	errored := false
+	for _, msg := range msgs {
+		p.producer.SendAsync(sendCtx, msg, func(_ pulsar.MessageID, _ *pulsar.ProducerMessage, err error) {
+			if err != nil {
+				logging.
+					WithStacktrace(ctx, err).
+					Error("error sending message to Pulsar")
+				errored = true
+			}
+			wg.Done()
+		})
+	}
+	wg.Wait()
+	cancel()
+	if errored {
+		return errors.New("One or more messages failed to send to Pulsar")
+	}
+	return nil
 }
 
 func (p *PulsarPublisher) Close() {