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
	EnqueueTimeout   time.Duration
	OperationTimeout time.Duration
	OnError          func(error)
	Metrics          *metrics.Registry
	Name             string
}

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

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

type asyncJob struct {
	kind       string
	env        envelope.FrameEnvelope
	enqueuedAt time.Time
}

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

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.EnqueueTimeout == 0 {
		cfg.EnqueueTimeout = time.Second
	}
	if cfg.EnqueueTimeout < 0 {
		cfg.EnqueueTimeout = 0
	}
	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),
		enqueueTimeout: cfg.EnqueueTimeout,
		timeout:        cfg.OperationTimeout,
		onError:        cfg.OnError,
		metrics:        cfg.Metrics,
		name:           cfg.Name,
		queueWait:      metrics.NewRecentLatencyByKey(512),
		closed:         make(chan struct{}),
		done:           make(chan struct{}),
	}
	s.wg.Add(cfg.Workers)
	for i := 0; i < cfg.Workers; i++ {
		go s.worker()
	}
	s.recordQueueCapacity()
	s.recordWorkers("default", cfg.Workers)
	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)
	})
	<-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:
	}
	var timeoutC <-chan time.Time
	var timer *time.Timer
	if s.enqueueTimeout > 0 {
		timer = time.NewTimer(s.enqueueTimeout)
		timeoutC = timer.C
		defer timer.Stop()
	}
	job.enqueuedAt = time.Now()
	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()
	case <-timeoutC:
		s.recordEnqueue(job.kind, "timeout")
		s.recordQueueDepth()
		return ErrAsyncSinkEnqueueTimeout
	}
}

func (s *AsyncSink) worker() {
	defer s.wg.Done()
	for {
		select {
		case job := <-s.jobs:
			s.publishJob(job)
		case <-s.closed:
			for {
				select {
				case job := <-s.jobs:
					s.publishJob(job)
				default:
					return
				}
			}
		}
	}
}

func (s *AsyncSink) publishJob(job asyncJob) {
	s.recordQueueDepth()
	s.recordQueueWait("default", job)
	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)))
}

func (s *AsyncSink) recordQueueCapacity() {
	if s.metrics == nil {
		return
	}
	s.metrics.SetGauge("vehicle_async_sink_queue_capacity", metrics.Labels{
		"sink": s.name,
	}, float64(cap(s.jobs)))
}

func (s *AsyncSink) recordQueueWait(queueName string, job asyncJob) {
	if s.metrics == nil || job.enqueuedAt.IsZero() {
		return
	}
	elapsedMS := float64(time.Since(job.enqueuedAt)) / float64(time.Millisecond)
	labels := metrics.Labels{
		"sink":  s.name,
		"queue": queueName,
		"kind":  job.kind,
	}
	s.metrics.ObserveHistogram("vehicle_async_sink_queue_wait_duration_ms_histogram", labels, asyncSinkPublishDurationBucketsMS, elapsedMS)
	p99, samples := s.queueWait.Observe(queueName+"\x00"+job.kind, elapsedMS)
	s.metrics.SetGauge("vehicle_async_sink_queue_wait_recent_p99_ms", labels, p99)
	s.metrics.SetGauge("vehicle_async_sink_queue_wait_recent_samples", labels, float64(samples))
}

func (s *AsyncSink) recordWorkers(queueName string, workers int) {
	if s.metrics == nil {
		return
	}
	s.metrics.SetGauge("vehicle_async_sink_workers", metrics.Labels{
		"sink":  s.name,
		"queue": queueName,
	}, float64(workers))
}