lingniu-vehicle-ingest/go/vehicle-gateway/internal/eventbus/async_sink.go

package eventbus

import (
	"context"
	"errors"
	"sync"
	"time"

	"lingniu-vehicle-ingest/go/vehicle-gateway/internal/envelope"
	"lingniu-vehicle-ingest/go/vehicle-gateway/internal/metrics"
)

type AsyncConfig struct {
	QueueSize        int
	Workers          int
	OperationTimeout time.Duration
	OnError          func(error)
	Metrics          *metrics.Registry
	Name             string
}

type AsyncSink struct {
	delegate Sink
	jobs     chan asyncJob
	timeout  time.Duration
	onError  func(error)
	metrics  *metrics.Registry
	name     string

	closeOnce sync.Once
	closed    chan struct{}
	done      chan struct{}
	wg        sync.WaitGroup
}

type asyncJob struct {
	kind string
	env  envelope.FrameEnvelope
}

var ErrAsyncSinkClosed = errors.New("async sink is closed")

var asyncSinkPublishDurationBucketsMS = []float64{1, 5, 10, 25, 50, 100, 250, 500, 1000, 5000}

func NewAsyncSink(delegate Sink, cfg AsyncConfig) *AsyncSink {
	if delegate == nil {
		panic("async delegate sink must not be nil")
	}
	if cfg.QueueSize <= 0 {
		cfg.QueueSize = 10_000
	}
	if cfg.Workers <= 0 {
		cfg.Workers = 1
	}
	if cfg.OperationTimeout <= 0 {
		cfg.OperationTimeout = 30 * time.Second
	}
	if cfg.Name == "" {
		cfg.Name = "async"
	}
	s := &AsyncSink{
		delegate: delegate,
		jobs:     make(chan asyncJob, cfg.QueueSize),
		timeout:  cfg.OperationTimeout,
		onError:  cfg.OnError,
		metrics:  cfg.Metrics,
		name:     cfg.Name,
		closed:   make(chan struct{}),
		done:     make(chan struct{}),
	}
	s.wg.Add(cfg.Workers)
	for i := 0; i < cfg.Workers; i++ {
		go s.worker()
	}
	go func() {
		s.wg.Wait()
		close(s.done)
	}()
	return s
}

func (s *AsyncSink) PublishRaw(ctx context.Context, env envelope.FrameEnvelope) error {
	return s.enqueue(ctx, asyncJob{kind: "raw", env: env})
}

func (s *AsyncSink) PublishUnified(ctx context.Context, env envelope.FrameEnvelope) error {
	return s.enqueue(ctx, asyncJob{kind: "unified", env: env})
}

func (s *AsyncSink) PublishFields(ctx context.Context, env envelope.FrameEnvelope) error {
	return s.enqueue(ctx, asyncJob{kind: "fields", env: env})
}

func (s *AsyncSink) Close() error {
	s.closeOnce.Do(func() {
		close(s.closed)
		close(s.jobs)
	})
	<-s.done
	return s.delegate.Close()
}

func (s *AsyncSink) enqueue(ctx context.Context, job asyncJob) error {
	select {
	case <-s.closed:
		s.recordEnqueue(job.kind, "closed")
		return ErrAsyncSinkClosed
	default:
	}
	select {
	case s.jobs <- job:
		s.recordEnqueue(job.kind, "queued")
		s.recordQueueDepth()
		return nil
	case <-s.closed:
		s.recordEnqueue(job.kind, "closed")
		return ErrAsyncSinkClosed
	case <-ctx.Done():
		s.recordEnqueue(job.kind, "timeout")
		s.recordQueueDepth()
		return ctx.Err()
	}
}

func (s *AsyncSink) worker() {
	defer s.wg.Done()
	for job := range s.jobs {
		s.recordQueueDepth()
		ctx, cancel := context.WithTimeout(context.Background(), s.timeout)
		started := time.Now()
		var err error
		switch job.kind {
		case "raw":
			err = s.delegate.PublishRaw(ctx, job.env)
		case "unified":
			err = s.delegate.PublishUnified(ctx, job.env)
		case "fields":
			err = s.delegate.PublishFields(ctx, job.env)
		}
		cancel()
		status := "ok"
		if err != nil {
			status = "error"
		}
		s.recordPublish(job.kind, status, time.Since(started))
		if err != nil && s.onError != nil {
			s.onError(err)
		}
		s.recordQueueDepth()
	}
}

func (s *AsyncSink) recordEnqueue(kind string, status string) {
	if s.metrics == nil {
		return
	}
	s.metrics.IncCounter("vehicle_async_sink_enqueue_total", metrics.Labels{
		"sink":   s.name,
		"kind":   kind,
		"status": status,
	})
}

func (s *AsyncSink) recordPublish(kind string, status string, elapsed time.Duration) {
	if s.metrics == nil {
		return
	}
	labels := metrics.Labels{
		"sink":   s.name,
		"kind":   kind,
		"status": status,
	}
	s.metrics.IncCounter("vehicle_async_sink_publish_total", labels)
	elapsedMS := float64(elapsed.Milliseconds())
	s.metrics.SetGauge("vehicle_async_sink_publish_duration_ms", labels, elapsedMS)
	s.metrics.ObserveHistogram("vehicle_async_sink_publish_duration_ms_histogram", labels, asyncSinkPublishDurationBucketsMS, elapsedMS)
}

func (s *AsyncSink) recordQueueDepth() {
	if s.metrics == nil {
		return
	}
	s.metrics.SetGauge("vehicle_async_sink_queue_depth", metrics.Labels{
		"sink": s.name,
	}, float64(len(s.jobs)))
}