126 lines
2.5 KiB
Go
126 lines
2.5 KiB
Go
package eventbus
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"sync"
|
|
"time"
|
|
|
|
"lingniu-vehicle-ingest/go/vehicle-gateway/internal/envelope"
|
|
)
|
|
|
|
type AsyncConfig struct {
|
|
QueueSize int
|
|
Workers int
|
|
OperationTimeout time.Duration
|
|
OnError func(error)
|
|
}
|
|
|
|
type AsyncSink struct {
|
|
delegate Sink
|
|
jobs chan asyncJob
|
|
timeout time.Duration
|
|
onError func(error)
|
|
|
|
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")
|
|
|
|
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
|
|
}
|
|
s := &AsyncSink{
|
|
delegate: delegate,
|
|
jobs: make(chan asyncJob, cfg.QueueSize),
|
|
timeout: cfg.OperationTimeout,
|
|
onError: cfg.OnError,
|
|
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:
|
|
return ErrAsyncSinkClosed
|
|
default:
|
|
}
|
|
select {
|
|
case s.jobs <- job:
|
|
return nil
|
|
case <-s.closed:
|
|
return ErrAsyncSinkClosed
|
|
case <-ctx.Done():
|
|
return ctx.Err()
|
|
}
|
|
}
|
|
|
|
func (s *AsyncSink) worker() {
|
|
defer s.wg.Done()
|
|
for job := range s.jobs {
|
|
ctx, cancel := context.WithTimeout(context.Background(), s.timeout)
|
|
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()
|
|
if err != nil && s.onError != nil {
|
|
s.onError(err)
|
|
}
|
|
}
|
|
}
|