mirror of
https://github.com/Rain-kl/OpenFlare.git
synced 2026-09-29 05:56:38 +08:00
feat(drivers): implement in-process async worker driver plugin
This commit is contained in:
@@ -0,0 +1,188 @@
|
||||
// Copyright 2026 Arctel.net
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
package driver_inproc_worker
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"Wavelet/core/extpoints"
|
||||
"Wavelet/pkg/idgen"
|
||||
"Wavelet/pkg/util"
|
||||
)
|
||||
|
||||
// TaskMessage represents an in-process task item in the queue.
|
||||
type TaskMessage struct {
|
||||
ID string
|
||||
TaskType string
|
||||
Payload []byte
|
||||
Source string
|
||||
CreatedAt time.Time
|
||||
RetryLeft int
|
||||
}
|
||||
|
||||
// InprocQueue manages in-memory task queuing and worker pool execution.
|
||||
type InprocQueue struct {
|
||||
mu sync.RWMutex
|
||||
concurrency int
|
||||
queue chan TaskMessage
|
||||
taskReg extpoints.TaskExtension
|
||||
running atomic.Bool
|
||||
stopCh chan struct{}
|
||||
wg sync.WaitGroup
|
||||
}
|
||||
|
||||
// NewInprocQueue creates a new InprocQueue with a given concurrency and queue capacity.
|
||||
func NewInprocQueue(concurrency int, queueCap int, taskReg extpoints.TaskExtension) *InprocQueue {
|
||||
if concurrency <= 0 {
|
||||
concurrency = 10
|
||||
}
|
||||
if queueCap <= 0 {
|
||||
queueCap = 1000
|
||||
}
|
||||
|
||||
return &InprocQueue{
|
||||
concurrency: concurrency,
|
||||
queue: make(chan TaskMessage, queueCap),
|
||||
taskReg: taskReg,
|
||||
stopCh: make(chan struct{}),
|
||||
}
|
||||
}
|
||||
|
||||
// Enqueue puts a new task into the in-process queue.
|
||||
func (q *InprocQueue) Enqueue(taskType string, payload []byte, source string) (string, error) {
|
||||
if !q.running.Load() {
|
||||
return "", errors.New("driver_inproc_worker: queue is not running")
|
||||
}
|
||||
|
||||
taskID := fmt.Sprintf("inproc_%d", idgen.NextUint64ID())
|
||||
msg := TaskMessage{
|
||||
ID: taskID,
|
||||
TaskType: taskType,
|
||||
Payload: payload,
|
||||
Source: source,
|
||||
CreatedAt: time.Now(),
|
||||
}
|
||||
|
||||
if q.taskReg != nil {
|
||||
if td, ok := q.taskReg.Get(taskType); ok {
|
||||
msg.RetryLeft = td.Retry
|
||||
}
|
||||
}
|
||||
|
||||
select {
|
||||
case q.queue <- msg:
|
||||
return taskID, nil
|
||||
default:
|
||||
return "", errors.New("driver_inproc_worker: queue is full")
|
||||
}
|
||||
}
|
||||
|
||||
// Start begins processing tasks with the worker pool.
|
||||
func (q *InprocQueue) Start() {
|
||||
if !q.running.CompareAndSwap(false, true) {
|
||||
return
|
||||
}
|
||||
|
||||
for i := 0; i < q.concurrency; i++ {
|
||||
q.wg.Add(1)
|
||||
util.Go(func() {
|
||||
defer q.wg.Done()
|
||||
q.workerLoop()
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// Stop gracefully waits for in-flight tasks and shuts down workers.
|
||||
func (q *InprocQueue) Stop(ctx context.Context) error {
|
||||
if !q.running.CompareAndSwap(true, false) {
|
||||
return nil
|
||||
}
|
||||
|
||||
close(q.stopCh)
|
||||
|
||||
done := make(chan struct{})
|
||||
util.Go(func() {
|
||||
q.wg.Wait()
|
||||
close(done)
|
||||
})
|
||||
|
||||
select {
|
||||
case <-done:
|
||||
return nil
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
}
|
||||
}
|
||||
|
||||
func (q *InprocQueue) workerLoop() {
|
||||
for {
|
||||
select {
|
||||
case <-q.stopCh:
|
||||
return
|
||||
case msg, ok := <-q.queue:
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
q.executeTask(msg)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (q *InprocQueue) executeTask(msg TaskMessage) {
|
||||
if q.taskReg == nil {
|
||||
return
|
||||
}
|
||||
|
||||
td, ok := q.taskReg.Get(msg.TaskType)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
|
||||
timeout := td.Timeout
|
||||
if timeout <= 0 {
|
||||
timeout = 5 * time.Minute
|
||||
}
|
||||
|
||||
taskCtx, cancel := context.WithTimeout(context.Background(), timeout)
|
||||
defer cancel()
|
||||
|
||||
err := invokeHandler(taskCtx, td.Handler, msg.Payload)
|
||||
if err != nil && msg.RetryLeft > 0 {
|
||||
msg.RetryLeft--
|
||||
// Retry with backoff
|
||||
util.Go(func() {
|
||||
time.Sleep(500 * time.Millisecond)
|
||||
if q.running.Load() {
|
||||
select {
|
||||
case q.queue <- msg:
|
||||
default:
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func invokeHandler(ctx context.Context, handler any, payload []byte) error {
|
||||
if handler == nil {
|
||||
return errors.New("nil task handler")
|
||||
}
|
||||
|
||||
switch fn := handler.(type) {
|
||||
case func(context.Context, []byte) error:
|
||||
return fn(ctx, payload)
|
||||
case func(context.Context) error:
|
||||
return fn(ctx)
|
||||
case func([]byte) error:
|
||||
return fn(payload)
|
||||
case func() error:
|
||||
return fn()
|
||||
default:
|
||||
return fmt.Errorf("unsupported handler type: %T", handler)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,145 @@
|
||||
// Copyright 2026 Arctel.net
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
// Package driver_inproc_worker provides the zero-dependency in-process async worker driver plugin for Cordis.
|
||||
package driver_inproc_worker
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"Wavelet/core"
|
||||
)
|
||||
|
||||
const (
|
||||
defaultConcurrency = 10
|
||||
defaultQueueCapacity = 2000
|
||||
defaultShutdownTimeout = 5 * time.Second
|
||||
)
|
||||
|
||||
var (
|
||||
globalMu sync.RWMutex
|
||||
globalQueue *InprocQueue
|
||||
)
|
||||
|
||||
// DispatchTask enqueues a background task to the active in-process worker queue.
|
||||
func DispatchTask(_ context.Context, taskType string, payload []byte, source string) (string, error) {
|
||||
globalMu.RLock()
|
||||
q := globalQueue
|
||||
globalMu.RUnlock()
|
||||
|
||||
if q == nil {
|
||||
return "", nil
|
||||
}
|
||||
return q.Enqueue(taskType, payload, source)
|
||||
}
|
||||
|
||||
// Option configures the in-process worker driver plugin.
|
||||
type Option func(*Plugin)
|
||||
|
||||
// WithConcurrency sets the number of concurrent worker goroutines.
|
||||
func WithConcurrency(concurrency int) Option {
|
||||
return func(p *Plugin) {
|
||||
p.concurrency = concurrency
|
||||
}
|
||||
}
|
||||
|
||||
// WithQueueCapacity sets the internal queue buffer capacity.
|
||||
func WithQueueCapacity(capacity int) Option {
|
||||
return func(p *Plugin) {
|
||||
p.queueCapacity = capacity
|
||||
}
|
||||
}
|
||||
|
||||
// WithShutdownTimeout sets the maximum duration to wait for in-flight tasks during graceful shutdown.
|
||||
func WithShutdownTimeout(d time.Duration) Option {
|
||||
return func(p *Plugin) {
|
||||
p.shutdownTimeout = d
|
||||
}
|
||||
}
|
||||
|
||||
// Plugin implements core.Plugin and core.Driver for in-process background worker execution.
|
||||
type Plugin struct {
|
||||
concurrency int
|
||||
queueCapacity int
|
||||
shutdownTimeout time.Duration
|
||||
coreCtx *core.Context
|
||||
queue *InprocQueue
|
||||
}
|
||||
|
||||
// New creates a new in-process worker driver plugin.
|
||||
func New(opts ...Option) *Plugin {
|
||||
p := &Plugin{
|
||||
concurrency: defaultConcurrency,
|
||||
queueCapacity: defaultQueueCapacity,
|
||||
shutdownTimeout: defaultShutdownTimeout,
|
||||
}
|
||||
for _, opt := range opts {
|
||||
if opt != nil {
|
||||
opt(p)
|
||||
}
|
||||
}
|
||||
return p
|
||||
}
|
||||
|
||||
// Name returns the unique identifier of the plugin.
|
||||
func (p *Plugin) Name() string {
|
||||
return "driver_inproc_worker"
|
||||
}
|
||||
|
||||
// Manifest returns the plugin metadata.
|
||||
func (p *Plugin) Manifest() core.Manifest {
|
||||
return core.Manifest{
|
||||
Name: "driver_inproc_worker",
|
||||
Version: "1.0.0",
|
||||
Description: "Zero-dependency in-process async worker driver plugin",
|
||||
Author: "Wavelet Team",
|
||||
}
|
||||
}
|
||||
|
||||
// Apply registers the worker driver into the Context.
|
||||
func (p *Plugin) Apply(ctx *core.Context) error {
|
||||
p.coreCtx = ctx
|
||||
|
||||
ctx.OnDispose(func() error {
|
||||
shutdownCtx, cancel := context.WithTimeout(context.Background(), p.shutdownTimeout)
|
||||
defer cancel()
|
||||
return p.Stop(shutdownCtx)
|
||||
})
|
||||
|
||||
return ctx.RegisterDriver(p)
|
||||
}
|
||||
|
||||
// Type returns DriverTypeWorker.
|
||||
func (p *Plugin) Type() core.DriverType {
|
||||
return core.DriverTypeWorker
|
||||
}
|
||||
|
||||
// Start initiates task consumption.
|
||||
func (p *Plugin) Start(_ context.Context) error {
|
||||
if p.queue == nil {
|
||||
p.queue = NewInprocQueue(p.concurrency, p.queueCapacity, p.coreCtx.Tasks())
|
||||
}
|
||||
|
||||
globalMu.Lock()
|
||||
globalQueue = p.queue
|
||||
globalMu.Unlock()
|
||||
|
||||
p.queue.Start()
|
||||
return nil
|
||||
}
|
||||
|
||||
// Stop gracefully terminates worker processing.
|
||||
func (p *Plugin) Stop(ctx context.Context) error {
|
||||
if p.queue != nil {
|
||||
err := p.queue.Stop(ctx)
|
||||
globalMu.Lock()
|
||||
if globalQueue == p.queue {
|
||||
globalQueue = nil
|
||||
}
|
||||
globalMu.Unlock()
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,52 @@
|
||||
// Copyright 2026 Arctel.net
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
package driver_inproc_worker_test
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"Wavelet/core"
|
||||
"Wavelet/core/extpoints"
|
||||
"Wavelet/plugins/drivers/driver_inproc_worker"
|
||||
)
|
||||
|
||||
func TestInprocWorkerPlugin(t *testing.T) {
|
||||
ctx := core.NewContext(context.Background())
|
||||
p := driver_inproc_worker.New(
|
||||
driver_inproc_worker.WithConcurrency(2),
|
||||
driver_inproc_worker.WithShutdownTimeout(time.Second),
|
||||
)
|
||||
|
||||
assert.Equal(t, "driver_inproc_worker", p.Name())
|
||||
assert.Equal(t, core.DriverTypeWorker, p.Type())
|
||||
require.NoError(t, p.Apply(ctx))
|
||||
|
||||
var executedCount atomic.Int32
|
||||
ctx.Tasks().Register("test:task", func(ctx context.Context, payload []byte) error {
|
||||
executedCount.Add(1)
|
||||
return nil
|
||||
}, extpoints.WithTaskTimeout(2*time.Second))
|
||||
|
||||
// Start worker driver
|
||||
require.NoError(t, p.Start(context.Background()))
|
||||
|
||||
// Enqueue tasks
|
||||
taskID, err := driver_inproc_worker.DispatchTask(context.Background(), "test:task", []byte("hello"), "test")
|
||||
require.NoError(t, err)
|
||||
assert.NotEmpty(t, taskID)
|
||||
|
||||
// Wait for execution
|
||||
require.Eventually(t, func() bool {
|
||||
return executedCount.Load() == 1
|
||||
}, 2*time.Second, 20*time.Millisecond)
|
||||
|
||||
// Stop driver
|
||||
require.NoError(t, p.Stop(context.Background()))
|
||||
}
|
||||
Reference in New Issue
Block a user