fix(go): batch durable kafka replay
This commit is contained in:
@@ -40,6 +40,10 @@ type durableRecordFile struct {
|
||||
record durableRecord
|
||||
}
|
||||
|
||||
type recordPublishingSink interface {
|
||||
PublishRecords(context.Context, []durableRecord) error
|
||||
}
|
||||
|
||||
const durableReplayOperationTimeout = 30 * time.Second
|
||||
|
||||
func NewDurableSink(delegate Sink, cfg DurableConfig) *DurableSink {
|
||||
@@ -93,6 +97,24 @@ func (s *DurableSink) replay(ctx context.Context, limit int) error {
|
||||
records = append(records, durableRecordFile{path: file, record: record})
|
||||
}
|
||||
sortDurableRecords(records)
|
||||
if publisher, ok := s.delegate.(recordPublishingSink); ok {
|
||||
durableRecords := make([]durableRecord, 0, len(records))
|
||||
for _, item := range records {
|
||||
durableRecords = append(durableRecords, item.record)
|
||||
}
|
||||
if err := publisher.PublishRecords(ctx, durableRecords); err != nil {
|
||||
return err
|
||||
}
|
||||
for _, item := range records {
|
||||
if err := os.Remove(item.path); err != nil {
|
||||
return err
|
||||
}
|
||||
if item.record.Kind == "raw" {
|
||||
s.clearRawPending(item.record.Envelope)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
for _, item := range records {
|
||||
if err := s.publishRecord(ctx, item.record); err != nil {
|
||||
return err
|
||||
@@ -163,10 +185,14 @@ func (s *DurableSink) publishRecord(ctx context.Context, record durableRecord) e
|
||||
case "unified":
|
||||
return s.delegate.PublishUnified(ctx, record.Envelope)
|
||||
default:
|
||||
return fmt.Errorf("unknown durable record kind %q", record.Kind)
|
||||
return errUnknownRecordKind(record.Kind)
|
||||
}
|
||||
}
|
||||
|
||||
func errUnknownRecordKind(kind string) error {
|
||||
return fmt.Errorf("unknown durable record kind %q", kind)
|
||||
}
|
||||
|
||||
func (s *DurableSink) spool(kind string, env envelope.FrameEnvelope) error {
|
||||
if s.dir == "" {
|
||||
return fmt.Errorf("durable spool directory is empty")
|
||||
|
||||
@@ -139,6 +139,34 @@ func TestDurableSinkReplayLoopTickRespectsBatchSize(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestDurableSinkReplayUsesBatchPublisherWhenAvailable(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
env := durableTestEnvelope()
|
||||
writeDurableRecord(t, filepath.Join(dir, "0001-raw.json"), durableRecord{Kind: "raw", Envelope: env})
|
||||
writeDurableRecord(t, filepath.Join(dir, "0002-unified.json"), durableRecord{Kind: "unified", Envelope: env})
|
||||
delegate := &batchRecordingSink{}
|
||||
sink := NewDurableSink(delegate, DurableConfig{Directory: dir})
|
||||
|
||||
if err := sink.ReplayOnce(context.Background()); err != nil {
|
||||
t.Fatalf("ReplayOnce() error = %v", err)
|
||||
}
|
||||
if delegate.batchCalls != 1 {
|
||||
t.Fatalf("batch calls = %d, want 1", delegate.batchCalls)
|
||||
}
|
||||
if delegate.rawCalls != 0 || delegate.unifiedCalls != 0 {
|
||||
t.Fatalf("individual publish should not be used, raw=%d unified=%d", delegate.rawCalls, delegate.unifiedCalls)
|
||||
}
|
||||
if got, want := len(delegate.records), 2; got != want {
|
||||
t.Fatalf("batch record count = %d, want %d", got, want)
|
||||
}
|
||||
if delegate.records[0].Kind != "raw" || delegate.records[1].Kind != "unified" {
|
||||
t.Fatalf("batch order = %#v", delegate.records)
|
||||
}
|
||||
if files := spoolFiles(t, dir); len(files) != 0 {
|
||||
t.Fatalf("spool files after replay = %#v, want none", files)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDurableFilesFromReaderStopsAfterLimitedJSONBatch(t *testing.T) {
|
||||
reader := &fakeNameReader{
|
||||
batches: [][]string{
|
||||
@@ -298,3 +326,30 @@ func (s *contextCheckingReplaySink) PublishUnified(ctx context.Context, _ envelo
|
||||
func (s *contextCheckingReplaySink) Close() error {
|
||||
return nil
|
||||
}
|
||||
|
||||
type batchRecordingSink struct {
|
||||
batchCalls int
|
||||
records []durableRecord
|
||||
rawCalls int
|
||||
unifiedCalls int
|
||||
}
|
||||
|
||||
func (s *batchRecordingSink) PublishRecords(_ context.Context, records []durableRecord) error {
|
||||
s.batchCalls++
|
||||
s.records = append(s.records, records...)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *batchRecordingSink) PublishRaw(context.Context, envelope.FrameEnvelope) error {
|
||||
s.rawCalls++
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *batchRecordingSink) PublishUnified(context.Context, envelope.FrameEnvelope) error {
|
||||
s.unifiedCalls++
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *batchRecordingSink) Close() error {
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -76,6 +76,25 @@ func (s *KafkaSink) PublishUnified(ctx context.Context, env envelope.FrameEnvelo
|
||||
return s.publish(ctx, s.unifiedTopic, env)
|
||||
}
|
||||
|
||||
func (s *KafkaSink) PublishRecords(ctx context.Context, records []durableRecord) error {
|
||||
messages := make([]kafka.Message, 0, len(records))
|
||||
for _, record := range records {
|
||||
topic, err := s.topicForRecord(record)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
message, err := kafkaMessage(topic, record.Envelope)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
messages = append(messages, message)
|
||||
}
|
||||
if len(messages) == 0 {
|
||||
return nil
|
||||
}
|
||||
return s.writer.WriteMessages(ctx, messages...)
|
||||
}
|
||||
|
||||
func (s *KafkaSink) Close() error {
|
||||
if s == nil || s.writer == nil {
|
||||
return nil
|
||||
@@ -84,13 +103,36 @@ func (s *KafkaSink) Close() error {
|
||||
}
|
||||
|
||||
func (s *KafkaSink) publish(ctx context.Context, topic string, env envelope.FrameEnvelope) error {
|
||||
payload, err := env.MarshalJSONBytes()
|
||||
message, err := kafkaMessage(topic, env)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return s.writer.WriteMessages(ctx, kafka.Message{
|
||||
return s.writer.WriteMessages(ctx, message)
|
||||
}
|
||||
|
||||
func (s *KafkaSink) topicForRecord(record durableRecord) (string, error) {
|
||||
switch record.Kind {
|
||||
case "raw":
|
||||
topic, ok := s.rawTopics[record.Envelope.Protocol]
|
||||
if !ok || topic == "" {
|
||||
return "", fmt.Errorf("raw topic not configured for protocol %s", record.Envelope.Protocol)
|
||||
}
|
||||
return topic, nil
|
||||
case "unified":
|
||||
return s.unifiedTopic, nil
|
||||
default:
|
||||
return "", fmt.Errorf("unknown durable record kind %q", record.Kind)
|
||||
}
|
||||
}
|
||||
|
||||
func kafkaMessage(topic string, env envelope.FrameEnvelope) (kafka.Message, error) {
|
||||
payload, err := env.MarshalJSONBytes()
|
||||
if err != nil {
|
||||
return kafka.Message{}, err
|
||||
}
|
||||
return kafka.Message{
|
||||
Topic: topic,
|
||||
Key: env.KafkaKey(),
|
||||
Value: payload,
|
||||
})
|
||||
}, nil
|
||||
}
|
||||
|
||||
@@ -64,6 +64,29 @@ func TestKafkaSinkRejectsUnknownProtocol(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestKafkaSinkPublishesDurableRecordsInSingleWriterCall(t *testing.T) {
|
||||
writer := &recordingWriter{}
|
||||
sink := newKafkaSinkWithWriter(writer, KafkaConfig{UnifiedTopic: "custom.unified"})
|
||||
env := envelope.FrameEnvelope{Protocol: envelope.ProtocolJT808, Phone: "13307795425", MessageID: "0x0200"}
|
||||
|
||||
err := sink.PublishRecords(context.Background(), []durableRecord{
|
||||
{Kind: "raw", Envelope: env},
|
||||
{Kind: "unified", Envelope: env},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("PublishRecords() error = %v", err)
|
||||
}
|
||||
if writer.writeCalls != 1 {
|
||||
t.Fatalf("WriteMessages calls = %d, want 1", writer.writeCalls)
|
||||
}
|
||||
if len(writer.messages) != 2 {
|
||||
t.Fatalf("messages = %d, want 2", len(writer.messages))
|
||||
}
|
||||
if writer.messages[0].Topic != "vehicle.raw.jt808.v1" || writer.messages[1].Topic != "custom.unified" {
|
||||
t.Fatalf("topics = %q, %q", writer.messages[0].Topic, writer.messages[1].Topic)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNewKafkaSinkUsesProductionDeliveryGuarantees(t *testing.T) {
|
||||
sink, err := NewKafkaSink(KafkaConfig{Brokers: []string{"127.0.0.1:9092"}})
|
||||
if err != nil {
|
||||
@@ -90,10 +113,12 @@ func TestNewKafkaSinkUsesProductionDeliveryGuarantees(t *testing.T) {
|
||||
}
|
||||
|
||||
type recordingWriter struct {
|
||||
messages []kafka.Message
|
||||
messages []kafka.Message
|
||||
writeCalls int
|
||||
}
|
||||
|
||||
func (w *recordingWriter) WriteMessages(_ context.Context, messages ...kafka.Message) error {
|
||||
w.writeCalls++
|
||||
w.messages = append(w.messages, messages...)
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -40,6 +40,29 @@ func (s *RetryingSink) PublishUnified(ctx context.Context, env envelope.FrameEnv
|
||||
})
|
||||
}
|
||||
|
||||
func (s *RetryingSink) PublishRecords(ctx context.Context, records []durableRecord) error {
|
||||
if publisher, ok := s.delegate.(recordPublishingSink); ok {
|
||||
return s.withRetry(ctx, func(attemptCtx context.Context) error {
|
||||
return publisher.PublishRecords(attemptCtx, records)
|
||||
})
|
||||
}
|
||||
for _, record := range records {
|
||||
switch record.Kind {
|
||||
case "raw":
|
||||
if err := s.PublishRaw(ctx, record.Envelope); err != nil {
|
||||
return err
|
||||
}
|
||||
case "unified":
|
||||
if err := s.PublishUnified(ctx, record.Envelope); err != nil {
|
||||
return err
|
||||
}
|
||||
default:
|
||||
return errUnknownRecordKind(record.Kind)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *RetryingSink) Close() error {
|
||||
return s.delegate.Close()
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user