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

package eventbus

import (
	"context"
	"sync"
	"time"

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

type PartitionedAsyncConfig struct {
	RawQueueSize          int
	DerivedQueueSize      int
	RawWorkers            int
	DerivedWorkers        int
	EnqueueTimeout        time.Duration
	RawEnqueueTimeout     time.Duration
	DerivedEnqueueTimeout time.Duration
	OperationTimeout      time.Duration
	OnError               func(error)
	Metrics               *metrics.Registry
	Name                  string
}

type PartitionedAsyncSink struct {
	delegate              Sink
	rawJobs               chan asyncJob
	derivedJobs           chan asyncJob
	rawEnqueueTimeout     time.Duration
	derivedEnqueueTimeout 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
}

func NewPartitionedAsyncSink(delegate Sink, cfg PartitionedAsyncConfig) *PartitionedAsyncSink {
	if delegate == nil {
		panic("partitioned async delegate sink must not be nil")
	}
	if cfg.RawQueueSize <= 0 {
		cfg.RawQueueSize = 100_000
	}
	if cfg.DerivedQueueSize <= 0 {
		cfg.DerivedQueueSize = 50_000
	}
	if cfg.RawWorkers <= 0 {
		cfg.RawWorkers = 4
	}
	if cfg.DerivedWorkers <= 0 {
		cfg.DerivedWorkers = 2
	}
	enqueueTimeout := normalizePartitionedEnqueueTimeout(cfg.EnqueueTimeout, time.Second)
	rawEnqueueTimeout := normalizePartitionedEnqueueTimeout(cfg.RawEnqueueTimeout, enqueueTimeout)
	derivedEnqueueTimeout := normalizePartitionedEnqueueTimeout(cfg.DerivedEnqueueTimeout, enqueueTimeout)
	if cfg.OperationTimeout <= 0 {
		cfg.OperationTimeout = 30 * time.Second
	}
	if cfg.Name == "" {
		cfg.Name = "partitioned-async"
	}
	s := &PartitionedAsyncSink{
		delegate:              delegate,
		rawJobs:               make(chan asyncJob, cfg.RawQueueSize),
		derivedJobs:           make(chan asyncJob, cfg.DerivedQueueSize),
		rawEnqueueTimeout:     rawEnqueueTimeout,
		derivedEnqueueTimeout: derivedEnqueueTimeout,
		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.startWorkers("raw", s.rawJobs, cfg.RawWorkers)
	s.startWorkers("derived", s.derivedJobs, cfg.DerivedWorkers)
	s.recordWorkers("raw", cfg.RawWorkers)
	s.recordWorkers("derived", cfg.DerivedWorkers)
	s.recordQueueCapacity("raw", cap(s.rawJobs))
	s.recordQueueCapacity("derived", cap(s.derivedJobs))
	go func() {
		s.wg.Wait()
		close(s.done)
	}()
	return s
}

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

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

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

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

func (s *PartitionedAsyncSink) startWorkers(queueName string, jobs <-chan asyncJob, workers int) {
	s.wg.Add(workers)
	for i := 0; i < workers; i++ {
		go s.worker(queueName, jobs)
	}
}

func (s *PartitionedAsyncSink) enqueue(ctx context.Context, queueName string, jobs chan<- asyncJob, enqueueTimeout time.Duration, 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 enqueueTimeout > 0 {
		timer = time.NewTimer(enqueueTimeout)
		timeoutC = timer.C
		defer timer.Stop()
	}
	job.enqueuedAt = time.Now()
	select {
	case jobs <- job:
		s.recordEnqueue(job.kind, "queued")
		s.recordQueueDepth(queueName)
		return nil
	case <-s.closed:
		s.recordEnqueue(job.kind, "closed")
		return ErrAsyncSinkClosed
	case <-ctx.Done():
		s.recordEnqueue(job.kind, "timeout")
		s.recordQueueDepth(queueName)
		return ctx.Err()
	case <-timeoutC:
		s.recordEnqueue(job.kind, "timeout")
		s.recordQueueDepth(queueName)
		return ErrAsyncSinkEnqueueTimeout
	}
}

func normalizePartitionedEnqueueTimeout(value time.Duration, fallback time.Duration) time.Duration {
	if value == 0 {
		value = fallback
	}
	if value < 0 {
		return 0
	}
	return value
}

func (s *PartitionedAsyncSink) worker(queueName string, jobs <-chan asyncJob) {
	defer s.wg.Done()
	for {
		select {
		case job := <-jobs:
			s.publishJob(queueName, job)
		case <-s.closed:
			for {
				select {
				case job := <-jobs:
					s.publishJob(queueName, job)
				default:
					return
				}
			}
		}
	}
}

func (s *PartitionedAsyncSink) publishJob(queueName string, job asyncJob) {
	s.recordQueueDepth(queueName)
	s.recordQueueWait(queueName, 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(queueName)
}

func (s *PartitionedAsyncSink) 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 *PartitionedAsyncSink) 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 *PartitionedAsyncSink) recordQueueDepth(queueName string) {
	if s.metrics == nil {
		return
	}
	s.metrics.SetGauge("vehicle_async_sink_queue_depth", metrics.Labels{
		"sink":  s.name,
		"queue": queueName,
	}, float64(s.queueDepth(queueName)))
}

func (s *PartitionedAsyncSink) recordQueueCapacity(queueName string, value int) {
	if s.metrics == nil {
		return
	}
	s.metrics.SetGauge("vehicle_async_sink_queue_capacity", metrics.Labels{
		"sink":  s.name,
		"queue": queueName,
	}, float64(value))
}

func (s *PartitionedAsyncSink) 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 *PartitionedAsyncSink) 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))
}

func (s *PartitionedAsyncSink) queueDepth(queueName string) int {
	switch queueName {
	case "raw":
		return len(s.rawJobs)
	case "derived":
		return len(s.derivedJobs)
	default:
		return 0
	}
}