mirror of
https://github.com/Rain-kl/OpenFlare.git
synced 2026-10-06 23:56:37 +08:00
feat(plugins): implement runtime drivers for http, asynq worker, and cron
This commit is contained in:
@@ -0,0 +1,216 @@
|
|||||||
|
package driver_asynq_cron
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"sync"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/hibiken/asynq"
|
||||||
|
|
||||||
|
"github.com/Rain-kl/Wavelet/core"
|
||||||
|
)
|
||||||
|
|
||||||
|
// Option configures the Asynq cron scheduler driver plugin.
|
||||||
|
type Option func(*Plugin)
|
||||||
|
|
||||||
|
// WithRedisOpt sets the Redis connection options for Asynq scheduler.
|
||||||
|
func WithRedisOpt(opt asynq.RedisConnOpt) Option {
|
||||||
|
return func(p *Plugin) {
|
||||||
|
p.redisOpt = opt
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// WithLocation sets the timezone location for the cron scheduler.
|
||||||
|
func WithLocation(loc *time.Location) Option {
|
||||||
|
return func(p *Plugin) {
|
||||||
|
p.location = loc
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// WithSchedulerOpts sets custom Asynq SchedulerOpts.
|
||||||
|
func WithSchedulerOpts(opts *asynq.SchedulerOpts) Option {
|
||||||
|
return func(p *Plugin) {
|
||||||
|
p.schedulerOpts = opts
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// WithScheduler injects a pre-configured Asynq Scheduler instance.
|
||||||
|
func WithScheduler(sched *asynq.Scheduler) Option {
|
||||||
|
return func(p *Plugin) {
|
||||||
|
p.scheduler = sched
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Plugin implements core.Plugin and core.Driver for Asynq Cron Scheduler.
|
||||||
|
type Plugin struct {
|
||||||
|
mu sync.RWMutex
|
||||||
|
redisOpt asynq.RedisConnOpt
|
||||||
|
location *time.Location
|
||||||
|
schedulerOpts *asynq.SchedulerOpts
|
||||||
|
scheduler *asynq.Scheduler
|
||||||
|
running bool
|
||||||
|
coreCtx *core.Context
|
||||||
|
}
|
||||||
|
|
||||||
|
// New creates a new Asynq Cron Scheduler driver plugin.
|
||||||
|
func New(opts ...Option) *Plugin {
|
||||||
|
p := &Plugin{
|
||||||
|
redisOpt: asynq.RedisClientOpt{Addr: "127.0.0.1:6379"},
|
||||||
|
location: time.Local,
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, opt := range opts {
|
||||||
|
if opt != nil {
|
||||||
|
opt(p)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
return p
|
||||||
|
}
|
||||||
|
|
||||||
|
// Name returns the unique plugin identifier.
|
||||||
|
func (p *Plugin) Name() string {
|
||||||
|
return "driver_asynq_cron"
|
||||||
|
}
|
||||||
|
|
||||||
|
// Apply mounts the Asynq Cron Scheduler driver into the micro-kernel Context.
|
||||||
|
func (p *Plugin) Apply(ctx *core.Context) error {
|
||||||
|
p.mu.Lock()
|
||||||
|
p.coreCtx = ctx
|
||||||
|
p.mu.Unlock()
|
||||||
|
|
||||||
|
ctx.OnDispose(func() error {
|
||||||
|
return p.Stop(context.Background())
|
||||||
|
})
|
||||||
|
|
||||||
|
return ctx.RegisterDriver(p)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Type returns DriverTypeScheduler.
|
||||||
|
func (p *Plugin) Type() core.DriverType {
|
||||||
|
return core.DriverTypeScheduler
|
||||||
|
}
|
||||||
|
|
||||||
|
// Start registers scheduled cron entries from Context and boots the Asynq scheduler.
|
||||||
|
func (p *Plugin) Start(_ context.Context) error {
|
||||||
|
p.mu.Lock()
|
||||||
|
defer p.mu.Unlock()
|
||||||
|
|
||||||
|
if p.running {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
if p.scheduler == nil {
|
||||||
|
opts := p.schedulerOpts
|
||||||
|
if opts == nil {
|
||||||
|
opts = &asynq.SchedulerOpts{
|
||||||
|
Location: p.location,
|
||||||
|
}
|
||||||
|
} else if opts.Location == nil && p.location != nil {
|
||||||
|
opts.Location = p.location
|
||||||
|
}
|
||||||
|
|
||||||
|
p.scheduler = asynq.NewScheduler(p.redisOpt, opts)
|
||||||
|
}
|
||||||
|
|
||||||
|
if p.coreCtx != nil && p.coreCtx.Schedules() != nil {
|
||||||
|
for _, sd := range p.coreCtx.Schedules().Schedules() {
|
||||||
|
payloadBytes, err := encodePayload(sd.Payload)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("driver_asynq_cron: encode payload for task %q failed: %w", sd.TaskType, err)
|
||||||
|
}
|
||||||
|
|
||||||
|
task := asynq.NewTask(sd.TaskType, payloadBytes)
|
||||||
|
asynqOpts := buildAsynqOptions(sd.Options)
|
||||||
|
|
||||||
|
if _, err := p.scheduler.Register(sd.Spec, task, asynqOpts...); err != nil {
|
||||||
|
return fmt.Errorf("driver_asynq_cron: register schedule %q for task %q failed: %w", sd.Spec, sd.TaskType, err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if err := p.scheduler.Start(); err != nil {
|
||||||
|
return fmt.Errorf("driver_asynq_cron: start scheduler failed: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
p.running = true
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Stop gracefully shuts down the Asynq cron scheduler.
|
||||||
|
func (p *Plugin) Stop(_ context.Context) error {
|
||||||
|
p.mu.Lock()
|
||||||
|
defer p.mu.Unlock()
|
||||||
|
|
||||||
|
if !p.running {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
p.running = false
|
||||||
|
|
||||||
|
if p.scheduler != nil {
|
||||||
|
p.scheduler.Shutdown()
|
||||||
|
p.scheduler = nil
|
||||||
|
}
|
||||||
|
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// IsRunning returns whether the scheduler is running.
|
||||||
|
func (p *Plugin) IsRunning() bool {
|
||||||
|
p.mu.RLock()
|
||||||
|
defer p.mu.RUnlock()
|
||||||
|
return p.running
|
||||||
|
}
|
||||||
|
|
||||||
|
// Scheduler returns the underlying Asynq scheduler instance.
|
||||||
|
func (p *Plugin) Scheduler() *asynq.Scheduler {
|
||||||
|
p.mu.RLock()
|
||||||
|
defer p.mu.RUnlock()
|
||||||
|
return p.scheduler
|
||||||
|
}
|
||||||
|
|
||||||
|
func encodePayload(payload any) ([]byte, error) {
|
||||||
|
if payload == nil {
|
||||||
|
return nil, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
switch val := payload.(type) {
|
||||||
|
case []byte:
|
||||||
|
return val, nil
|
||||||
|
case string:
|
||||||
|
return []byte(val), nil
|
||||||
|
default:
|
||||||
|
return json.Marshal(val)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func buildAsynqOptions(opts map[string]any) []asynq.Option {
|
||||||
|
if len(opts) == 0 {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
var res []asynq.Option
|
||||||
|
|
||||||
|
if q, ok := opts["queue"].(string); ok && q != "" {
|
||||||
|
res = append(res, asynq.Queue(q))
|
||||||
|
}
|
||||||
|
|
||||||
|
if retry, ok := opts["retry"].(int); ok {
|
||||||
|
res = append(res, asynq.MaxRetry(retry))
|
||||||
|
} else if maxRetry, ok := opts["max_retry"].(int); ok {
|
||||||
|
res = append(res, asynq.MaxRetry(maxRetry))
|
||||||
|
}
|
||||||
|
|
||||||
|
if timeout, ok := opts["timeout"].(time.Duration); ok {
|
||||||
|
res = append(res, asynq.Timeout(timeout))
|
||||||
|
}
|
||||||
|
|
||||||
|
if retention, ok := opts["retention"].(time.Duration); ok {
|
||||||
|
res = append(res, asynq.Retention(retention))
|
||||||
|
}
|
||||||
|
|
||||||
|
return res
|
||||||
|
}
|
||||||
@@ -0,0 +1,236 @@
|
|||||||
|
package driver_asynq_worker
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"errors"
|
||||||
|
"fmt"
|
||||||
|
"sync"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/hibiken/asynq"
|
||||||
|
|
||||||
|
"github.com/Rain-kl/Wavelet/core"
|
||||||
|
)
|
||||||
|
|
||||||
|
const (
|
||||||
|
defaultConcurrency = 10
|
||||||
|
defaultShutdownTimeout = 10 * time.Second
|
||||||
|
)
|
||||||
|
|
||||||
|
// Option configures the Asynq worker driver plugin.
|
||||||
|
type Option func(*Plugin)
|
||||||
|
|
||||||
|
// WithRedisOpt sets the Redis connection options for Asynq.
|
||||||
|
func WithRedisOpt(opt asynq.RedisConnOpt) Option {
|
||||||
|
return func(p *Plugin) {
|
||||||
|
p.redisOpt = opt
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// WithConcurrency sets the worker concurrency limit.
|
||||||
|
func WithConcurrency(concurrency int) Option {
|
||||||
|
return func(p *Plugin) {
|
||||||
|
p.concurrency = concurrency
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// WithQueues sets the queue priorities mapping.
|
||||||
|
func WithQueues(queues map[string]int) Option {
|
||||||
|
return func(p *Plugin) {
|
||||||
|
p.queues = queues
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// WithStrictPriority sets whether to process queues strictly in priority order.
|
||||||
|
func WithStrictPriority(strict bool) Option {
|
||||||
|
return func(p *Plugin) {
|
||||||
|
p.strictPriority = strict
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// WithShutdownTimeout sets the timeout for graceful worker shutdown.
|
||||||
|
func WithShutdownTimeout(d time.Duration) Option {
|
||||||
|
return func(p *Plugin) {
|
||||||
|
p.shutdownTimeout = d
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// WithServer injects an existing Asynq server instance.
|
||||||
|
func WithServer(srv *asynq.Server) Option {
|
||||||
|
return func(p *Plugin) {
|
||||||
|
p.server = srv
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Plugin implements core.Plugin and core.Driver for Asynq background worker server.
|
||||||
|
type Plugin struct {
|
||||||
|
mu sync.RWMutex
|
||||||
|
redisOpt asynq.RedisConnOpt
|
||||||
|
concurrency int
|
||||||
|
queues map[string]int
|
||||||
|
strictPriority bool
|
||||||
|
shutdownTimeout time.Duration
|
||||||
|
server *asynq.Server
|
||||||
|
mux *asynq.ServeMux
|
||||||
|
running bool
|
||||||
|
coreCtx *core.Context
|
||||||
|
}
|
||||||
|
|
||||||
|
// New creates a new Asynq Worker driver plugin.
|
||||||
|
func New(opts ...Option) *Plugin {
|
||||||
|
p := &Plugin{
|
||||||
|
redisOpt: asynq.RedisClientOpt{Addr: "127.0.0.1:6379"},
|
||||||
|
concurrency: defaultConcurrency,
|
||||||
|
shutdownTimeout: defaultShutdownTimeout,
|
||||||
|
queues: map[string]int{"default": 1},
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, opt := range opts {
|
||||||
|
if opt != nil {
|
||||||
|
opt(p)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
return p
|
||||||
|
}
|
||||||
|
|
||||||
|
// Name returns the unique plugin identifier.
|
||||||
|
func (p *Plugin) Name() string {
|
||||||
|
return "driver_asynq_worker"
|
||||||
|
}
|
||||||
|
|
||||||
|
// Apply mounts the Asynq Worker driver into the micro-kernel Context.
|
||||||
|
func (p *Plugin) Apply(ctx *core.Context) error {
|
||||||
|
p.mu.Lock()
|
||||||
|
p.coreCtx = ctx
|
||||||
|
p.mu.Unlock()
|
||||||
|
|
||||||
|
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 boots the Asynq worker server and starts processing background tasks.
|
||||||
|
func (p *Plugin) Start(ctx context.Context) error {
|
||||||
|
p.mu.Lock()
|
||||||
|
defer p.mu.Unlock()
|
||||||
|
|
||||||
|
if p.running {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
mux := asynq.NewServeMux()
|
||||||
|
|
||||||
|
if p.coreCtx != nil && p.coreCtx.Tasks() != nil {
|
||||||
|
for _, td := range p.coreCtx.Tasks().Tasks() {
|
||||||
|
handler, err := toAsynqHandler(td.Handler)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("driver_asynq_worker: invalid handler for task pattern %q: %w", td.Pattern, err)
|
||||||
|
}
|
||||||
|
mux.Handle(td.Pattern, handler)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if p.server == nil {
|
||||||
|
p.server = asynq.NewServer(
|
||||||
|
p.redisOpt,
|
||||||
|
asynq.Config{
|
||||||
|
Concurrency: p.concurrency,
|
||||||
|
Queues: p.queues,
|
||||||
|
StrictPriority: p.strictPriority,
|
||||||
|
ShutdownTimeout: p.shutdownTimeout,
|
||||||
|
},
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
if err := p.server.Start(mux); err != nil {
|
||||||
|
return fmt.Errorf("driver_asynq_worker: start server failed: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
p.mux = mux
|
||||||
|
p.running = true
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Stop gracefully shuts down the Asynq worker server.
|
||||||
|
func (p *Plugin) Stop(_ context.Context) error {
|
||||||
|
p.mu.Lock()
|
||||||
|
defer p.mu.Unlock()
|
||||||
|
|
||||||
|
if !p.running {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
p.running = false
|
||||||
|
|
||||||
|
if p.server != nil {
|
||||||
|
p.server.Stop()
|
||||||
|
p.server.Shutdown()
|
||||||
|
p.server = nil
|
||||||
|
}
|
||||||
|
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// IsRunning returns whether the worker server is running.
|
||||||
|
func (p *Plugin) IsRunning() bool {
|
||||||
|
p.mu.RLock()
|
||||||
|
defer p.mu.RUnlock()
|
||||||
|
return p.running
|
||||||
|
}
|
||||||
|
|
||||||
|
// Server returns the underlying Asynq server instance.
|
||||||
|
func (p *Plugin) Server() *asynq.Server {
|
||||||
|
p.mu.RLock()
|
||||||
|
defer p.mu.RUnlock()
|
||||||
|
return p.server
|
||||||
|
}
|
||||||
|
|
||||||
|
// Mux returns the underlying Asynq serve mux.
|
||||||
|
func (p *Plugin) Mux() *asynq.ServeMux {
|
||||||
|
p.mu.RLock()
|
||||||
|
defer p.mu.RUnlock()
|
||||||
|
return p.mux
|
||||||
|
}
|
||||||
|
|
||||||
|
func toAsynqHandler(h any) (asynq.Handler, error) {
|
||||||
|
if h == nil {
|
||||||
|
return nil, errors.New("nil handler")
|
||||||
|
}
|
||||||
|
|
||||||
|
switch fn := h.(type) {
|
||||||
|
case asynq.Handler:
|
||||||
|
return fn, nil
|
||||||
|
case asynq.HandlerFunc:
|
||||||
|
return fn, nil
|
||||||
|
case func(context.Context, *asynq.Task) error:
|
||||||
|
return asynq.HandlerFunc(fn), nil
|
||||||
|
case func(context.Context, []byte) error:
|
||||||
|
return asynq.HandlerFunc(func(c context.Context, t *asynq.Task) error {
|
||||||
|
return fn(c, t.Payload())
|
||||||
|
}), nil
|
||||||
|
case func(context.Context) error:
|
||||||
|
return asynq.HandlerFunc(func(c context.Context, _ *asynq.Task) error {
|
||||||
|
return fn(c)
|
||||||
|
}), nil
|
||||||
|
case func([]byte) error:
|
||||||
|
return asynq.HandlerFunc(func(_ context.Context, t *asynq.Task) error {
|
||||||
|
return fn(t.Payload())
|
||||||
|
}), nil
|
||||||
|
case func() error:
|
||||||
|
return asynq.HandlerFunc(func(_ context.Context, _ *asynq.Task) error {
|
||||||
|
return fn()
|
||||||
|
}), nil
|
||||||
|
default:
|
||||||
|
return nil, fmt.Errorf("unsupported task handler type: %T", h)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,243 @@
|
|||||||
|
package driver_http
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"errors"
|
||||||
|
"fmt"
|
||||||
|
"net"
|
||||||
|
"net/http"
|
||||||
|
"sync"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/gin-gonic/gin"
|
||||||
|
|
||||||
|
"github.com/Rain-kl/Wavelet/core"
|
||||||
|
"github.com/Rain-kl/Wavelet/pkg/util"
|
||||||
|
)
|
||||||
|
|
||||||
|
const (
|
||||||
|
defaultAddr = ":8080"
|
||||||
|
defaultReadHeaderTimeout = 10 * time.Second
|
||||||
|
defaultShutdownTimeout = 5 * time.Second
|
||||||
|
)
|
||||||
|
|
||||||
|
// Option configures the HTTP driver plugin.
|
||||||
|
type Option func(*Plugin)
|
||||||
|
|
||||||
|
// WithAddr sets the TCP address for the HTTP server to listen on.
|
||||||
|
func WithAddr(addr string) Option {
|
||||||
|
return func(p *Plugin) {
|
||||||
|
p.addr = addr
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// WithEngine sets a pre-configured Gin engine for the HTTP server.
|
||||||
|
func WithEngine(engine *gin.Engine) Option {
|
||||||
|
return func(p *Plugin) {
|
||||||
|
p.engine = engine
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// WithReadHeaderTimeout sets the ReadHeaderTimeout for http.Server.
|
||||||
|
func WithReadHeaderTimeout(d time.Duration) Option {
|
||||||
|
return func(p *Plugin) {
|
||||||
|
p.readHeaderTimeout = d
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// WithShutdownTimeout sets the fallback timeout for graceful server shutdown.
|
||||||
|
func WithShutdownTimeout(d time.Duration) Option {
|
||||||
|
return func(p *Plugin) {
|
||||||
|
p.shutdownTimeout = d
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Plugin implements core.Plugin and core.Driver for Gin HTTP Web Server.
|
||||||
|
type Plugin struct {
|
||||||
|
mu sync.RWMutex
|
||||||
|
addr string
|
||||||
|
engine *gin.Engine
|
||||||
|
server *http.Server
|
||||||
|
listener net.Listener
|
||||||
|
running bool
|
||||||
|
readHeaderTimeout time.Duration
|
||||||
|
shutdownTimeout time.Duration
|
||||||
|
coreCtx *core.Context
|
||||||
|
}
|
||||||
|
|
||||||
|
// New creates a new Gin HTTP server driver plugin.
|
||||||
|
func New(opts ...Option) *Plugin {
|
||||||
|
p := &Plugin{
|
||||||
|
addr: defaultAddr,
|
||||||
|
readHeaderTimeout: defaultReadHeaderTimeout,
|
||||||
|
shutdownTimeout: defaultShutdownTimeout,
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, opt := range opts {
|
||||||
|
if opt != nil {
|
||||||
|
opt(p)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
return p
|
||||||
|
}
|
||||||
|
|
||||||
|
// Name returns the unique plugin identifier.
|
||||||
|
func (p *Plugin) Name() string {
|
||||||
|
return "driver_http"
|
||||||
|
}
|
||||||
|
|
||||||
|
// Apply mounts the HTTP driver plugin into the micro-kernel Context.
|
||||||
|
func (p *Plugin) Apply(ctx *core.Context) error {
|
||||||
|
p.mu.Lock()
|
||||||
|
p.coreCtx = ctx
|
||||||
|
p.mu.Unlock()
|
||||||
|
|
||||||
|
ctx.OnDispose(func() error {
|
||||||
|
shutdownCtx, cancel := context.WithTimeout(context.Background(), p.shutdownTimeout)
|
||||||
|
defer cancel()
|
||||||
|
return p.Stop(shutdownCtx)
|
||||||
|
})
|
||||||
|
|
||||||
|
return ctx.RegisterDriver(p)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Type returns DriverTypeHTTP.
|
||||||
|
func (p *Plugin) Type() core.DriverType {
|
||||||
|
return core.DriverTypeHTTP
|
||||||
|
}
|
||||||
|
|
||||||
|
// Start boots the Gin HTTP server, binds routes collected from ctx.Router(), and starts listening.
|
||||||
|
func (p *Plugin) Start(ctx context.Context) error {
|
||||||
|
p.mu.Lock()
|
||||||
|
defer p.mu.Unlock()
|
||||||
|
|
||||||
|
if p.running {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
if p.engine == nil {
|
||||||
|
p.engine = gin.New()
|
||||||
|
}
|
||||||
|
|
||||||
|
// Mount routes collected in Context RouterExtension
|
||||||
|
if p.coreCtx != nil && p.coreCtx.Router() != nil {
|
||||||
|
for _, rd := range p.coreCtx.Router().Routes() {
|
||||||
|
allHandlers := make([]gin.HandlerFunc, 0, len(rd.Middlewares)+len(rd.Handlers))
|
||||||
|
|
||||||
|
for _, m := range rd.Middlewares {
|
||||||
|
gh, err := toGinHandler(m)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("driver_http: invalid middleware for route %s %s: %w", rd.Method, rd.Path, err)
|
||||||
|
}
|
||||||
|
allHandlers = append(allHandlers, gh)
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, h := range rd.Handlers {
|
||||||
|
gh, err := toGinHandler(h)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("driver_http: invalid handler for route %s %s: %w", rd.Method, rd.Path, err)
|
||||||
|
}
|
||||||
|
allHandlers = append(allHandlers, gh)
|
||||||
|
}
|
||||||
|
|
||||||
|
p.engine.Handle(rd.Method, rd.Path, allHandlers...)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
p.server = &http.Server{
|
||||||
|
Addr: p.addr,
|
||||||
|
Handler: p.engine,
|
||||||
|
ReadHeaderTimeout: p.readHeaderTimeout,
|
||||||
|
}
|
||||||
|
|
||||||
|
lc := &net.ListenConfig{}
|
||||||
|
listener, err := lc.Listen(ctx, "tcp", p.addr)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("driver_http: listen on %s failed: %w", p.addr, err)
|
||||||
|
}
|
||||||
|
|
||||||
|
p.listener = listener
|
||||||
|
p.addr = listener.Addr().String()
|
||||||
|
p.running = true
|
||||||
|
|
||||||
|
srv := p.server
|
||||||
|
util.Go(func() {
|
||||||
|
if serveErr := srv.Serve(listener); serveErr != nil && !errors.Is(serveErr, http.ErrServerClosed) {
|
||||||
|
_ = serveErr
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Stop gracefully stops the HTTP server.
|
||||||
|
func (p *Plugin) Stop(ctx context.Context) error {
|
||||||
|
p.mu.Lock()
|
||||||
|
defer p.mu.Unlock()
|
||||||
|
|
||||||
|
if !p.running {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
p.running = false
|
||||||
|
|
||||||
|
var err error
|
||||||
|
if p.server != nil {
|
||||||
|
if ctx == nil {
|
||||||
|
var cancel context.CancelFunc
|
||||||
|
ctx, cancel = context.WithTimeout(context.Background(), p.shutdownTimeout)
|
||||||
|
defer cancel()
|
||||||
|
}
|
||||||
|
err = p.server.Shutdown(ctx)
|
||||||
|
}
|
||||||
|
|
||||||
|
if p.listener != nil {
|
||||||
|
_ = p.listener.Close()
|
||||||
|
p.listener = nil
|
||||||
|
}
|
||||||
|
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
// Addr returns the current listening address (or configured address if not yet started).
|
||||||
|
func (p *Plugin) Addr() string {
|
||||||
|
p.mu.RLock()
|
||||||
|
defer p.mu.RUnlock()
|
||||||
|
return p.addr
|
||||||
|
}
|
||||||
|
|
||||||
|
// Engine returns the underlying Gin engine.
|
||||||
|
func (p *Plugin) Engine() *gin.Engine {
|
||||||
|
p.mu.RLock()
|
||||||
|
defer p.mu.RUnlock()
|
||||||
|
return p.engine
|
||||||
|
}
|
||||||
|
|
||||||
|
// IsRunning returns whether the HTTP server is currently running.
|
||||||
|
func (p *Plugin) IsRunning() bool {
|
||||||
|
p.mu.RLock()
|
||||||
|
defer p.mu.RUnlock()
|
||||||
|
return p.running
|
||||||
|
}
|
||||||
|
|
||||||
|
func toGinHandler(h any) (gin.HandlerFunc, error) {
|
||||||
|
if h == nil {
|
||||||
|
return nil, errors.New("nil handler")
|
||||||
|
}
|
||||||
|
|
||||||
|
switch fn := h.(type) {
|
||||||
|
case gin.HandlerFunc:
|
||||||
|
return fn, nil
|
||||||
|
case func(*gin.Context):
|
||||||
|
return gin.HandlerFunc(fn), nil
|
||||||
|
case http.HandlerFunc:
|
||||||
|
return gin.WrapF(fn), nil
|
||||||
|
case func(http.ResponseWriter, *http.Request):
|
||||||
|
return gin.WrapF(fn), nil
|
||||||
|
case http.Handler:
|
||||||
|
return gin.WrapH(fn), nil
|
||||||
|
default:
|
||||||
|
return nil, fmt.Errorf("unsupported handler type: %T", h)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,288 @@
|
|||||||
|
package drivers_test
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bytes"
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"io"
|
||||||
|
"net/http"
|
||||||
|
"sync/atomic"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/alicebob/miniredis/v2"
|
||||||
|
"github.com/gin-gonic/gin"
|
||||||
|
"github.com/hibiken/asynq"
|
||||||
|
"github.com/stretchr/testify/assert"
|
||||||
|
"github.com/stretchr/testify/require"
|
||||||
|
|
||||||
|
"github.com/Rain-kl/Wavelet/core"
|
||||||
|
"github.com/Rain-kl/Wavelet/core/extpoints"
|
||||||
|
"github.com/Rain-kl/Wavelet/plugins/drivers/driver_asynq_cron"
|
||||||
|
"github.com/Rain-kl/Wavelet/plugins/drivers/driver_asynq_worker"
|
||||||
|
"github.com/Rain-kl/Wavelet/plugins/drivers/driver_http"
|
||||||
|
)
|
||||||
|
|
||||||
|
func init() {
|
||||||
|
gin.SetMode(gin.TestMode)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestHTTPDriverLifecycle(t *testing.T) {
|
||||||
|
ctx := core.NewContext(context.Background())
|
||||||
|
|
||||||
|
var globalMiddlewareCalled atomic.Bool
|
||||||
|
var groupMiddlewareCalled atomic.Bool
|
||||||
|
|
||||||
|
// Register global middleware
|
||||||
|
ctx.Router().Use(func(c *gin.Context) {
|
||||||
|
globalMiddlewareCalled.Store(true)
|
||||||
|
c.Next()
|
||||||
|
})
|
||||||
|
|
||||||
|
// Register standard gin handler
|
||||||
|
ctx.Router().GET("/ping", func(c *gin.Context) {
|
||||||
|
c.JSON(http.StatusOK, gin.H{"message": "pong"})
|
||||||
|
})
|
||||||
|
|
||||||
|
// Register route group with middleware
|
||||||
|
v1 := ctx.Router().Group("/api/v1", func(c *gin.Context) {
|
||||||
|
groupMiddlewareCalled.Store(true)
|
||||||
|
c.Next()
|
||||||
|
})
|
||||||
|
|
||||||
|
v1.POST("/echo", func(c *gin.Context) {
|
||||||
|
var req map[string]any
|
||||||
|
if err := c.ShouldBindJSON(&req); err != nil {
|
||||||
|
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||||
|
return
|
||||||
|
}
|
||||||
|
c.JSON(http.StatusOK, req)
|
||||||
|
})
|
||||||
|
|
||||||
|
// Register standard http.HandlerFunc
|
||||||
|
v1.GET("/legacy", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||||
|
w.WriteHeader(http.StatusOK)
|
||||||
|
_, _ = w.Write([]byte("legacy response"))
|
||||||
|
}))
|
||||||
|
|
||||||
|
// Create and apply HTTP driver with dynamic port
|
||||||
|
httpPlugin := driver_http.New(driver_http.WithAddr("127.0.0.1:0"))
|
||||||
|
require.Equal(t, "driver_http", httpPlugin.Name())
|
||||||
|
|
||||||
|
err := httpPlugin.Apply(ctx)
|
||||||
|
require.NoError(t, err)
|
||||||
|
|
||||||
|
// Verify driver registered in context
|
||||||
|
d, ok := ctx.Driver(core.DriverTypeHTTP)
|
||||||
|
require.True(t, ok)
|
||||||
|
require.Equal(t, core.DriverTypeHTTP, d.Type())
|
||||||
|
|
||||||
|
// Start HTTP Server
|
||||||
|
err = d.Start(context.Background())
|
||||||
|
require.NoError(t, err)
|
||||||
|
assert.True(t, httpPlugin.IsRunning())
|
||||||
|
|
||||||
|
addr := httpPlugin.Addr()
|
||||||
|
require.NotEmpty(t, addr)
|
||||||
|
|
||||||
|
// Verify GET /ping
|
||||||
|
resp, err := http.Get(fmt.Sprintf("http://%s/ping", addr))
|
||||||
|
require.NoError(t, err)
|
||||||
|
defer resp.Body.Close()
|
||||||
|
|
||||||
|
assert.Equal(t, http.StatusOK, resp.StatusCode)
|
||||||
|
body, err := io.ReadAll(resp.Body)
|
||||||
|
require.NoError(t, err)
|
||||||
|
assert.Contains(t, string(body), "pong")
|
||||||
|
assert.True(t, globalMiddlewareCalled.Load())
|
||||||
|
|
||||||
|
// Verify POST /api/v1/echo
|
||||||
|
echoPayload := []byte(`{"title":"test-echo","value":42}`)
|
||||||
|
respEcho, err := http.Post(fmt.Sprintf("http://%s/api/v1/echo", addr), "application/json", bytes.NewReader(echoPayload))
|
||||||
|
require.NoError(t, err)
|
||||||
|
defer respEcho.Body.Close()
|
||||||
|
|
||||||
|
assert.Equal(t, http.StatusOK, respEcho.StatusCode)
|
||||||
|
var echoResult map[string]any
|
||||||
|
err = json.NewDecoder(respEcho.Body).Decode(&echoResult)
|
||||||
|
require.NoError(t, err)
|
||||||
|
assert.Equal(t, "test-echo", echoResult["title"])
|
||||||
|
assert.Equal(t, float64(42), echoResult["value"])
|
||||||
|
assert.True(t, groupMiddlewareCalled.Load())
|
||||||
|
|
||||||
|
// Verify GET /api/v1/legacy
|
||||||
|
respLegacy, err := http.Get(fmt.Sprintf("http://%s/api/v1/legacy", addr))
|
||||||
|
require.NoError(t, err)
|
||||||
|
defer respLegacy.Body.Close()
|
||||||
|
|
||||||
|
assert.Equal(t, http.StatusOK, respLegacy.StatusCode)
|
||||||
|
bodyLegacy, err := io.ReadAll(respLegacy.Body)
|
||||||
|
require.NoError(t, err)
|
||||||
|
assert.Equal(t, "legacy response", string(bodyLegacy))
|
||||||
|
|
||||||
|
// Stop HTTP Server
|
||||||
|
stopCtx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
|
||||||
|
defer cancel()
|
||||||
|
|
||||||
|
err = d.Stop(stopCtx)
|
||||||
|
require.NoError(t, err)
|
||||||
|
assert.False(t, httpPlugin.IsRunning())
|
||||||
|
|
||||||
|
// Idempotent Stop
|
||||||
|
err = d.Stop(stopCtx)
|
||||||
|
require.NoError(t, err)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAsynqWorkerDriverLifecycle(t *testing.T) {
|
||||||
|
mr, err := miniredis.Run()
|
||||||
|
require.NoError(t, err)
|
||||||
|
defer mr.Close()
|
||||||
|
|
||||||
|
ctx := core.NewContext(context.Background())
|
||||||
|
|
||||||
|
var emailTaskProcessed atomic.Bool
|
||||||
|
var simpleTaskProcessed atomic.Bool
|
||||||
|
|
||||||
|
// Register tasks with various handler signatures
|
||||||
|
ctx.Tasks().Register("email:send", func(c context.Context, payload []byte) error {
|
||||||
|
var data map[string]string
|
||||||
|
if err := json.Unmarshal(payload, &data); err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
if data["to"] == "user@example.com" {
|
||||||
|
emailTaskProcessed.Store(true)
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}, extpoints.WithTaskConcurrency(2))
|
||||||
|
|
||||||
|
ctx.Tasks().Register("maintenance:cleanup", func() error {
|
||||||
|
simpleTaskProcessed.Store(true)
|
||||||
|
return nil
|
||||||
|
})
|
||||||
|
|
||||||
|
workerPlugin := driver_asynq_worker.New(
|
||||||
|
driver_asynq_worker.WithRedisOpt(asynq.RedisClientOpt{Addr: mr.Addr()}),
|
||||||
|
driver_asynq_worker.WithConcurrency(2),
|
||||||
|
driver_asynq_worker.WithShutdownTimeout(1*time.Second),
|
||||||
|
)
|
||||||
|
require.Equal(t, "driver_asynq_worker", workerPlugin.Name())
|
||||||
|
|
||||||
|
err = workerPlugin.Apply(ctx)
|
||||||
|
require.NoError(t, err)
|
||||||
|
|
||||||
|
// Verify driver registered
|
||||||
|
d, ok := ctx.Driver(core.DriverTypeWorker)
|
||||||
|
require.True(t, ok)
|
||||||
|
require.Equal(t, core.DriverTypeWorker, d.Type())
|
||||||
|
|
||||||
|
// Start Worker Driver
|
||||||
|
err = d.Start(context.Background())
|
||||||
|
require.NoError(t, err)
|
||||||
|
assert.True(t, workerPlugin.IsRunning())
|
||||||
|
|
||||||
|
// Enqueue tasks using asynq.Client
|
||||||
|
client := asynq.NewClient(asynq.RedisClientOpt{Addr: mr.Addr()})
|
||||||
|
defer client.Close()
|
||||||
|
|
||||||
|
emailPayload, _ := json.Marshal(map[string]string{"to": "user@example.com"})
|
||||||
|
_, err = client.Enqueue(asynq.NewTask("email:send", emailPayload))
|
||||||
|
require.NoError(t, err)
|
||||||
|
|
||||||
|
_, err = client.Enqueue(asynq.NewTask("maintenance:cleanup", nil))
|
||||||
|
require.NoError(t, err)
|
||||||
|
|
||||||
|
// Wait for worker to consume tasks
|
||||||
|
require.Eventually(t, func() bool {
|
||||||
|
return emailTaskProcessed.Load() && simpleTaskProcessed.Load()
|
||||||
|
}, 3*time.Second, 50*time.Millisecond)
|
||||||
|
|
||||||
|
// Stop Worker Driver
|
||||||
|
stopCtx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
|
||||||
|
defer cancel()
|
||||||
|
|
||||||
|
err = d.Stop(stopCtx)
|
||||||
|
require.NoError(t, err)
|
||||||
|
assert.False(t, workerPlugin.IsRunning())
|
||||||
|
|
||||||
|
// Idempotent Stop
|
||||||
|
err = d.Stop(stopCtx)
|
||||||
|
require.NoError(t, err)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestAsynqCronDriverLifecycle(t *testing.T) {
|
||||||
|
mr, err := miniredis.Run()
|
||||||
|
require.NoError(t, err)
|
||||||
|
defer mr.Close()
|
||||||
|
|
||||||
|
ctx := core.NewContext(context.Background())
|
||||||
|
|
||||||
|
// Register cron schedules
|
||||||
|
ctx.Schedules().RegisterCron("@every 1s", "sync:stats", map[string]string{"type": "daily"},
|
||||||
|
extpoints.WithScheduleOption("queue", "default"),
|
||||||
|
extpoints.WithScheduleOption("retry", 3),
|
||||||
|
)
|
||||||
|
ctx.Schedules().RegisterCron("0 0 * * *", "report:generate", "raw_payload")
|
||||||
|
|
||||||
|
cronPlugin := driver_asynq_cron.New(
|
||||||
|
driver_asynq_cron.WithRedisOpt(asynq.RedisClientOpt{Addr: mr.Addr()}),
|
||||||
|
driver_asynq_cron.WithLocation(time.UTC),
|
||||||
|
)
|
||||||
|
require.Equal(t, "driver_asynq_cron", cronPlugin.Name())
|
||||||
|
|
||||||
|
err = cronPlugin.Apply(ctx)
|
||||||
|
require.NoError(t, err)
|
||||||
|
|
||||||
|
// Verify driver registered
|
||||||
|
d, ok := ctx.Driver(core.DriverTypeScheduler)
|
||||||
|
require.True(t, ok)
|
||||||
|
require.Equal(t, core.DriverTypeScheduler, d.Type())
|
||||||
|
|
||||||
|
// Start Cron Scheduler Driver
|
||||||
|
err = d.Start(context.Background())
|
||||||
|
require.NoError(t, err)
|
||||||
|
assert.True(t, cronPlugin.IsRunning())
|
||||||
|
|
||||||
|
// Stop Cron Scheduler Driver
|
||||||
|
stopCtx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
|
||||||
|
defer cancel()
|
||||||
|
|
||||||
|
err = d.Stop(stopCtx)
|
||||||
|
require.NoError(t, err)
|
||||||
|
assert.False(t, cronPlugin.IsRunning())
|
||||||
|
|
||||||
|
// Idempotent Stop
|
||||||
|
err = d.Stop(stopCtx)
|
||||||
|
require.NoError(t, err)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestMultipleDriversInContext(t *testing.T) {
|
||||||
|
mr, err := miniredis.Run()
|
||||||
|
require.NoError(t, err)
|
||||||
|
defer mr.Close()
|
||||||
|
|
||||||
|
ctx := core.NewContext(context.Background())
|
||||||
|
|
||||||
|
httpPlugin := driver_http.New(driver_http.WithAddr("127.0.0.1:0"))
|
||||||
|
workerPlugin := driver_asynq_worker.New(driver_asynq_worker.WithRedisOpt(asynq.RedisClientOpt{Addr: mr.Addr()}))
|
||||||
|
cronPlugin := driver_asynq_cron.New(driver_asynq_cron.WithRedisOpt(asynq.RedisClientOpt{Addr: mr.Addr()}))
|
||||||
|
|
||||||
|
require.NoError(t, httpPlugin.Apply(ctx))
|
||||||
|
require.NoError(t, workerPlugin.Apply(ctx))
|
||||||
|
require.NoError(t, cronPlugin.Apply(ctx))
|
||||||
|
|
||||||
|
drivers := ctx.Drivers()
|
||||||
|
assert.Len(t, drivers, 3)
|
||||||
|
|
||||||
|
_, ok := ctx.Driver(core.DriverTypeHTTP)
|
||||||
|
assert.True(t, ok)
|
||||||
|
|
||||||
|
_, ok = ctx.Driver(core.DriverTypeWorker)
|
||||||
|
assert.True(t, ok)
|
||||||
|
|
||||||
|
_, ok = ctx.Driver(core.DriverTypeScheduler)
|
||||||
|
assert.True(t, ok)
|
||||||
|
|
||||||
|
_, ok = ctx.Driver("non_existent")
|
||||||
|
assert.False(t, ok)
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user