diff --git a/cmd/stream/main.go b/cmd/stream/main.go index 17e2b6fe12..37174f32af 100644 --- a/cmd/stream/main.go +++ b/cmd/stream/main.go @@ -85,8 +85,9 @@ func main() { } server := &http.Server{ - Addr: cfg.ServeURL, - Handler: handler, + Addr: cfg.ServeURL, + Handler: handler, + ReadHeaderTimeout: 15 * time.Second, } var wg sync.WaitGroup @@ -107,8 +108,21 @@ func main() { config := queue.ParseConfig(cfg.QueueConfig, "stream", cfg.Region) logger.Info("using queue config", zap.Any("config", config)) - consumer := queue.NewConsumer(ctx, queue.TopicTranscodeRequest, "stream", config) - producer := queue.NewProducer(ctx, queue.TopicTranscodeResult, config) + + consumerOptions := queue.ConsumerOptions{ + Topic: queue.TopicTranscodeRequest, + Group: "stream", + Config: config, + } + consumer := queue.NewConsumer(ctx, consumerOptions) + + producerOptions := queue.ProducerOptions{ + Topic: queue.TopicTranscodeResult, + Group: "stream", + Config: config, + RetryCount: 3, + } + producer := queue.NewProducer(ctx, producerOptions) wg.Add(1) go func() { diff --git a/internal/pkg/queue/queue.go b/internal/pkg/queue/queue.go index 9786581529..64389e8335 100644 --- a/internal/pkg/queue/queue.go +++ b/internal/pkg/queue/queue.go @@ -45,6 +45,13 @@ type TranscodeResult struct { Thumbnail string } +// ConsumerOptions represents options for the consumer +type ConsumerOptions struct { + Topic string + Group string + Config Config +} + // Consumer provides a consumer interface to a Kafka queue type Consumer interface { Read(ctx context.Context) (kafka.Message, error) @@ -53,6 +60,14 @@ type Consumer interface { Close() error } +// ProducerOptions represents options for the producer +type ProducerOptions struct { + Topic string + Group string + Config Config + RetryCount int +} + // Producer provides a producer interface to a Kafka queue type Producer interface { Send(ctx context.Context, workspaceID string, data any) error @@ -61,14 +76,16 @@ type Producer interface { // TConsumer implements Consumer interface type TConsumer struct { - topic string - reader *kafka.Reader + topic string + reader *kafka.Reader + options ConsumerOptions } // TProducer implements Producer interface type TProducer struct { - topic string - writer *kafka.Writer + topic string + writer *kafka.Writer + options ProducerOptions } // Logger is kafka.Logger implementation @@ -93,12 +110,12 @@ func (l *Logger) Printf(msg string, args ...any) { } // NewConsumer creates a new transcoding request consumer -func NewConsumer(ctx context.Context, topic, group string, config Config) Consumer { - platformTopic := makeTopicID(topic, config) - groupID := makeGroupID(group, topic, config) +func NewConsumer(ctx context.Context, options ConsumerOptions) Consumer { + platformTopic := makeTopicID(options.Topic, options.Config) + groupID := makeGroupID(options.Group, options.Topic, options.Config) reader := kafka.NewReader(kafka.ReaderConfig{ - Brokers: config.Brokers, + Brokers: options.Config.Brokers, GroupID: groupID, Topic: platformTopic, Logger: NewLogger(ctx), @@ -107,8 +124,9 @@ func NewConsumer(ctx context.Context, topic, group string, config Config) Consum }) return &TConsumer{ - topic: platformTopic, - reader: reader, + options: options, + topic: platformTopic, + reader: reader, } } @@ -133,17 +151,18 @@ func (c *TConsumer) Close() error { } // NewProducer creates a new transcoding result producer -func NewProducer(ctx context.Context, topic string, config Config) Producer { - platformTopic := makeTopicID(topic, config) +func NewProducer(ctx context.Context, options ProducerOptions) Producer { + platformTopic := makeTopicID(options.Topic, options.Config) writer := &kafka.Writer{ - Addr: kafka.TCP(config.Brokers...), + Addr: kafka.TCP(options.Config.Brokers...), Topic: platformTopic, AllowAutoTopicCreation: true, Logger: NewLogger(ctx), } return &TProducer{ - topic: platformTopic, - writer: writer, + options: options, + topic: platformTopic, + writer: writer, } } @@ -154,13 +173,20 @@ func (p *TProducer) Send(ctx context.Context, workspaceID string, data any) erro return err } - return p.writer.WriteMessages( - ctx, - kafka.Message{ - Key: []byte(workspaceID), - Value: value, - }, - ) + for range p.options.RetryCount { + err = p.writer.WriteMessages( + ctx, + kafka.Message{ + Key: []byte(workspaceID), + Value: value, + }, + ) + if err == nil { + break + } + } + + return err } // Close closes the producer. diff --git a/internal/pkg/queue/worker.go b/internal/pkg/queue/worker.go index 58f5506f46..0137f48be9 100644 --- a/internal/pkg/queue/worker.go +++ b/internal/pkg/queue/worker.go @@ -112,7 +112,7 @@ func (w *Worker) processMessage(ctx context.Context, msg kafka.Message, logger * res, err := transcoder.Transcode(ctx, &task) if err == nil { - if err := w.producer.Send(ctx, req.WorkspaceUUID, res); err != nil { + if err = w.producer.Send(ctx, req.WorkspaceUUID, res); err != nil { logger.Error("failed to send transcode result", zap.Error(err)) } }