From fe66221a9c9b1731f92c8e662d92d4b76158cdc6 Mon Sep 17 00:00:00 2001 From: k2safe Date: Wed, 23 Sep 2026 13:38:25 +0800 Subject: [PATCH] feat: add cache, outbound HTTP, request protection, tasks and module scaffolding --- MODULES.md | 168 +++++++++++ README.md | 7 + cache/cache.go | 135 +++++++++ cache/cache_test.go | 118 ++++++++ cache/memory.go | 135 +++++++++ cache/redis/redis.go | 68 +++++ cache/redis/redis_test.go | 125 +++++++++ cmd/go-omni/README.md | 7 + cmd/go-omni/main.go | 34 +++ cmd/go-omni/main_test.go | 24 ++ cmd/go-omni/scaffold/addModule.go | 109 ++++++++ cmd/go-omni/scaffold/addModule_test.go | 137 +++++++++ .../scaffold/moduleTemplates/README.md.tmpl | 15 + .../moduleTemplates/handler-handler.go.tmpl | 52 ++++ .../moduleTemplates/model-model.go.tmpl | 15 + .../scaffold/moduleTemplates/module.go.tmpl | 19 ++ .../repository-repository.go.tmpl | 60 ++++ .../moduleTemplates/service-service.go.tmpl | 27 ++ cmd/go-omni/scaffold/scaffold.go | 2 +- examples/modules/README.md | 33 +++ examples/modules/main.go | 133 +++++++++ examples/modules/main_test.go | 126 +++++++++ go.mod | 4 + go.sum | 12 + httpc/client.go | 247 +++++++++++++++++ httpc/client_test.go | 131 +++++++++ httpx/context.go | 31 +++ middleware/idempotency.go | 235 ++++++++++++++++ middleware/protection_test.go | 227 +++++++++++++++ middleware/rateLimit.go | 109 ++++++++ middleware/timeout.go | 35 +++ scripts/verifyModules.py | 157 +++++++++++ task/task.go | 262 ++++++++++++++++++ task/task_test.go | 102 +++++++ 34 files changed, 3100 insertions(+), 1 deletion(-) create mode 100644 MODULES.md create mode 100644 cache/cache.go create mode 100644 cache/cache_test.go create mode 100644 cache/memory.go create mode 100644 cache/redis/redis.go create mode 100644 cache/redis/redis_test.go create mode 100644 cmd/go-omni/scaffold/addModule.go create mode 100644 cmd/go-omni/scaffold/addModule_test.go create mode 100644 cmd/go-omni/scaffold/moduleTemplates/README.md.tmpl create mode 100644 cmd/go-omni/scaffold/moduleTemplates/handler-handler.go.tmpl create mode 100644 cmd/go-omni/scaffold/moduleTemplates/model-model.go.tmpl create mode 100644 cmd/go-omni/scaffold/moduleTemplates/module.go.tmpl create mode 100644 cmd/go-omni/scaffold/moduleTemplates/repository-repository.go.tmpl create mode 100644 cmd/go-omni/scaffold/moduleTemplates/service-service.go.tmpl create mode 100644 examples/modules/README.md create mode 100644 examples/modules/main.go create mode 100644 examples/modules/main_test.go create mode 100644 httpc/client.go create mode 100644 httpc/client_test.go create mode 100644 middleware/idempotency.go create mode 100644 middleware/protection_test.go create mode 100644 middleware/rateLimit.go create mode 100644 middleware/timeout.go create mode 100644 scripts/verifyModules.py create mode 100644 task/task.go create mode 100644 task/task_test.go diff --git a/MODULES.md b/MODULES.md new file mode 100644 index 0000000..9163c6b --- /dev/null +++ b/MODULES.md @@ -0,0 +1,168 @@ +# 业务基础模块 + +这批能力位于 main，尚未包含在 `v0.6.0-preview` 标签中。统一示例见 +[examples/modules](examples/modules/README.md)。业务入口继续使用 `omni.New`，中间件依赖 +`httpx`，不依赖具体服务端网络实现。各模块按需创建、注入 Service，再注册到应用关闭流程。 + +## 缓存：cache 与 cache/redis + +```go +store, err := cache.NewMemory(cache.MemoryConfig{MaxEntries: 4096, MaxBytes: 32 << 20}) +if err != nil { return err } +application.OnShutdown(store.Close) +loader, err := cache.NewLoader(store, cache.LoaderConfig{MaxLoads: 128, Timeout: 5 * time.Second}) +if err != nil { return err } +value, err := loader.GetOrLoad(ctx, "tenant:42:product:7", time.Minute, + func(ctx context.Context) ([]byte, error) { return repository.LoadProduct(ctx, "42", "7") }) +``` + +- `Store` 统一 `Get/Set/Delete`，返回字节切片；业务可用 JSON 编解码。 +- 内存缓存使用 LRU，同时限制条目数和 key/value 总字节数；字节上限不包含 Go 对象开销。 + 读写复制切片，过期时懒清理，没有每个条目一个 goroutine。 +- TTL 为零表示不过期，负值拒绝；缺失返回 `cache.ErrMiss`。删除不存在的 key 成功。 +- Loader 合并**同进程、同 key**的并发回源，最多同时运行 `MaxLoads` 个不同 key 的加载。 + 单个调用者取消只停止等待；共享加载有独立超时，加载函数必须响应 context。 + 超限返回 `ErrBusy`，缓存故障直接返回错误，不自动当成 miss 冲击数据库。 +- 同 key 的调用必须具有相同数据含义、权限范围和 TTL；加载采用首个调用者的 context 值。 + 多租户 key 必须含租户标识。进程间合并、分布式锁、缓存更新与数据库事务一致性由业务另外设计。 + +Redis 适配器借用调用方创建的 go-redis 客户端和连接池： + +```go +client := goredis.NewClient(&goredis.Options{ + Addr: "127.0.0.1:6379", PoolSize: 16, MaxRetries: -1, + DialTimeout: time.Second, ReadTimeout: time.Second, + WriteTimeout: time.Second, PoolTimeout: time.Second, + ContextTimeoutEnabled: true, +}) +// 启动时先使用有截止时间的 context 执行 client.Ping(ctx)。 +store, err := rediscache.New(client, "billing:") // cache/redis 的导入别名 +if err != nil { _ = client.Close(); return err } +application.OnShutdown(client.Close) +// store 可直接传入 cache.NewLoader。 +``` + +连接失败与 Redis 错误保持原样；只有 Redis nil 映射为 `ErrMiss`。Redis TTL 精度为毫秒， +正的亚毫秒 TTL 提升到 1ms。适配器不负责关闭借用的客户端，Redis 服务端容量策略由部署配置。 +连接池和超时配置参考 [go-redis 官方生产使用说明](https://redis.io/docs/latest/develop/clients/go/produsage/)。 + +## 出站 HTTP：httpc + +```go +client, err := httpc.New(httpc.Config{ + Timeout: 3*time.Second, MaxRetries: 2, MaxConnsPerHost: 32, + MaxResponseBytes: 1 << 20, +}) +if err != nil { return err } +application.OnShutdown(client.Close) +err = client.JSON(ctx, "GET", serviceURL, nil, &result) +``` + +客户端应在服务启动时创建并复用。默认使用独立的 `net/http.Transport` 连接池，每个目标主机 +有连接上限；该上限不是整个进程的全局连接数上限。也可以注入自有 `http.RoundTripper`， +其连接限制和生命周期由调用方负责。这样出站实现同样可以替换。 + +总超时覆盖建连、重试、等待和读取响应体。`Do` 返回后调用方必须读取/关闭 `response.Body`； +`JSON` 自动关闭，并限制解码前读取的字节数。`Do` 允许流式读取，本身不限制响应体大小。 +`Close` 拒绝新请求、关闭自有池的空闲连接；在途请求应先排空。 + +`MaxRetries` 默认 0，最多 5。封装只对可重放的 GET/HEAD/OPTIONS 重试网络错误及 +429/502/503/504；POST/PUT/PATCH/DELETE 不做应用层自动重试。标准库 Transport 和借用的 +RoundTripper 自身的重放规则仍生效，不能据此承诺“只执行一次”。退避带抖动，尊重 +`Retry-After`；服务器要求等待超过 `MaxBackoff` 时直接返回该响应，不提前重试。 +重定向交给业务处理，3xx 不会自动跟随。 + +默认从 context 传递 `X-Request-ID`，日志只包含方法、主机、状态、尝试次数和收到响应头的耗时， +不记录完整 URL、查询参数、凭证和响应体。非 2xx 的 JSON 调用返回只带状态码的 `StatusError`。 + +## 请求保护：middleware + +```go +limiter, err := middleware.NewRateLimiter(middleware.RateLimitConfig{Rate: 20, Burst: 40}) +if err != nil { return err } +idempotency, err := middleware.NewIdempotency(middleware.IdempotencyConfig{ + Scope: func(c *httpx.Context) (string, error) { + principal, _ := auth.FromContext(c) + return principal.UserID, nil // 多租户应用应组合租户和用户 ID + }, +}) +if err != nil { return err } +api := application.Group("/api", middleware.Timeout(3*time.Second), + auth.RequireAuth(authenticator), limiter.Middleware()) +api.POST("/orders", createOrder, idempotency.Middleware()) +``` + +限流为进程内令牌桶，`Rate` 为每秒额度，`Burst` 为突发容量。默认按已解析的可信客户端 IP， +可用 `Key` 按已认证身份分桶。额度不足返回 **429 + Retry-After**；key 数超过 `MaxKeys` +返回 **503**。满容量时不会踢掉活跃桶让攻击者绕过限流；空闲桶经过完整补充周期后才允许回收。 + +Timeout 设置协作式 deadline，不额外启动 Handler goroutine。业务及下游调用必须响应 context； +超时后尚未提交响应时返回 **504**，已经提交的响应保持原状态。它不能强制中断任意 Go 代码， +也不能替代网络层读写超时。不要在业务代码中异步持有 `httpx.Context`。 + +幂等保护适用于小型、非流式写接口： + +- 必须先认证/鉴权，再执行幂等中间件；`Scope` 必须取自服务端验证过的身份。 + GET/HEAD/OPTIONS 直接通过，写请求要求 `Idempotency-Key`，长度 1..128 个非空格 ASCII 字符。 +- 同 Scope + key + 方法/URL/Content-Type/body 重放首次响应；内容改变返回 **409**。 + 正在执行或结果不可重放的相同请求也返回 **409**，不并发执行第二次。 +- 默认最多 32 个处理中请求、256 条记录、16 MiB 保存的响应头/body，单次请求体 1 MiB、 + 响应 body 64 KiB。请求超限 **413**，容量满 **503**，响应超限 **500**。 +- panic、5xx、超时、响应超限、Set-Cookie、Cache-Control:no-store 或保存容量不足时保留失败记录， + 阻止在 TTL 内自动重执行业务。业务可能已经提交，调用方应查询实际结果后处理。 +- 完成记录默认保留 10 分钟；在途记录不会因 TTL 到期而允许重复执行。 + 重放保留原响应 body、原 `requestId` 与原 `X-Request-ID`，增加 `Idempotency-Replayed:true`； + 本次访问日志仍使用本次请求 ID。 +- 记录仅在当前进程有效，重启、TTL 到期和多副本切换后均无法保证去重。支付、库存等关键写入 + 应在数据库事务内配合唯一业务键。该实现不提供跨进程 exactly-once，也不支持 SSE/流式/升级连接。 + +## 后台执行：task + +```go +jobs, err := task.New(task.Config{Workers: 4, QueueSize: 128, TaskTimeout: 30*time.Second}) +if err != nil { return err } +application.OnShutdown(jobs.Close) // 在注册数据库/缓存关闭之后注册，先排空任务 +future, err := jobs.Submit(serviceContext, "refreshCatalog", refreshCatalog) +if err != nil { return err } // 满队列立即返回 ErrQueueFull +err = future.Wait(ctx) // 可选；不等待时，任务失败仍通过 OnError/logx 报告 +err = jobs.Every("refreshCatalog", time.Minute, refreshCatalog) +``` + +固定 worker 和有界队列控制资源。任务保留提交 context；需要独立于 HTTP 请求执行时使用服务 +context，或明确使用 `context.WithoutCancel` 保留日志关联值。任务仍受 Runner 的超时/退出控制。 +任务函数与 `OnError` 回调应及时返回，尤其要响应取消；panic 被转为失败，worker 继续工作。 + +`Every` 使用固定间隔，第一次在一个间隔后执行；同一计划上轮尚未完成、排队或队列满时跳过。 +默认最多 128 个计划，名称不可重复。`Stats` 提供 active/queued/completed/failed/skipped。 +它是进程内任务执行器，没有任务持久化、失败自动重试、cron 表达式、跨副本选主或崩溃恢复。 + +`Shutdown(ctx)` 停止接收和调度，排空已接受任务；截止时间到后取消任务并及时返回。 +忽略取消的任务不能被 Go 强杀，可能继续运行。`Close` 默认给排空 10 秒。 + +## 增加业务模块：go-omni add + +```sh +go build -o /tmp/go-omni ./cmd/go-omni +/tmp/go-omni add orderItems --dir /path/to/billing +``` + +生成 `internal/orderItems/{module.go,model,repository,service,handler,README.md}`，包含 name/createdAt +示例模型、内存 CRUD、输入校验、错误映射及可注入 Repository。模块名使用 lowerCamelCase。 + +在现有启动代码中导入模块后调用 `orderItems.Register(application, repository, authMiddleware)`。 +脚手架不自动修改现有路由、main、go.mod 或数据库，已有目录不会覆盖。 +内存 Repository 仅为可运行示例：重启丢数据，列表未分页；正式业务应实现自己的持久化 Repository。 + +## 验证入口 + +```sh +GO_OMNI_SCAFFOLD_INTEGRATION=1 go test -race ./... +go vet ./... +go build ./... +python3 scripts/verifyModules.py --goproxy https://goproxy.cn,direct +``` + +Linux runner 顺序运行两个容器：全量 race/vet/build/独立生成项目，随后真实 Redis 适配验证。 +每次最多一个容器，限制 2 CPU、2 GiB（不额外使用 swap）、256 pids；Redis 测试容器禁外网， +只用内部 loopback，不发布端口。Redis 集成二进制为静态非 race，适配器单测在全量 race 中覆盖。 +结束清理本轮容器，并记录提交、源文件摘要、镜像 ID、日志摘要和原有容器是否仍运行。 diff --git a/README.md b/README.md index b8de915..474d6b1 100644 --- a/README.md +++ b/README.md @@ -43,6 +43,9 @@ go-omni/ ├── validate/ # 参数校验 ├── auth/ # 认证和权限接口 ├── database/ # database/sql 生命周期和事务 +├── cache/ # 有界内存缓存、Redis 适配与合并回源 +├── httpc/ # 出站 HTTP 连接池、超时和重试 +├── task/ # 有界后台任务与间隔调度 ├── transport/nethttp/ # 标准库网络适配器 ├── transport/native/ # 自研 HTTP/1.1 服务端预览 ├── transport/http2/ # opt-in HTTP/2 实验适配器 @@ -50,6 +53,10 @@ go-omni/ └── cmd/ # go-omni 项目脚手架 ``` +main 分支新增缓存、出站 HTTP、限流/幂等/超时中间件、后台任务和 `go-omni add` 模块脚手架。 +使用方式、默认容量和进程内能力的适用范围见 [业务基础模块](MODULES.md)， +可以直接运行 [组合示例](examples/modules/README.md)。这些新增能力尚未包含在上面的历史版本标签中。 + ## 快速开始 ```go diff --git a/cache/cache.go b/cache/cache.go new file mode 100644 index 0000000..e170100 --- /dev/null +++ b/cache/cache.go @@ -0,0 +1,135 @@ +// Package cache provides replaceable byte-oriented caches and bounded local +// request coalescing. Storage errors are never silently treated as misses. +package cache + +import ( + "context" + "errors" + "fmt" + "sync" + "time" +) + +var ( + ErrMiss = errors.New("cache: key not found") + ErrClosed = errors.New("cache: closed") + ErrTooLarge = errors.New("cache: entry exceeds capacity") + ErrBusy = errors.New("cache: too many concurrent loads") +) + +// Store implementations copy values and accept ttl=0 for no expiration. +// Negative TTLs are invalid. Delete is successful for missing keys. +type Store interface { + Get(context.Context, string) ([]byte, error) + Set(context.Context, string, []byte, time.Duration) error + Delete(context.Context, ...string) error +} + +func Check(ctx context.Context, key string, ttl time.Duration) error { + if ctx == nil { + return errors.New("cache: nil context") + } + if err := ctx.Err(); err != nil { + return err + } + if key == "" || len(key) > 1024 { + return errors.New("cache: key must contain 1..1024 bytes") + } + if ttl < 0 { + return errors.New("cache: negative TTL") + } + return nil +} + +type LoaderConfig struct { + MaxLoads int `json:"maxLoads"` + Timeout time.Duration `json:"timeout"` +} + +type flight struct { + done chan struct{} + value []byte + err error +} + +// Loader coalesces concurrent misses for the same key within this process. +// Caller cancellation stops waiting without canceling other callers' load. +// Load functions must observe their context, which has a bounded deadline. +type Loader struct { + store Store + config LoaderConfig + mu sync.Mutex + flights map[string]*flight +} + +func NewLoader(store Store, cfg LoaderConfig) (*Loader, error) { + if store == nil || cfg.MaxLoads < 0 || cfg.Timeout < 0 { + return nil, errors.New("cache: invalid loader configuration") + } + if cfg.MaxLoads == 0 { + cfg.MaxLoads = 128 + } + if cfg.Timeout == 0 { + cfg.Timeout = 5 * time.Second + } + return &Loader{store: store, config: cfg, flights: make(map[string]*flight)}, nil +} + +func (l *Loader) GetOrLoad(ctx context.Context, key string, ttl time.Duration, load func(context.Context) ([]byte, error)) ([]byte, error) { + if err := Check(ctx, key, ttl); err != nil { + return nil, err + } + if load == nil { + return nil, errors.New("cache: nil loader") + } + value, err := l.store.Get(ctx, key) + if !errors.Is(err, ErrMiss) { + return value, err + } + l.mu.Lock() + f := l.flights[key] + if f == nil { + if len(l.flights) >= l.config.MaxLoads { + l.mu.Unlock() + return nil, ErrBusy + } + f = &flight{done: make(chan struct{})} + l.flights[key] = f + go l.run(context.WithoutCancel(ctx), key, ttl, load, f) + } + l.mu.Unlock() + select { + case <-ctx.Done(): + return nil, ctx.Err() + case <-f.done: + return append([]byte(nil), f.value...), f.err + } +} + +func (l *Loader) run(parent context.Context, key string, ttl time.Duration, load func(context.Context) ([]byte, error), f *flight) { + ctx, cancel := context.WithTimeout(parent, l.config.Timeout) + defer cancel() + defer func() { + if value := recover(); value != nil { + f.err = fmt.Errorf("cache: loader panic: %v", value) + } + l.mu.Lock() + delete(l.flights, key) + close(f.done) + l.mu.Unlock() + }() + // Another request may have populated the key between the first Get and + // flight acquisition. Recheck inside the coalesced load. + f.value, f.err = l.store.Get(ctx, key) + if !errors.Is(f.err, ErrMiss) { + return + } + f.value, f.err = load(ctx) + if f.err == nil { + f.err = ctx.Err() + } + if f.err == nil { + f.err = l.store.Set(ctx, key, f.value, ttl) + } + f.value = append([]byte(nil), f.value...) +} diff --git a/cache/cache_test.go b/cache/cache_test.go new file mode 100644 index 0000000..8b13662 --- /dev/null +++ b/cache/cache_test.go @@ -0,0 +1,118 @@ +package cache + +import ( + "context" + "errors" + "sync/atomic" + "testing" + "time" +) + +func TestMemoryOwnershipCapacityExpiryAndClose(t *testing.T) { + m, err := NewMemory(MemoryConfig{MaxEntries: 2, MaxBytes: 10}) + if err != nil { + t.Fatal(err) + } + now := time.Now() + m.now = func() time.Time { return now } + ctx := context.Background() + value := []byte("aa") + if err := m.Set(ctx, "a", value, time.Second); err != nil { + t.Fatal(err) + } + value[0] = 'x' + got, _ := m.Get(ctx, "a") + if string(got) != "aa" { + t.Fatal("retained caller memory") + } + got[0] = 'z' + got, _ = m.Get(ctx, "a") + if string(got) != "aa" { + t.Fatal("returned mutable cache memory") + } + m.Set(ctx, "b", []byte("bb"), 0) + m.Get(ctx, "a") + m.Set(ctx, "c", []byte("cc"), 0) + if _, err := m.Get(ctx, "b"); !errors.Is(err, ErrMiss) { + t.Fatal("LRU eviction failed") + } + now = now.Add(time.Second) + if _, err := m.Get(ctx, "a"); !errors.Is(err, ErrMiss) { + t.Fatal("expiration boundary ignored") + } + if err := m.Set(ctx, "large", []byte("123456"), 0); !errors.Is(err, ErrTooLarge) { + t.Fatal(err) + } + if _, err := m.Get(ctx, "c"); err != nil { + t.Fatal("oversized write evicted valid data") + } + canceled, cancel := context.WithCancel(ctx) + cancel() + if err := m.Delete(canceled, "c"); !errors.Is(err, context.Canceled) { + t.Fatal(err) + } + m.Close() + m.Close() + if _, err := m.Get(ctx, "c"); !errors.Is(err, ErrClosed) { + t.Fatal(err) + } +} + +func TestLoaderCoalescesAndSeparatesCallerCancellation(t *testing.T) { + m, _ := NewMemory(MemoryConfig{}) + loader, _ := NewLoader(m, LoaderConfig{MaxLoads: 1, Timeout: time.Second}) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + started, release := make(chan struct{}), make(chan struct{}) + var calls atomic.Int32 + load := func(ctx context.Context) ([]byte, error) { + calls.Add(1) + close(started) + select { + case <-release: + return []byte("value"), nil + case <-ctx.Done(): + return nil, ctx.Err() + } + } + first := make(chan error, 1) + go func() { _, err := loader.GetOrLoad(ctx, "key", time.Minute, load); first <- err }() + <-started + second := make(chan error, 1) + go func() { + v, err := loader.GetOrLoad(context.Background(), "key", time.Minute, load) + if err == nil && string(v) != "value" { + err = errors.New("wrong value") + } + second <- err + }() + if _, err := loader.GetOrLoad(context.Background(), "other", 0, load); !errors.Is(err, ErrBusy) { + t.Fatal(err) + } + cancel() + if err := <-first; !errors.Is(err, context.Canceled) { + t.Fatal(err) + } + close(release) + if err := <-second; err != nil { + t.Fatal(err) + } + if calls.Load() != 1 { + t.Fatalf("loads=%d", calls.Load()) + } +} + +func TestLoaderPanicDeadlineAndStoreFailure(t *testing.T) { + m, _ := NewMemory(MemoryConfig{}) + l, _ := NewLoader(m, LoaderConfig{Timeout: 20 * time.Millisecond}) + if _, err := l.GetOrLoad(context.Background(), "panic", 0, func(context.Context) ([]byte, error) { panic("failure") }); err == nil { + t.Fatal("panic escaped") + } + if _, err := l.GetOrLoad(context.Background(), "panic", 0, func(ctx context.Context) ([]byte, error) { <-ctx.Done(); return nil, ctx.Err() }); !errors.Is(err, context.DeadlineExceeded) { + t.Fatal(err) + } + m.Close() + if _, err := l.GetOrLoad(context.Background(), "closed", 0, func(context.Context) ([]byte, error) { t.Fatal("store outage triggered load"); return nil, nil }); !errors.Is(err, ErrClosed) { + t.Fatal(err) + } +} diff --git a/cache/memory.go b/cache/memory.go new file mode 100644 index 0000000..458eef6 --- /dev/null +++ b/cache/memory.go @@ -0,0 +1,135 @@ +package cache + +import ( + "container/list" + "context" + "errors" + "sync" + "time" +) + +type MemoryConfig struct { + MaxEntries int `json:"maxEntries"` + MaxBytes int `json:"maxBytes"` // Sum of key and value bytes, excluding Go object overhead. +} + +type entry struct { + key string + value []byte + expires time.Time +} + +// Memory is an LRU cache with lazy expiration and no background goroutines. +type Memory struct { + mu sync.Mutex + config MemoryConfig + items map[string]*list.Element + lru *list.List + bytes int + closed bool + now func() time.Time +} + +func NewMemory(cfg MemoryConfig) (*Memory, error) { + if cfg.MaxEntries < 0 || cfg.MaxBytes < 0 { + return nil, errors.New("cache: negative capacity") + } + if cfg.MaxEntries == 0 { + cfg.MaxEntries = 4096 + } + if cfg.MaxBytes == 0 { + cfg.MaxBytes = 32 << 20 + } + return &Memory{config: cfg, items: make(map[string]*list.Element), lru: list.New(), now: time.Now}, nil +} + +func (m *Memory) Get(ctx context.Context, key string) ([]byte, error) { + if err := Check(ctx, key, 0); err != nil { + return nil, err + } + m.mu.Lock() + defer m.mu.Unlock() + if m.closed { + return nil, ErrClosed + } + e := m.items[key] + if e == nil { + return nil, ErrMiss + } + item := e.Value.(*entry) + if !item.expires.IsZero() && !m.now().Before(item.expires) { + m.remove(e) + return nil, ErrMiss + } + m.lru.MoveToFront(e) + return append([]byte(nil), item.value...), nil +} + +func (m *Memory) Set(ctx context.Context, key string, value []byte, ttl time.Duration) error { + if err := Check(ctx, key, ttl); err != nil { + return err + } + if len(value) > m.config.MaxBytes-len(key) { + return ErrTooLarge + } + m.mu.Lock() + defer m.mu.Unlock() + if m.closed { + return ErrClosed + } + if e := m.items[key]; e != nil { + m.remove(e) + } + for len(m.items) >= m.config.MaxEntries || m.bytes+len(key)+len(value) > m.config.MaxBytes { + m.remove(m.lru.Back()) + } + item := &entry{key: key, value: append([]byte(nil), value...)} + if ttl > 0 { + item.expires = m.now().Add(ttl) + } + m.items[key] = m.lru.PushFront(item) + m.bytes += len(key) + len(value) + return nil +} + +func (m *Memory) Delete(ctx context.Context, keys ...string) error { + if ctx == nil { + return errors.New("cache: nil context") + } + if err := ctx.Err(); err != nil { + return err + } + for _, key := range keys { + if err := Check(ctx, key, 0); err != nil { + return err + } + } + m.mu.Lock() + defer m.mu.Unlock() + if m.closed { + return ErrClosed + } + for _, key := range keys { + if e := m.items[key]; e != nil { + m.remove(e) + } + } + return nil +} + +func (m *Memory) remove(e *list.Element) { + item := e.Value.(*entry) + m.bytes -= len(item.key) + len(item.value) + delete(m.items, item.key) + m.lru.Remove(e) +} + +func (m *Memory) Close() error { + m.mu.Lock() + defer m.mu.Unlock() + m.closed = true + clear(m.items) + m.lru.Init() + m.bytes = 0 + return nil +} diff --git a/cache/redis/redis.go b/cache/redis/redis.go new file mode 100644 index 0000000..c932d50 --- /dev/null +++ b/cache/redis/redis.go @@ -0,0 +1,68 @@ +// Package redis adapts go-redis's pooled client to cache.Store. +package redis + +import ( + "context" + "errors" + "strings" + "time" + + "git.yuntuo.hk/k2safe/go-omni/cache" + goredis "github.com/redis/go-redis/v9" +) + +// Store borrows the client; its owner must close the client after users stop. +type Store struct { + client goredis.Cmdable + prefix string +} + +func New(client goredis.Cmdable, prefix string) (*Store, error) { + if client == nil || len(prefix) > 1024 || strings.ContainsAny(prefix, "\r\n") { + return nil, errors.New("cache/redis: invalid client or prefix") + } + return &Store{client: client, prefix: prefix}, nil +} + +func (s *Store) Get(ctx context.Context, key string) ([]byte, error) { + if err := cache.Check(ctx, key, 0); err != nil { + return nil, err + } + value, err := s.client.Get(ctx, s.prefix+key).Bytes() + if errors.Is(err, goredis.Nil) { + return nil, cache.ErrMiss + } + return value, err +} + +func (s *Store) Set(ctx context.Context, key string, value []byte, ttl time.Duration) error { + if err := cache.Check(ctx, key, ttl); err != nil { + return err + } + // Redis expiration precision is milliseconds; round up, never truncate a + // positive duration to an unbounded lifetime. + if ttl > 0 && ttl < time.Millisecond { + ttl = time.Millisecond + } + return s.client.Set(ctx, s.prefix+key, value, ttl).Err() +} + +func (s *Store) Delete(ctx context.Context, keys ...string) error { + if ctx == nil { + return errors.New("cache/redis: nil context") + } + if err := ctx.Err(); err != nil { + return err + } + names := make([]string, len(keys)) + for i, key := range keys { + if err := cache.Check(ctx, key, 0); err != nil { + return err + } + names[i] = s.prefix + key + } + if len(names) == 0 { + return nil + } + return s.client.Del(ctx, names...).Err() +} diff --git a/cache/redis/redis_test.go b/cache/redis/redis_test.go new file mode 100644 index 0000000..e4986fc --- /dev/null +++ b/cache/redis/redis_test.go @@ -0,0 +1,125 @@ +package redis + +import ( + "context" + "errors" + "fmt" + "os" + "testing" + "time" + + "git.yuntuo.hk/k2safe/go-omni/cache" + goredis "github.com/redis/go-redis/v9" +) + +type fakeClient struct { + goredis.Cmdable + key string + ttl time.Duration + failure error +} + +func (f *fakeClient) Get(ctx context.Context, key string) *goredis.StringCmd { + f.key = key + return goredis.NewStringResult("", f.failure) +} +func (f *fakeClient) Set(ctx context.Context, key string, value interface{}, ttl time.Duration) *goredis.StatusCmd { + f.key, f.ttl = key, ttl + return goredis.NewStatusResult("OK", f.failure) +} +func (f *fakeClient) Del(ctx context.Context, keys ...string) *goredis.IntCmd { + f.key = keys[0] + return goredis.NewIntResult(1, f.failure) +} + +func TestAdapterPreservesErrorsAndExpiration(t *testing.T) { + f := &fakeClient{failure: goredis.Nil} + store, _ := New(f, "app:") + if _, err := store.Get(context.Background(), "key"); !errors.Is(err, cache.ErrMiss) || f.key != "app:key" { + t.Fatal(err, f.key) + } + f.failure = errors.New("unavailable") + if _, err := store.Get(context.Background(), "key"); !errors.Is(err, f.failure) { + t.Fatal(err) + } + f.failure = nil + if err := store.Set(context.Background(), "key", []byte("value"), time.Nanosecond); err != nil || f.ttl != time.Millisecond { + t.Fatal(err, f.ttl) + } + if err := store.Set(context.Background(), "key", nil, -time.Second); err == nil { + t.Fatal("negative TTL accepted") + } + if err := store.Delete(context.Background(), "key"); err != nil || f.key != "app:key" { + t.Fatal(err) + } +} + +func TestRedisIntegration(t *testing.T) { + address := os.Getenv("GO_OMNI_REDIS_TEST_ADDR") + if address == "" { + t.Skip("set GO_OMNI_REDIS_TEST_ADDR to an isolated test Redis") + } + client := goredis.NewClient(&goredis.Options{Addr: address, PoolSize: 2, MaxRetries: -1, ContextTimeoutEnabled: true, DialTimeout: time.Second, ReadTimeout: time.Second, WriteTimeout: time.Second}) + defer client.Close() + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + if err := client.Ping(ctx).Err(); err != nil { + t.Fatal(err) + } + prefix := fmt.Sprintf("goOmniTest:%d:", time.Now().UnixNano()) + store, _ := New(client, prefix) + defer store.Delete(context.Background(), "key", "forever", "list", "canceled") + if err := store.Set(ctx, "key", []byte{0, 1, 255}, 80*time.Millisecond); err != nil { + t.Fatal(err) + } + v, err := store.Get(ctx, "key") + if err != nil || len(v) != 3 || v[2] != 255 { + t.Fatal(v, err) + } + deadline := time.Now().Add(time.Second) + for { + _, err := store.Get(ctx, "key") + if errors.Is(err, cache.ErrMiss) { + break + } + if err != nil || time.Now().After(deadline) { + t.Fatal("TTL failed", err) + } + time.Sleep(5 * time.Millisecond) + } + if err := store.Set(ctx, "forever", []byte("ok"), 0); err != nil { + t.Fatal(err) + } + if ttl := client.PTTL(ctx, prefix+"forever").Val(); ttl != -1 { + t.Fatalf("unexpected persistent TTL %v", ttl) + } + if err := store.Delete(ctx, "forever"); err != nil { + t.Fatal(err) + } + if _, err := store.Get(ctx, "forever"); !errors.Is(err, cache.ErrMiss) { + t.Fatal(err) + } + if err := client.RPush(ctx, prefix+"list", "value").Err(); err != nil { + t.Fatal(err) + } + if _, err := store.Get(ctx, "list"); err == nil || errors.Is(err, cache.ErrMiss) { + t.Fatal("Redis error treated as miss", err) + } + canceled, stop := context.WithCancel(ctx) + stop() + if err := store.Set(canceled, "canceled", []byte("value"), 0); !errors.Is(err, context.Canceled) { + t.Fatal(err) + } + if _, err := store.Get(ctx, "canceled"); !errors.Is(err, cache.ErrMiss) { + t.Fatal("canceled write reached Redis", err) + } + if err := store.Delete(ctx, "list"); err != nil { + t.Fatal(err) + } + if err := client.Close(); err != nil { + t.Fatal(err) + } + if _, err := store.Get(ctx, "key"); err == nil || errors.Is(err, cache.ErrMiss) { + t.Fatal("closed client treated as miss", err) + } +} diff --git a/cmd/go-omni/README.md b/cmd/go-omni/README.md index ff7c4a6..eda9c85 100644 --- a/cmd/go-omni/README.md +++ b/cmd/go-omni/README.md @@ -1,5 +1,12 @@ # go-omni CLI +main 分支还提供 `go-omni add orderItems --dir /path/to/project`，向已有 GoOmni 项目增加 +Handler / Service / Repository / Model 与模块注册入口。模块名为 lowerCamelCase， +目标目录存在时拒绝覆盖，已有启动代码保持原样。生成模块带可替换的内存 CRUD Repository， +导入后使用 `orderItems.Register(application, repository, authMiddleware)` 注册。 +该命令尚未包含在下方历史标签中；当前源码可用 `go build -o /tmp/go-omni ./cmd/go-omni` 构建。 +边界和验证方式见 [业务基础模块](../../MODULES.md)。 + 从远程安装，生成一个带 Handler / Service / Repository / Model 分层的服务骨架： ```sh diff --git a/cmd/go-omni/main.go b/cmd/go-omni/main.go index 7a60227..e46f111 100644 --- a/cmd/go-omni/main.go +++ b/cmd/go-omni/main.go @@ -17,6 +17,11 @@ func main() { } switch os.Args[1] { + case "add": + if err := runAdd(os.Args[2:]); err != nil { + fmt.Fprintln(os.Stderr, "go-omni add:", err) + os.Exit(1) + } case "new": if err := runNew(os.Args[2:]); err != nil { fmt.Fprintln(os.Stderr, "go-omni new:", err) @@ -104,7 +109,36 @@ func usage() { fmt.Fprintln(os.Stderr, "") fmt.Fprintln(os.Stderr, "Usage:") fmt.Fprintln(os.Stderr, " go-omni new [--module ] [--force] ") + fmt.Fprintln(os.Stderr, " go-omni add [--dir ] Generate an isolated business module") fmt.Fprintln(os.Stderr, " --framework-version Pin a remote framework version (defaults to CLI version)") fmt.Fprintln(os.Stderr, " --framework-dir Use a local framework checkout") fmt.Fprintln(os.Stderr, " --database Select generated storage (default memory)") } + +func runAdd(args []string) error { + directory, name := ".", "" + for index := 0; index < len(args); index++ { + arg := args[index] + switch { + case arg == "--dir": + if index+1 >= len(args) { + return errors.New("--dir requires a project directory") + } + index++ + directory = args[index] + case strings.HasPrefix(arg, "--dir="): + directory = strings.TrimPrefix(arg, "--dir=") + case strings.HasPrefix(arg, "-"): + return fmt.Errorf("unknown option %q", arg) + case name == "": + name = arg + default: + return fmt.Errorf("unexpected argument %q", arg) + } + } + if err := scaffold.AddModule(directory, name); err != nil { + return err + } + fmt.Printf("created internal/%s; see its README.md for route registration\n", name) + return nil +} diff --git a/cmd/go-omni/main_test.go b/cmd/go-omni/main_test.go index 64f56df..3f5fad4 100644 --- a/cmd/go-omni/main_test.go +++ b/cmd/go-omni/main_test.go @@ -32,6 +32,30 @@ func TestNewRejectsEmptyFrameworkVersion(t *testing.T) { } } +func TestAddModuleCLI(t *testing.T) { + for _, equals := range []bool{false, true} { + project := filepath.Join(t.TempDir(), "billing") + if err := runNew([]string{"--module", "example.com/billing", project}); err != nil { + t.Fatal(err) + } + args := []string{"orders", "--dir", project} + if equals { + args = []string{"--dir=" + project, "orders"} + } + if err := runAdd(args); err != nil { + t.Fatal(err) + } + if _, err := os.Stat(filepath.Join(project, "internal/orders/module.go")); err != nil { + t.Fatal(err) + } + } + for _, args := range [][]string{nil, {"orders", "--dir"}, {"orders", "--dir="}, {"orders", "extra"}, {"orders", "--unknown"}} { + if err := runAdd(args); err == nil { + t.Fatalf("accepted invalid arguments %v", args) + } + } +} + func TestNewMySQLDatabaseOptions(t *testing.T) { for _, flag := range [][]string{{"--database", "mysql"}, {"--database=mysql"}} { project := filepath.Join(t.TempDir(), "billing") diff --git a/cmd/go-omni/scaffold/addModule.go b/cmd/go-omni/scaffold/addModule.go new file mode 100644 index 0000000..b4ce249 --- /dev/null +++ b/cmd/go-omni/scaffold/addModule.go @@ -0,0 +1,109 @@ +package scaffold + +import ( + "bytes" + "errors" + "fmt" + "go/format" + "go/token" + "os" + "path/filepath" + "regexp" + "strings" + "text/template" + + "golang.org/x/mod/modfile" +) + +var moduleName = regexp.MustCompile(`^[a-z][a-zA-Z0-9]{0,47}$`) + +// AddModule creates a complete, injectable memory-backed CRUD module. It never +// edits existing routes, go.mod, migrations or files in another module. +func AddModule(directory, name string) (err error) { + if strings.TrimSpace(directory) == "" || !moduleName.MatchString(name) || token.Lookup(name).IsKeyword() || name == "main" || name == "init" || name == "internal" { + return errors.New("module name must be a non-keyword lowerCamelCase Go identifier (1..48 characters)") + } + root, err := filepath.Abs(directory) + if err != nil { + return err + } + root, err = filepath.EvalSymlinks(root) + if err != nil { + return err + } + data, err := os.ReadFile(filepath.Join(root, "go.mod")) + if err != nil { + return fmt.Errorf("read project go.mod: %w", err) + } + manifest, err := modfile.Parse("go.mod", data, nil) + if err != nil || manifest.Module == nil { + return errors.New("project go.mod must declare a valid module") + } + found := false + for _, dependency := range manifest.Require { + if dependency.Mod.Path == frameworkModule { + found = true + } + } + if !found { + return errors.New("project must already require go-omni; create it with go-omni new") + } + module := manifest.Module.Mod.Path + if strings.ContainsAny(module, "\"\\\r\n") { + return errors.New("invalid project module path") + } + paths := []string{"module.go", "model/model.go", "repository/repository.go", "service/service.go", "handler/handler.go", "README.md"} + rendered := make(map[string][]byte, len(paths)) + for _, path := range paths { + source, readErr := templateFiles.ReadFile("moduleTemplates/" + strings.ReplaceAll(path, "/", "-") + ".tmpl") + if readErr != nil { + return readErr + } + tmpl, parseErr := template.New(path).Parse(string(source)) + if parseErr != nil { + return parseErr + } + var output bytes.Buffer + if err := tmpl.Execute(&output, struct{ Module, Name string }{module, name}); err != nil { + return err + } + content := output.Bytes() + if strings.HasSuffix(path, ".go") { + content, err = format.Source(content) + if err != nil { + return err + } + } + rendered[path] = content + } + parent := filepath.Join(root, "internal") + if info, statErr := os.Lstat(parent); statErr == nil { + if !info.IsDir() || info.Mode()&os.ModeSymlink != 0 { + return errors.New("internal must be a real directory, not a symlink") + } + } else if !errors.Is(statErr, os.ErrNotExist) { + return statErr + } else if err := os.Mkdir(parent, 0755); err != nil { + return err + } + target := filepath.Join(parent, name) + // Exclusive creation protects existing files, including empty directories. + if err := os.Mkdir(target, 0755); err != nil { + return fmt.Errorf("create module (existing modules are never overwritten): %w", err) + } + defer func() { + if err != nil { + _ = os.RemoveAll(target) + } + }() + for _, path := range paths { + file := filepath.Join(target, filepath.FromSlash(path)) + if err = os.MkdirAll(filepath.Dir(file), 0755); err != nil { + return err + } + if err = os.WriteFile(file, rendered[path], 0644); err != nil { + return err + } + } + return nil +} diff --git a/cmd/go-omni/scaffold/addModule_test.go b/cmd/go-omni/scaffold/addModule_test.go new file mode 100644 index 0000000..10ba863 --- /dev/null +++ b/cmd/go-omni/scaffold/addModule_test.go @@ -0,0 +1,137 @@ +package scaffold + +import ( + "context" + "os" + "os/exec" + "path/filepath" + "strings" + "testing" + "time" +) + +func moduleProject(t *testing.T) string { + t.Helper() + root := filepath.Join(t.TempDir(), "billing") + if err := Generate(Options{Dir: root, Module: "example.com/billing", FrameworkDir: frameworkRoot(t)}); err != nil { + t.Fatal(err) + } + return root +} + +func TestAddModulePreservesProjectAndRejectsOverwrite(t *testing.T) { + root := moduleProject(t) + manifest, _ := os.ReadFile(filepath.Join(root, "go.mod")) + main, _ := os.ReadFile(filepath.Join(root, "cmd/server/main.go")) + if err := AddModule(root, "orderItems"); err != nil { + t.Fatal(err) + } + for _, path := range []string{"module.go", "model/model.go", "handler/handler.go", "service/service.go", "repository/repository.go", "README.md"} { + if _, err := os.Stat(filepath.Join(root, "internal/orderItems", path)); err != nil { + t.Fatal(err) + } + } + marker := filepath.Join(root, "internal/orderItems/keep.txt") + if err := os.WriteFile(marker, []byte("keep"), 0600); err != nil { + t.Fatal(err) + } + if err := AddModule(root, "orderItems"); err == nil { + t.Fatal("overwrote existing module") + } + if got, _ := os.ReadFile(marker); string(got) != "keep" { + t.Fatal("existing file changed") + } + if got, _ := os.ReadFile(filepath.Join(root, "go.mod")); string(got) != string(manifest) { + t.Fatal("go.mod changed") + } + if got, _ := os.ReadFile(filepath.Join(root, "cmd/server/main.go")); string(got) != string(main) { + t.Fatal("main changed") + } + for _, name := range []string{"../escape", "Order", "order_items", "func", "main", "init", "internal", "", strings.Repeat("a", 49)} { + if err := AddModule(root, name); err == nil { + t.Fatalf("accepted %q", name) + } + } +} + +func TestAddModuleRejectsForeignProjectAndSymlink(t *testing.T) { + root := t.TempDir() + if err := os.WriteFile(filepath.Join(root, "go.mod"), []byte("module example.com/foreign\ngo 1.22\n"), 0600); err != nil { + t.Fatal(err) + } + if err := AddModule(root, "orders"); err == nil { + t.Fatal("accepted foreign project") + } + if _, err := os.Stat(filepath.Join(root, "internal")); !os.IsNotExist(err) { + t.Fatal("created files before validation") + } + root = moduleProject(t) + if err := os.Rename(filepath.Join(root, "internal"), filepath.Join(root, "original")); err != nil { + t.Fatal(err) + } + if err := os.Symlink(filepath.Join(root, "original"), filepath.Join(root, "internal")); err != nil { + t.Fatal(err) + } + if err := AddModule(root, "orders"); err == nil { + t.Fatal("accepted symlink parent") + } +} + +// Opt-in because this compiles a separate consuming module and may need module +// downloads. The Linux feature acceptance runner enables it explicitly. +func TestAddModuleIntegration(t *testing.T) { + if os.Getenv("GO_OMNI_SCAFFOLD_INTEGRATION") != "1" { + t.Skip("set GO_OMNI_SCAFFOLD_INTEGRATION=1 for generated-project CRUD") + } + root := moduleProject(t) + if err := AddModule(root, "orderItems"); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(root, "internal/orderItems/module_test.go"), []byte(generatedModuleTest), 0600); err != nil { + t.Fatal(err) + } + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute) + defer cancel() + cmd := exec.CommandContext(ctx, "go", "test", "-mod=mod", "-race", "./...") + cmd.Dir = root + cmd.Env = append(os.Environ(), "GOWORK=off") + if output, err := cmd.CombinedOutput(); err != nil { + t.Fatalf("generated project: %v\n%s", err, output) + } +} + +const generatedModuleTest = `package orderItems_test +import ( + "context" + "net/http/httptest" + "strings" + "testing" + omni "git.yuntuo.hk/k2safe/go-omni" + "example.com/billing/internal/orderItems" +) +func TestGeneratedCRUD(t *testing.T) { + app:=omni.New() + repo:=orderItems.NewMemory() + orderItems.Register(app,repo,func(next omni.Handler)omni.Handler{return func(c *omni.Context)error{ + if c.Header("Authorization")!="Bearer demo" {return omni.Unauthorized("DENIED","denied")};return next(c) + }}) + request:=func(method,path,body string,authorized bool)*httptest.ResponseRecorder{ + r:=httptest.NewRequest(method,path,strings.NewReader(body));r.Header.Set("Content-Type","application/json") + if authorized {r.Header.Set("Authorization","Bearer demo")} + w:=httptest.NewRecorder();app.Engine().ServeHTTP(w,r);return w + } + if w:=request("POST","/orderItems","{\"name\":\"first\"}",false);w.Code!=401 {t.Fatal(w.Code)} + cases:=[]struct{method,path,body string;status int;contains string}{ + {"POST","/orderItems","{\"name\":\" first \"}",201,"createdAt"}, + {"GET","/orderItems/1","",200,"first"}, + {"PUT","/orderItems/1","{\"name\":\"changed\"}",200,"changed"}, + {"GET","/orderItems","",200,"changed"}, + {"POST","/orderItems","{\"name\":\" \"}",400,"INVALID_NAME"}, + {"DELETE","/orderItems/1","",204,""}, + {"GET","/orderItems/1","",404,"NOT_FOUND"}, + } + for _,tc:=range cases{w:=request(tc.method,tc.path,tc.body,true);if w.Code!=tc.status||!strings.Contains(w.Body.String(),tc.contains){t.Fatalf("%s %s: %d %s",tc.method,tc.path,w.Code,w.Body.String())}} + ctx,cancel:=context.WithCancel(context.Background());cancel() + if _,err:=repo.List(ctx);err!=context.Canceled {t.Fatal(err)} +} +` diff --git a/cmd/go-omni/scaffold/moduleTemplates/README.md.tmpl b/cmd/go-omni/scaffold/moduleTemplates/README.md.tmpl new file mode 100644 index 0000000..423e525 --- /dev/null +++ b/cmd/go-omni/scaffold/moduleTemplates/README.md.tmpl @@ -0,0 +1,15 @@ +# {{.Name}} + +在业务启动文件导入并注册： + +```go +import "{{.Module}}/internal/{{.Name}}" + +{{.Name}}.Register(application, {{.Name}}.NewMemory()) +``` + +提供 `POST /{{.Name}}`、`GET /{{.Name}}`、`GET /{{.Name}}/:id`、`PUT /{{.Name}}/:id`、 +`DELETE /{{.Name}}/:id`。写入 JSON 为 `{"name":"example"}`；输出字段为 id、name、createdAt。 +代码按 Handler / Service / Repository / Model 分层，`Register` 可注入自己的 Repository 和认证中间件。 +内存示例不持久化、不提供分页；列表适合小规模演示。接实际业务时定义领域字段、分页、鉴权和数据库实现。 +生成器不会自动修改主入口、已有模块或数据库结构，也不会覆盖已有文件。 diff --git a/cmd/go-omni/scaffold/moduleTemplates/handler-handler.go.tmpl b/cmd/go-omni/scaffold/moduleTemplates/handler-handler.go.tmpl new file mode 100644 index 0000000..ffa5819 --- /dev/null +++ b/cmd/go-omni/scaffold/moduleTemplates/handler-handler.go.tmpl @@ -0,0 +1,52 @@ +package handler + +import ( + "errors" + omni "git.yuntuo.hk/k2safe/go-omni" + "{{.Module}}/internal/{{.Name}}/model" + "{{.Module}}/internal/{{.Name}}/service" +) + +type Handler struct { service *service.Service } +type WriteRequest struct { Name string `json:"name" binding:"required,max=200"` } +func New(service *service.Service) *Handler {return &Handler{service:service}} +func (h *Handler) Register(routes *omni.RouterGroup) { + routes.POST("",h.Create) + routes.GET("",h.List) + routes.GET("/:id",h.Get) + routes.PUT("/:id",h.Update) + routes.DELETE("/:id",h.Delete) +} +func (h *Handler) Create(c *omni.Context) error { + var request WriteRequest + if err:=c.BindJSON(&request); err!=nil {return err} + item,err:=h.service.Create(c.Context(),request.Name) + if err!=nil {return publicError(err)} + return c.Created(item) +} +func (h *Handler) Get(c *omni.Context) error { + item,err:=h.service.Get(c.Context(),c.Param("id")) + if err!=nil {return publicError(err)} + return c.OK(item) +} +func (h *Handler) List(c *omni.Context) error { + items,err:=h.service.List(c.Context()) + if err!=nil {return publicError(err)} + return c.OK(items) +} +func (h *Handler) Update(c *omni.Context) error { + var request WriteRequest + if err:=c.BindJSON(&request); err!=nil {return err} + item,err:=h.service.Update(c.Context(),c.Param("id"),request.Name) + if err!=nil {return publicError(err)} + return c.OK(item) +} +func (h *Handler) Delete(c *omni.Context) error { + if err:=h.service.Delete(c.Context(),c.Param("id")); err!=nil {return publicError(err)} + return c.NoContent() +} +func publicError(err error) error { + if errors.Is(err,model.ErrNotFound) {return omni.NotFound("NOT_FOUND","record not found").WithCause(err)} + if errors.Is(err,model.ErrInvalidName) {return omni.BadRequest("INVALID_NAME","name must contain 1..200 bytes").WithCause(err)} + return err +} diff --git a/cmd/go-omni/scaffold/moduleTemplates/model-model.go.tmpl b/cmd/go-omni/scaffold/moduleTemplates/model-model.go.tmpl new file mode 100644 index 0000000..a911858 --- /dev/null +++ b/cmd/go-omni/scaffold/moduleTemplates/model-model.go.tmpl @@ -0,0 +1,15 @@ +package model + +import ( + "errors" + "time" +) + +var ErrNotFound = errors.New("{{.Name}}: record not found") +var ErrInvalidName = errors.New("{{.Name}}: name must contain 1..200 bytes") + +type Item struct { + ID string `json:"id"` + Name string `json:"name"` + CreatedAt time.Time `json:"createdAt"` +} diff --git a/cmd/go-omni/scaffold/moduleTemplates/module.go.tmpl b/cmd/go-omni/scaffold/moduleTemplates/module.go.tmpl new file mode 100644 index 0000000..6c981b6 --- /dev/null +++ b/cmd/go-omni/scaffold/moduleTemplates/module.go.tmpl @@ -0,0 +1,19 @@ +package {{.Name}} + +import ( + omni "git.yuntuo.hk/k2safe/go-omni" + "{{.Module}}/internal/{{.Name}}/handler" + "{{.Module}}/internal/{{.Name}}/repository" + "{{.Module}}/internal/{{.Name}}/service" +) + +type Repository = repository.Repository + +func NewMemory() Repository { return repository.NewMemory() } + +// Register wires this module with caller-owned storage. Pass auth middleware +// here to apply it to every route before module-specific middleware. +func Register(app *omni.App, repo Repository, middleware ...omni.Middleware) { + if repo == nil { panic("{{.Name}}: repository is required") } + handler.New(service.New(repo)).Register(app.Group("/{{.Name}}", middleware...)) +} diff --git a/cmd/go-omni/scaffold/moduleTemplates/repository-repository.go.tmpl b/cmd/go-omni/scaffold/moduleTemplates/repository-repository.go.tmpl new file mode 100644 index 0000000..99b98e7 --- /dev/null +++ b/cmd/go-omni/scaffold/moduleTemplates/repository-repository.go.tmpl @@ -0,0 +1,60 @@ +package repository + +import ( + "context" + "sort" + "strconv" + "sync" + "time" + "{{.Module}}/internal/{{.Name}}/model" +) + +type Repository interface { + Create(context.Context,string) (model.Item,error) + Get(context.Context,string) (model.Item,error) + List(context.Context) ([]model.Item,error) + Update(context.Context,string,string) (model.Item,error) + Delete(context.Context,string) error +} + +// Memory is demonstration storage; data and IDs reset with the process. +type Memory struct { mu sync.Mutex; next uint64; items map[string]model.Item } +func NewMemory() *Memory { return &Memory{items:make(map[string]model.Item)} } +func (m *Memory) Create(ctx context.Context, name string) (model.Item,error) { + m.mu.Lock(); defer m.mu.Unlock() + if err:=ctx.Err(); err!=nil {return model.Item{},err} + m.next++ + item:=model.Item{ID:strconv.FormatUint(m.next,10),Name:name,CreatedAt:time.Now().UTC()} + m.items[item.ID]=item + return item,nil +} +func (m *Memory) Get(ctx context.Context,id string) (model.Item,error) { + m.mu.Lock(); defer m.mu.Unlock() + if err:=ctx.Err(); err!=nil {return model.Item{},err} + item,ok:=m.items[id] + if !ok {return model.Item{},model.ErrNotFound} + return item,nil +} +func (m *Memory) List(ctx context.Context) ([]model.Item,error) { + m.mu.Lock(); defer m.mu.Unlock() + if err:=ctx.Err(); err!=nil {return nil,err} + items:=make([]model.Item,0,len(m.items)) + for _,item:=range m.items {items=append(items,item)} + sort.Slice(items,func(i,j int)bool{return items[i].ID200 {return "",model.ErrInvalidName} + return value,nil +} +func (s *Service) Create(ctx context.Context,value string) (model.Item,error) { + value,err:=name(value); if err!=nil {return model.Item{},err} + return s.repo.Create(ctx,value) +} +func (s *Service) Get(ctx context.Context,id string) (model.Item,error) { return s.repo.Get(ctx,id) } +func (s *Service) List(ctx context.Context) ([]model.Item,error) { return s.repo.List(ctx) } +func (s *Service) Update(ctx context.Context,id,value string) (model.Item,error) { + value,err:=name(value); if err!=nil {return model.Item{},err} + return s.repo.Update(ctx,id,value) +} +func (s *Service) Delete(ctx context.Context,id string) error {return s.repo.Delete(ctx,id)} diff --git a/cmd/go-omni/scaffold/scaffold.go b/cmd/go-omni/scaffold/scaffold.go index da903b4..52d32dc 100644 --- a/cmd/go-omni/scaffold/scaffold.go +++ b/cmd/go-omni/scaffold/scaffold.go @@ -14,7 +14,7 @@ import ( "text/template" ) -//go:embed templates +//go:embed templates moduleTemplates var templateFiles embed.FS // Options controls project generation. diff --git a/examples/modules/README.md b/examples/modules/README.md new file mode 100644 index 0000000..7bec617 --- /dev/null +++ b/examples/modules/README.md @@ -0,0 +1,33 @@ +# 基础模块组合示例 + +在仓库根目录执行： + +```sh +go run ./examples/modules +``` + +默认监听 `127.0.0.1:8080`，使用仅限开发的 `Bearer dev-token`。 +配置 `GO_OMNI_TRANSPORT=native` 可直接切换自研网络层，TLS 配置与主框架一致。 + +```sh +curl -H 'Authorization: Bearer dev-token' http://127.0.0.1:8080/api/catalog +curl -H 'Authorization: Bearer dev-token' -H 'Content-Type: application/json' \ + -H 'Idempotency-Key: report-1' -d '{"name":"report"}' http://127.0.0.1:8080/api/jobs +curl -H 'Authorization: Bearer dev-token' http://127.0.0.1:8080/api/jobs +``` + +catalog 缓存 30 秒；可通过启动环境变量 `GO_OMNI_DEMO_UPSTREAM` 指定可信 JSON 上游， +演示 HTTP 连接池、请求 ID 传递与合并回源。URL 来自服务配置，不接受请求参数传入。 + +jobs 展示固定 worker、队列、结构化日志与状态查询；相同 Idempotency-Key 和请求内容重放 +第一次 202 响应，内容不同返回 409。每分钟的统计任务演示间隔调度。退出先排空任务，再关闭 +HTTP 客户端和缓存。这个示例的任务、缓存和幂等记录均保存在进程内。 + +组合测试通过真实 TLS 连接覆盖 nethttp、native 和实验 HTTP/2，检查缓存回源次数、 +幂等重放、认证不可被重放绕过、任务只提交一次与关闭流程： + +```sh +go test -race ./examples/modules +``` + +完整 API、默认上限和适用范围见 [业务基础模块](../../MODULES.md)。 diff --git a/examples/modules/main.go b/examples/modules/main.go new file mode 100644 index 0000000..6807f32 --- /dev/null +++ b/examples/modules/main.go @@ -0,0 +1,133 @@ +// modules combines caching, outbound calls, middleware and background jobs. +package main + +import ( + "context" + "encoding/json" + "errors" + "log" + "os" + "os/signal" + "sync/atomic" + "syscall" + "time" + + omni "git.yuntuo.hk/k2safe/go-omni" + "git.yuntuo.hk/k2safe/go-omni/auth" + "git.yuntuo.hk/k2safe/go-omni/cache" + "git.yuntuo.hk/k2safe/go-omni/httpc" + "git.yuntuo.hk/k2safe/go-omni/logx" + "git.yuntuo.hk/k2safe/go-omni/middleware" + "git.yuntuo.hk/k2safe/go-omni/task" +) + +func newApp(settings omni.Config, token, upstream string) (application *omni.App, err error) { + application = omni.New(omni.WithConfig(settings)) + defer func(app *omni.App) { + if err != nil { + _ = app.Close() + } + }(application) + store, err := cache.NewMemory(cache.MemoryConfig{MaxEntries: 128, MaxBytes: 1 << 20}) + if err != nil { + return nil, err + } + application.OnShutdown(store.Close) + loader, err := cache.NewLoader(store, cache.LoaderConfig{MaxLoads: 8, Timeout: time.Second}) + if err != nil { + return application, err + } + client, err := httpc.New(httpc.Config{Timeout: time.Second, MaxRetries: 1, MaxConnsPerHost: 8}) + if err != nil { + return application, err + } + application.OnShutdown(client.Close) + jobs, err := task.New(task.Config{Workers: 2, QueueSize: 16, TaskTimeout: 2 * time.Second}) + if err != nil { + return application, err + } + // Reverse cleanup order: drain jobs before closing their dependencies. + application.OnShutdown(jobs.Close) + if err = jobs.Every("reportStats", time.Minute, func(ctx context.Context) error { + logx.InfoContext(ctx, "jobs.stats", "stats", jobs.Stats()) + return nil + }); err != nil { + return application, err + } + limiter, err := middleware.NewRateLimiter(middleware.RateLimitConfig{Rate: 20, Burst: 40, MaxKeys: 128, Key: func(c *omni.Context) string { p, _ := auth.FromContext(c); return p.UserID }}) + if err != nil { + return application, err + } + idempotency, err := middleware.NewIdempotency(middleware.IdempotencyConfig{Scope: func(c *omni.Context) (string, error) { p, _ := auth.FromContext(c); return p.UserID, nil }}) + if err != nil { + return application, err + } + api := application.Group("/api", middleware.Timeout(2*time.Second), auth.RequireAuth(auth.StaticAuthenticator{token: {UserID: "demo-user"}}), limiter.Middleware()) + api.GET("/catalog", func(c *omni.Context) error { + value, err := loader.GetOrLoad(c.Context(), "publicCatalog:v1", 30*time.Second, func(ctx context.Context) ([]byte, error) { + data := map[string]any{"name": "demo catalog"} + if upstream != "" { + if err := client.JSON(ctx, "GET", upstream, nil, &data); err != nil { + return nil, err + } + } + return json.Marshal(data) + }) + if err != nil { + return err + } + return c.OK(json.RawMessage(value)) + }) + var nextID atomic.Uint64 + api.POST("/jobs", func(c *omni.Context) error { + var input struct { + Name string `json:"name" binding:"required,max=100"` + } + if err := c.BindJSON(&input); err != nil { + return err + } + id := nextID.Add(1) + // This job deliberately outlives the HTTP request. Runner still applies + // its own deadline and shutdown cancellation; log correlation survives. + _, err := jobs.Submit(context.WithoutCancel(c.Context()), "demoJob", func(ctx context.Context) error { + logx.InfoContext(ctx, "job.finished", "jobId", id, "name", input.Name) + return nil + }) + if errors.Is(err, task.ErrQueueFull) { + return omni.NewError(503, "QUEUE_FULL", "job queue is full") + } + if err != nil { + return err + } + return c.Accepted(map[string]uint64{"jobId": id}) + }, idempotency.Middleware()) + api.GET("/jobs", func(c *omni.Context) error { return c.OK(jobs.Stats()) }) + return application, nil +} + +func run(ctx context.Context) error { + settings, err := omni.LoadConfig("") + if err != nil { + return err + } + if os.Getenv("GO_OMNI_HTTP_ADDR") == "" { + settings.Address = "127.0.0.1:8080" + } + token := os.Getenv("GO_OMNI_DEV_TOKEN") + if token == "" { + token = "dev-token" + } + application, err := newApp(settings, token, os.Getenv("GO_OMNI_DEMO_UPSTREAM")) + if err != nil { + return err + } + return application.RunContext(ctx) +} + +func main() { + ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM) + defer stop() + if err := run(ctx); err != nil { + log.Fatal(err) + } +} diff --git a/examples/modules/main_test.go b/examples/modules/main_test.go new file mode 100644 index 0000000..b85c1f8 --- /dev/null +++ b/examples/modules/main_test.go @@ -0,0 +1,126 @@ +package main + +import ( + "context" + "io" + "net" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" + "time" + + omni "git.yuntuo.hk/k2safe/go-omni" + "git.yuntuo.hk/k2safe/go-omni/internal/testutil" +) + +// Exercise the composed feature chain through real sockets and all runners. +func TestModulesConfiguredTransports(t *testing.T) { + for _, transport := range []string{omni.TransportNetHTTP, omni.TransportNative, omni.TransportHTTP2} { + t.Run(transport, func(t *testing.T) { + var loads atomic.Int32 + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + loads.Add(1) + if r.Header.Get("X-Request-ID") == "" { + t.Error("missing outbound correlation") + } + io.WriteString(w, `{"name":"upstream catalog"}`) + })) + defer upstream.Close() + settings, err := omni.LoadConfig("MODULES_TEST_") + if err != nil { + t.Fatal(err) + } + cert, key, tlsClient := testutil.CertificateFiles(t) + settings.Transport = transport + settings.TLS.CertFile = cert + settings.TLS.KeyFile = key + application, err := newApp(settings, "demo", upstream.URL) + if err != nil { + t.Fatal(err) + } + listener, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan error, 1) + go func() { done <- application.ServeContext(ctx, listener) }() + defer func() { + cancel() + select { + case err := <-done: + if err != nil { + t.Error(err) + } + case <-time.After(5 * time.Second): + listener.Close() + t.Error("server failed to stop") + } + }() + client := &http.Client{Timeout: 3 * time.Second, Transport: &http.Transport{TLSClientConfig: tlsClient, ForceAttemptHTTP2: true}} + defer client.CloseIdleConnections() + request := func(method, path, body, key, requestID, token string) (*http.Response, string) { + r, err := http.NewRequest(method, "https://"+listener.Addr().String()+path, strings.NewReader(body)) + if err != nil { + t.Fatal(err) + } + r.Header.Set("Content-Type", "application/json") + r.Header.Set("Idempotency-Key", key) + r.Header.Set("X-Request-ID", requestID) + r.Header.Set("Authorization", "Bearer "+token) + response, err := client.Do(r) + if err != nil { + t.Fatal(err) + } + defer response.Body.Close() + data, err := io.ReadAll(response.Body) + if err != nil { + t.Fatal(err) + } + if response.TLS == nil || len(response.TLS.VerifiedChains) == 0 { + t.Fatal("unverified TLS") + } + if transport == omni.TransportHTTP2 && response.ProtoMajor != 2 { + t.Fatal(response.Proto) + } + return response, string(data) + } + for n := 0; n < 2; n++ { + response, body := request("GET", "/api/catalog", "", "", "catalogRequest", "demo") + if response.StatusCode != 200 || !strings.Contains(body, "upstream catalog") { + t.Fatal(response.StatusCode, body) + } + } + if loads.Load() != 1 { + t.Fatal("cache missed", loads.Load()) + } + first, original := request("POST", "/api/jobs", `{"name":"report"}`, "job1", "firstRequest", "demo") + second, replayed := request("POST", "/api/jobs", `{"name":"report"}`, "job1", "secondRequest", "demo") + if first.StatusCode != 202 || second.StatusCode != 202 || original != replayed || second.Header.Get("Idempotency-Replayed") != "true" || second.Header.Get("X-Request-ID") != "firstRequest" { + t.Fatal(first.StatusCode, second.StatusCode, original, replayed) + } + if response, _ := request("POST", "/api/jobs", `{"name":"other"}`, "job1", "conflictRequest", "demo"); response.StatusCode != 409 { + t.Fatal(response.StatusCode) + } + if response, _ := request("POST", "/api/jobs", `{"name":"report"}`, "job1", "unauthorized", "wrong"); response.StatusCode != 401 { + t.Fatal("replay bypassed authentication", response.StatusCode) + } + deadline := time.Now().Add(time.Second) + for { + response, body := request("GET", "/api/jobs", "", "", "statsRequest", "demo") + if response.StatusCode != 200 { + t.Fatal(response.StatusCode) + } + if strings.Contains(body, `"completed":1`) { + break + } + if time.Now().After(deadline) { + t.Fatal("job did not finish once", body) + } + time.Sleep(10 * time.Millisecond) + } + }) + } +} diff --git a/go.mod b/go.mod index 44f6ab4..ba19c3d 100644 --- a/go.mod +++ b/go.mod @@ -5,11 +5,15 @@ go 1.22 require ( github.com/go-playground/validator/v10 v10.20.0 github.com/go-sql-driver/mysql v1.9.3 + github.com/redis/go-redis/v9 v9.17.3 + golang.org/x/mod v0.17.0 golang.org/x/net v0.25.0 ) require ( filippo.io/edwards25519 v1.1.0 // indirect + github.com/cespare/xxhash/v2 v2.3.0 // indirect + github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect github.com/gabriel-vasile/mimetype v1.4.3 // indirect github.com/go-playground/locales v0.14.1 // indirect github.com/go-playground/universal-translator v0.18.1 // indirect diff --git a/go.sum b/go.sum index 45c9961..2179257 100644 --- a/go.sum +++ b/go.sum @@ -1,7 +1,15 @@ filippo.io/edwards25519 v1.1.0 h1:FNf4tywRC1HmFuKW5xopWpigGjJKiJSV0Cqo0cJWDaA= filippo.io/edwards25519 v1.1.0/go.mod h1:BxyFTGdWcka3PhytdK4V28tE5sGfRvvvRV7EaN4VDT4= +github.com/bsm/ginkgo/v2 v2.12.0 h1:Ny8MWAHyOepLGlLKYmXG4IEkioBysk6GpaRTLC8zwWs= +github.com/bsm/ginkgo/v2 v2.12.0/go.mod h1:SwYbGRRDovPVboqFv0tPTcG1sN61LM1Z4ARdbAV9g4c= +github.com/bsm/gomega v1.27.10 h1:yeMWxP2pV2fG3FgAODIY8EiRE3dy0aeFYt4l7wh6yKA= +github.com/bsm/gomega v1.27.10/go.mod h1:JyEr/xRbxbtgWNi8tIEVPUYZ5Dzef52k01W3YH0H+O0= +github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs= +github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f h1:lO4WD4F/rVNCu3HqELle0jiPLLBs70cWOduZpkS1E78= +github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f/go.mod h1:cuUVRXasLTGF7a8hSLbxyZXjz+1KgoB3wDUb6vlszIc= github.com/gabriel-vasile/mimetype v1.4.3 h1:in2uUcidCuFcDKtdcBxlR0rJ1+fsokWf+uqxgUFjbI0= github.com/gabriel-vasile/mimetype v1.4.3/go.mod h1:d8uq/6HKRL6CGdk+aubisF/M5GcPfT7nKyLpA0lbSSk= github.com/go-playground/assert/v2 v2.2.0 h1:JvknZsQTYeFEAhQwI4qEt9cyV5ONwRHC+lYKSsYSR8s= @@ -18,10 +26,14 @@ github.com/leodido/go-urn v1.4.0 h1:WT9HwE9SGECu3lg4d/dIA+jxlljEa1/ffXKmRjqdmIQ= github.com/leodido/go-urn v1.4.0/go.mod h1:bvxc+MVxLKB4z00jd1z+Dvzr47oO32F/QSNjSBOlFxI= github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/redis/go-redis/v9 v9.17.3 h1:fN29NdNrE17KttK5Ndf20buqfDZwGNgoUr9qjl1DQx4= +github.com/redis/go-redis/v9 v9.17.3/go.mod h1:u410H11HMLoB+TP67dz8rL9s6QW2j76l0//kSOd3370= github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg= github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY= golang.org/x/crypto v0.23.0 h1:dIJU/v2J8Mdglj/8rJ6UUOM3Zc9zLZxVZwwxMooUSAI= golang.org/x/crypto v0.23.0/go.mod h1:CKFgDieR+mRhux2Lsu27y0fO304Db0wZe70UKqHu0v8= +golang.org/x/mod v0.17.0 h1:zY54UmvipHiNd+pm+m0x9KhZ9hl1/7QNMyxXbc6ICqA= +golang.org/x/mod v0.17.0/go.mod h1:hTbmBsO62+eylJbnUtE2MGJUyE7QWk4xUqPFrRgJ+7c= golang.org/x/net v0.25.0 h1:d/OCCoBEUq33pjydKrGQhw7IlUPI2Oylr+8qLx49kac= golang.org/x/net v0.25.0/go.mod h1:JkAGAh7GEvH74S6FOH42FLoXpXbE/aqXSrIQjXgsiwM= golang.org/x/sys v0.20.0 h1:Od9JTbYCk261bKm4M/mw7AklTlFYIa0bIp9BgSm1S8Y= diff --git a/httpc/client.go b/httpc/client.go new file mode 100644 index 0000000..7cc2540 --- /dev/null +++ b/httpc/client.go @@ -0,0 +1,247 @@ +// Package httpc provides a pooled HTTP client with an end-to-end deadline and +// conservative, bounded retries. Redirects are returned to the caller. +package httpc + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "math/rand" + "net/http" + "strconv" + "strings" + "sync/atomic" + "time" + + "git.yuntuo.hk/k2safe/go-omni/logx" +) + +var ErrClosed = errors.New("httpc: closed") +var ErrResponseTooLarge = errors.New("httpc: response exceeds limit") + +type Config struct { + Timeout time.Duration `json:"timeout"` + MaxRetries int `json:"maxRetries"` + Backoff time.Duration `json:"backoff"` + MaxBackoff time.Duration `json:"maxBackoff"` + MaxConnsPerHost int `json:"maxConnsPerHost"` + MaxResponseBytes int64 `json:"maxResponseBytes"` + DisableRequestID bool `json:"disableRequestId"` + DisableLogging bool `json:"disableLogging"` + Transport http.RoundTripper `json:"-"` // Borrowed. The caller owns its lifecycle. +} + +type Client struct { + config Config + client *http.Client + owned *http.Transport + closed atomic.Bool +} + +func New(cfg Config) (*Client, error) { + if cfg.Timeout < 0 || cfg.Backoff < 0 || cfg.MaxBackoff < 0 || cfg.MaxRetries < 0 || cfg.MaxRetries > 5 || cfg.MaxConnsPerHost < 0 || cfg.MaxResponseBytes < 0 { + return nil, errors.New("httpc: invalid configuration") + } + if cfg.Timeout == 0 { + cfg.Timeout = 10 * time.Second + } + if cfg.Backoff == 0 { + cfg.Backoff = 100 * time.Millisecond + } + if cfg.MaxBackoff == 0 { + cfg.MaxBackoff = 2 * time.Second + } + if cfg.Backoff > cfg.MaxBackoff { + return nil, errors.New("httpc: backoff exceeds maxBackoff") + } + if cfg.MaxConnsPerHost == 0 { + cfg.MaxConnsPerHost = 64 + } + if cfg.MaxResponseBytes == 0 { + cfg.MaxResponseBytes = 1 << 20 + } + if cfg.MaxResponseBytes > 1<<30 || cfg.MaxConnsPerHost > 65536 { + return nil, errors.New("httpc: response or connection limit exceeds supported capacity") + } + c := &Client{config: cfg} + transport := cfg.Transport + if transport == nil { + // Own a private pool; never mutate the process-wide DefaultTransport. + base, ok := http.DefaultTransport.(*http.Transport) + if !ok { + return nil, errors.New("httpc: provide Transport when DefaultTransport is replaced") + } + c.owned = base.Clone() + c.owned.MaxConnsPerHost = cfg.MaxConnsPerHost + c.owned.MaxIdleConnsPerHost = cfg.MaxConnsPerHost + c.owned.MaxIdleConns = cfg.MaxConnsPerHost * 2 + transport = c.owned + } + c.client = &http.Client{Transport: transport, CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }} + return c, nil +} + +// Do keeps the total deadline active until the response body is closed or +// consumed. Callers own response.Body. Only GET, HEAD and OPTIONS are retried, +// and only when their bodies can be replayed. MaxRetries=0 disables retries. +func (c *Client) Do(request *http.Request) (*http.Response, error) { + if c.closed.Load() { + return nil, ErrClosed + } + if request == nil || request.URL == nil { + return nil, errors.New("httpc: nil request or URL") + } + ctx, cancel := context.WithTimeout(request.Context(), c.config.Timeout) + base := request.Clone(ctx) + if base.Header == nil { + base.Header = make(http.Header) + } + if !c.config.DisableRequestID && base.Header.Get("X-Request-ID") == "" { + if id := logx.RequestID(ctx); id != "" { + base.Header.Set("X-Request-ID", id) + } + } + replayable := (base.Method == "" || base.Method == "GET" || base.Method == "HEAD" || base.Method == "OPTIONS") && (base.Body == nil || base.Body == http.NoBody || base.GetBody != nil) + for attempt := 0; ; attempt++ { + req := base.Clone(ctx) + if attempt > 0 && base.Body != nil && base.Body != http.NoBody { + body, err := base.GetBody() + if err != nil { + cancel() + return nil, err + } + req.Body = body + } + started := time.Now() + response, err := c.client.Do(req) + status := 0 + if response != nil { + status = response.StatusCode + } + if !c.config.DisableLogging { + logx.InfoContext(ctx, "http.client", "method", req.Method, "host", req.URL.Hostname(), "status", status, "attempt", attempt+1, "durationMs", time.Since(started).Seconds()*1000, "failed", err != nil) + } + retry := replayable && attempt < c.config.MaxRetries && ctx.Err() == nil && (err != nil || status == 429 || status == 502 || status == 503 || status == 504) + delay := c.config.Backoff + for i := 0; i < attempt && delay < c.config.MaxBackoff; i++ { + if delay > c.config.MaxBackoff/2 { + delay = c.config.MaxBackoff + } else { + delay *= 2 + } + } + if response != nil && response.Header.Get("Retry-After") != "" { + if value, ok := retryAfter(response.Header.Get("Retry-After")); ok { + delay = value + } + } else if delay > 1 { + delay = delay/2 + time.Duration(rand.Int63n(int64(delay/2)+1)) + } + if delay > c.config.MaxBackoff { + retry = false + } + if !retry { + if err != nil { + cancel() + return response, err + } + response.Body = &responseBody{ReadCloser: response.Body, cancel: cancel} + return response, nil + } + if response != nil { + if response.ContentLength >= 0 && response.ContentLength <= 32<<10 { + _, _ = io.Copy(io.Discard, response.Body) + } + _ = response.Body.Close() + } + timer := time.NewTimer(delay) + select { + case <-ctx.Done(): + timer.Stop() + cancel() + return nil, ctx.Err() + case <-timer.C: + } + } +} + +type responseBody struct { + io.ReadCloser + cancel context.CancelFunc +} + +func (b *responseBody) Read(data []byte) (int, error) { + n, err := b.ReadCloser.Read(data) + if err != nil { + b.cancel() + } + return n, err +} +func (b *responseBody) Close() error { defer b.cancel(); return b.ReadCloser.Close() } + +func retryAfter(value string) (time.Duration, bool) { + if n, err := strconv.ParseInt(strings.TrimSpace(value), 10, 32); err == nil && n >= 0 { + return time.Duration(n) * time.Second, true + } + if at, err := http.ParseTime(value); err == nil { + return max(time.Until(at), 0), true + } + return 0, false +} + +type StatusError struct{ StatusCode int } + +func (e *StatusError) Error() string { return fmt.Sprintf("httpc: HTTP status %d", e.StatusCode) } + +// JSON bounds response memory, closes the body, and decodes only 2xx responses. +// Error bodies, URLs, credentials and query strings are not included in logs. +func (c *Client) JSON(ctx context.Context, method, url string, input, output any) error { + if ctx == nil { + return errors.New("httpc: nil context") + } + var body io.Reader + if input != nil { + data, err := json.Marshal(input) + if err != nil { + return err + } + body = bytes.NewReader(data) + } + request, err := http.NewRequestWithContext(ctx, method, url, body) + if err != nil { + return err + } + if input != nil { + request.Header.Set("Content-Type", "application/json") + } + response, err := c.Do(request) + if err != nil { + return err + } + defer response.Body.Close() + if response.StatusCode < 200 || response.StatusCode >= 300 { + return &StatusError{StatusCode: response.StatusCode} + } + data, err := io.ReadAll(io.LimitReader(response.Body, c.config.MaxResponseBytes+1)) + if err != nil { + return err + } + if int64(len(data)) > c.config.MaxResponseBytes { + return ErrResponseTooLarge + } + if output == nil || len(bytes.TrimSpace(data)) == 0 { + return nil + } + return json.Unmarshal(data, output) +} + +func (c *Client) Close() error { + c.closed.Store(true) + if c.owned != nil { + c.owned.CloseIdleConnections() + } + return nil +} diff --git a/httpc/client_test.go b/httpc/client_test.go new file mode 100644 index 0000000..4ccd644 --- /dev/null +++ b/httpc/client_test.go @@ -0,0 +1,131 @@ +package httpc + +import ( + "context" + "errors" + "io" + "net" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" + "time" + + "git.yuntuo.hk/k2safe/go-omni/logx" +) + +func TestRetryCorrelationPoolAndWriteSafety(t *testing.T) { + var calls, connections atomic.Int32 + server := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + n := calls.Add(1) + if r.Header.Get("X-Request-ID") != "request-1" { + t.Error("missing correlation") + } + if n == 1 || r.Method == "POST" { + w.WriteHeader(503) + io.WriteString(w, "busy") + return + } + io.WriteString(w, `{"ok":true}`) + })) + server.Config.ConnState = func(_ net.Conn, s http.ConnState) { + if s == http.StateNew { + connections.Add(1) + } + } + server.Start() + defer server.Close() + client, err := New(Config{MaxRetries: 2, Backoff: time.Millisecond, DisableLogging: true}) + if err != nil { + t.Fatal(err) + } + defer client.Close() + ctx := logx.WithRequestID(context.Background(), "request-1") + var output struct { + OK bool `json:"ok"` + } + if err := client.JSON(ctx, "GET", server.URL, nil, &output); err != nil || !output.OK { + t.Fatal(err) + } + if calls.Load() != 2 || connections.Load() != 1 { + t.Fatalf("calls=%d connections=%d", calls.Load(), connections.Load()) + } + if err := client.JSON(ctx, "POST", server.URL, map[string]string{"name": "value"}, nil); err == nil { + t.Fatal("expected 503") + } + if calls.Load() != 3 { + t.Fatal("POST was retried") + } + client.Close() + if err := client.JSON(ctx, "GET", server.URL, nil, nil); !errors.Is(err, ErrClosed) { + t.Fatal(err) + } +} + +func TestTotalDeadlineIncludesBackoffAndBody(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/retry" { + w.Header().Set("Retry-After", "1") + w.WriteHeader(503) + return + } + w.WriteHeader(200) + w.(http.Flusher).Flush() + <-r.Context().Done() + })) + defer server.Close() + client, _ := New(Config{Timeout: 40 * time.Millisecond, MaxRetries: 3, DisableLogging: true}) + defer client.Close() + started := time.Now() + err := client.JSON(context.Background(), "GET", server.URL+"/retry", nil, nil) + if !errors.Is(err, context.DeadlineExceeded) || time.Since(started) > time.Second { + t.Fatal(err, time.Since(started)) + } + req, _ := http.NewRequest("GET", server.URL+"/body", nil) + response, err := client.Do(req) + if err != nil { + t.Fatal(err) + } + defer response.Body.Close() + if _, err := io.ReadAll(response.Body); err == nil { + t.Fatal("body deadline was canceled too early") + } +} + +func TestLimitsRedirectAndNonReplayableBody(t *testing.T) { + var calls atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + calls.Add(1) + switch r.URL.Path { + case "/redirect": + http.Redirect(w, r, "/large", 302) + case "/large": + io.WriteString(w, strings.Repeat("x", 17)) + default: + w.WriteHeader(503) + } + })) + defer server.Close() + client, _ := New(Config{MaxRetries: 2, MaxResponseBytes: 16, DisableLogging: true}) + defer client.Close() + if err := client.JSON(context.Background(), "GET", server.URL+"/large", nil, nil); !errors.Is(err, ErrResponseTooLarge) { + t.Fatal(err) + } + var status *StatusError + if err := client.JSON(context.Background(), "GET", server.URL+"/redirect", nil, nil); !errors.As(err, &status) || status.StatusCode != 302 { + t.Fatal(err) + } + if calls.Load() != 2 { + t.Fatal("followed redirect") + } + req, _ := http.NewRequest("GET", server.URL, io.NopCloser(strings.NewReader("body"))) + res, err := client.Do(req) + if err != nil { + t.Fatal(err) + } + res.Body.Close() + if calls.Load() != 3 { + t.Fatal("non-replayable body was retried") + } +} diff --git a/httpx/context.go b/httpx/context.go index 9f11b8f..612298e 100644 --- a/httpx/context.go +++ b/httpx/context.go @@ -117,6 +117,37 @@ func (c *Context) Context() context.Context { // Prefer the transport-neutral methods on Context for portable handlers. func (c *Context) Raw() http.ResponseWriter { return c.writer } +// Clone gives a buffering middleware its own response state and values map. +// Stored values themselves are shared; middleware must not mutate them concurrently. +func (c *Context) Clone(ctx context.Context, writer http.ResponseWriter) *Context { + clone := NewContext(c.request.Clone(ctx), writer, c.params) + for key, value := range c.values { + clone.values[key] = value + } + return clone +} + +// WriteResponse commits a buffered response through GoOmni's response tracking. +func (c *Context) WriteResponse(status int, headers http.Header, body []byte) error { + if c.responded { + return nil + } + if c.writer == nil { + return apierrors.Internal(nil) + } + for key, values := range headers { + c.writer.Header()[key] = append([]string(nil), values...) + } + c.responded, c.status = true, status + c.writer.WriteHeader(status) + if len(body) > 0 { + if _, err := c.writer.Write(body); err != nil { + return apierrors.Internal(err) + } + } + return nil +} + func (c *Context) Param(name string) string { return c.params[name] } func (c *Context) Query(name string) string { diff --git a/middleware/idempotency.go b/middleware/idempotency.go new file mode 100644 index 0000000..95965cc --- /dev/null +++ b/middleware/idempotency.go @@ -0,0 +1,235 @@ +package middleware + +import ( + "bytes" + "crypto/sha256" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "strings" + "sync" + "time" + + apierrors "git.yuntuo.hk/k2safe/go-omni/errors" + "git.yuntuo.hk/k2safe/go-omni/httpx" +) + +type IdempotencyConfig struct { + TTL time.Duration `json:"ttl"` + MaxEntries int `json:"maxEntries"` + MaxBytes int `json:"maxBytes"` + MaxRequestBytes int64 `json:"maxRequestBytes"` + MaxResponseBytes int `json:"maxResponseBytes"` + MaxInFlight int `json:"maxInFlight"` + // Scope must return a server-verified tenant/user identity. Register auth + // before this middleware. Client-provided identity is not authorization. + Scope func(*httpx.Context) (string, error) `json:"-"` +} + +type replay struct { + fingerprint [32]byte + status int + headers http.Header + body []byte + expires time.Time + running, failed bool + size int +} +type Idempotency struct { + mu sync.Mutex + config IdempotencyConfig + items map[[32]byte]*replay + bytes int + slots chan struct{} + now func() time.Time +} + +func NewIdempotency(cfg IdempotencyConfig) (*Idempotency, error) { + if cfg.Scope == nil || cfg.TTL < 0 || cfg.MaxEntries < 0 || cfg.MaxBytes < 0 || cfg.MaxRequestBytes < 0 || cfg.MaxResponseBytes < 0 || cfg.MaxInFlight < 0 { + return nil, errors.New("middleware: invalid idempotency configuration or missing scope") + } + if cfg.TTL == 0 { + cfg.TTL = 10 * time.Minute + } + if cfg.MaxEntries == 0 { + cfg.MaxEntries = 256 + } + if cfg.MaxBytes == 0 { + cfg.MaxBytes = 16 << 20 + } + if cfg.MaxRequestBytes == 0 { + cfg.MaxRequestBytes = 1 << 20 + } + if cfg.MaxResponseBytes == 0 { + cfg.MaxResponseBytes = 64 << 10 + } + if cfg.MaxInFlight == 0 { + cfg.MaxInFlight = 32 + } + if cfg.MaxRequestBytes > 1<<30 { + return nil, errors.New("middleware: maxRequestBytes exceeds 1 GiB") + } + return &Idempotency{config: cfg, items: make(map[[32]byte]*replay), slots: make(chan struct{}, cfg.MaxInFlight), now: time.Now}, nil +} + +func (i *Idempotency) Middleware() httpx.Middleware { + return func(next httpx.Handler) httpx.Handler { + return func(c *httpx.Context) error { + r := c.Request() + if r.Method == "GET" || r.Method == "HEAD" || r.Method == "OPTIONS" { + return next(c) + } + select { + case i.slots <- struct{}{}: + defer func() { <-i.slots }() + default: + return apierrors.New(503, "IDEMPOTENCY_BUSY", "idempotency capacity reached") + } + key := r.Header.Get("Idempotency-Key") + if len(r.Header.Values("Idempotency-Key")) != 1 || len(key) < 1 || len(key) > 128 || strings.IndexFunc(key, func(r rune) bool { return r < 33 || r > 126 }) >= 0 { + return apierrors.BadRequest("INVALID_IDEMPOTENCY_KEY", "Idempotency-Key must contain 1..128 printable ASCII characters") + } + scope, err := i.config.Scope(c) + if err != nil { + return err + } + if scope == "" || len(scope) > 1024 { + return apierrors.Unauthorized("IDEMPOTENCY_SCOPE_REQUIRED", "authenticated scope is required") + } + if r.ContentLength > i.config.MaxRequestBytes { + return apierrors.New(413, "BODY_TOO_LARGE", "request body too large") + } + var body []byte + if r.Body != nil { + body, err = io.ReadAll(io.LimitReader(r.Body, i.config.MaxRequestBytes+1)) + if err != nil { + return apierrors.BadRequest("INVALID_BODY", "unable to read request body").WithCause(err) + } + if int64(len(body)) > i.config.MaxRequestBytes { + return apierrors.New(413, "BODY_TOO_LARGE", "request body too large") + } + } + if err := c.Context().Err(); err != nil { + return err + } + originalBody := r.Body + r.Body = io.NopCloser(bytes.NewReader(body)) + defer func() { r.Body = originalBody }() + parts, _ := json.Marshal([]string{r.Method, r.URL.RequestURI(), r.Header.Get("Content-Type")}) + hash := sha256.New() + hash.Write(parts) + hash.Write(body) + var fingerprint [32]byte + copy(fingerprint[:], hash.Sum(nil)) + namespace, _ := json.Marshal([]string{scope, key}) + id := sha256.Sum256(namespace) + i.mu.Lock() + now := i.now() + for k, v := range i.items { + if !v.running && !now.Before(v.expires) { + delete(i.items, k) + i.bytes -= v.size + } + } + if saved := i.items[id]; saved != nil { + if saved.fingerprint != fingerprint { + i.mu.Unlock() + return apierrors.New(409, "IDEMPOTENCY_CONFLICT", "key was used with a different request") + } + if saved.running || saved.failed { + i.mu.Unlock() + return apierrors.New(409, "IDEMPOTENCY_UNAVAILABLE", "request is in progress or its outcome cannot be replayed") + } + status, headers, data := saved.status, saved.headers.Clone(), append([]byte(nil), saved.body...) + i.mu.Unlock() + headers.Set("Idempotency-Replayed", "true") + return c.WriteResponse(status, headers, data) + } + if len(i.items) >= i.config.MaxEntries { + i.mu.Unlock() + return apierrors.New(503, "IDEMPOTENCY_BUSY", "idempotency capacity reached") + } + saved := &replay{fingerprint: fingerprint, running: true} + i.items[id] = saved + i.mu.Unlock() + buffer := &boundedResponse{headers: c.Raw().Header().Clone(), limit: i.config.MaxResponseBytes} + // Always retain an outcome tombstone on panic/error/overflow. Deleting + // the key could execute a partially committed business operation twice. + defer func() { + i.mu.Lock() + defer i.mu.Unlock() + saved.running = false + saved.failed = true + saved.expires = i.now().Add(i.config.TTL) + size := buffer.body.Len() + for k, values := range buffer.headers { + size += len(k) + for _, v := range values { + size += len(v) + } + } + if buffer.complete && buffer.err == nil && buffer.status < 500 && buffer.headers.Get("Set-Cookie") == "" && !strings.Contains(strings.ToLower(buffer.headers.Get("Cache-Control")), "no-store") && size <= i.config.MaxBytes-i.bytes { + saved.failed = false + saved.status = buffer.status + saved.headers = buffer.headers.Clone() + saved.body = append([]byte(nil), buffer.body.Bytes()...) + saved.size = size + i.bytes += size + } + }() + child := c.Clone(c.Context(), buffer) + err = next(child) + // Keep an ambiguous outcome tombstone, but let an outer Timeout + // middleware turn deadline expiry into 504 before committing bytes. + if ctxErr := child.Context().Err(); ctxErr != nil { + return ctxErr + } + if err != nil { + child.WriteError(err) + } + if buffer.err != nil { + return apierrors.Internal(buffer.err) + } + if buffer.status == 0 { + buffer.status = 200 + } + buffer.complete = true + return c.WriteResponse(buffer.status, buffer.headers, buffer.body.Bytes()) + } + } +} + +// Intentionally lacks Flush/Hijack: only buffered, bounded endpoints are valid. +type boundedResponse struct { + headers http.Header + status, limit int + body bytes.Buffer + err error + complete bool +} + +func (b *boundedResponse) Header() http.Header { return b.headers } +func (b *boundedResponse) WriteHeader(status int) { + if status < 200 { + b.err = errors.New("idempotency: informational/upgrade responses are unsupported") + return + } + if b.status == 0 { + b.status = status + } +} +func (b *boundedResponse) Write(data []byte) (int, error) { + if b.status == 0 { + b.status = 200 + } + if b.err != nil { + return 0, b.err + } + if len(data) > b.limit-b.body.Len() { + b.err = fmt.Errorf("idempotency: response exceeds %d bytes", b.limit) + return 0, b.err + } + return b.body.Write(data) +} diff --git a/middleware/protection_test.go b/middleware/protection_test.go new file mode 100644 index 0000000..80717c3 --- /dev/null +++ b/middleware/protection_test.go @@ -0,0 +1,227 @@ +package middleware + +import ( + "context" + "errors" + "net/http/httptest" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "git.yuntuo.hk/k2safe/go-omni/httpx" +) + +func invoke(h httpx.Handler, method, body, scope, key string) *httptest.ResponseRecorder { + r := httptest.NewRequest(method, "/records?version=1", strings.NewReader(body)) + r.Header.Set("Content-Type", "application/json") + if key != "" { + r.Header.Set("Idempotency-Key", key) + } + w := httptest.NewRecorder() + c := httpx.NewContext(r, w, nil) + c.Set("verifiedIdentity", scope) + if err := h(c); err != nil { + c.WriteError(err) + } + return w +} + +func identity(c *httpx.Context) (string, error) { + v, _ := c.Get("verifiedIdentity") + s, _ := v.(string) + return s, nil +} + +func TestRateLimitRefillCapacityAndIdleExpiry(t *testing.T) { + l, err := NewRateLimiter(RateLimitConfig{Rate: 2, Burst: 2, MaxKeys: 1, IdleTTL: time.Second, Key: func(c *httpx.Context) string { s, _ := identity(c); return s }}) + if err != nil { + t.Fatal(err) + } + now := time.Unix(100, 0) + l.now = func() time.Time { return now } + h := l.Middleware()(func(c *httpx.Context) error { return c.NoContent() }) + for _, status := range []int{204, 204, 429} { + w := invoke(h, "GET", "", "a", "") + if w.Code != status { + t.Fatal(w.Code, w.Body.String()) + } + if status == 429 && w.Header().Get("Retry-After") != "1" { + t.Fatal(w.Header()) + } + } + if w := invoke(h, "GET", "", "b", ""); w.Code != 503 { + t.Fatal(w.Code) + } + now = now.Add(500 * time.Millisecond) + if w := invoke(h, "GET", "", "a", ""); w.Code != 204 { + t.Fatal(w.Code) + } + now = now.Add(time.Second) + if w := invoke(h, "GET", "", "b", ""); w.Code != 204 { + t.Fatal(w.Code) + } + if _, err := NewRateLimiter(RateLimitConfig{Rate: 1, Burst: 2, IdleTTL: time.Second}); err == nil { + t.Fatal("early bucket expiry accepted") + } +} + +func TestRateLimitConcurrentBurst(t *testing.T) { + l, _ := NewRateLimiter(RateLimitConfig{Rate: 1, Burst: 7}) + l.now = func() time.Time { return time.Unix(100, 0) } + var allowed atomic.Int32 + var wg sync.WaitGroup + for n := 0; n < 100; n++ { + wg.Add(1) + go func() { + defer wg.Done() + delay, full := l.allow("same") + if delay == 0 && !full { + allowed.Add(1) + } + }() + } + wg.Wait() + if allowed.Load() != 7 { + t.Fatal(allowed.Load()) + } +} + +func TestTimeoutCooperativeAndCommittedResponse(t *testing.T) { + for _, committed := range []bool{false, true} { + h := Timeout(5 * time.Millisecond)(func(c *httpx.Context) error { + if committed { + _ = c.Accepted("accepted") + } + <-c.Context().Done() + return c.Context().Err() + }) + w := invoke(h, "GET", "", "", "") + want := 504 + if committed { + want = 202 + } + if w.Code != want { + t.Fatal(w.Code, w.Body.String()) + } + } +} + +func TestIdempotencyReplayScopeConflictAndExpiry(t *testing.T) { + i, err := NewIdempotency(IdempotencyConfig{Scope: identity, TTL: time.Minute}) + if err != nil { + t.Fatal(err) + } + now := time.Unix(100, 0) + i.now = func() time.Time { return now } + calls := 0 + h := i.Middleware()(func(c *httpx.Context) error { calls++; return c.Created(map[string]int{"calls": calls}) }) + first := invoke(h, "POST", `{"name":"one"}`, "alice", "key1") + replay := invoke(h, "POST", `{"name":"one"}`, "alice", "key1") + if first.Code != 201 || replay.Code != 201 || first.Body.String() != replay.Body.String() || replay.Header().Get("Idempotency-Replayed") != "true" || calls != 1 { + t.Fatal(first, replay, calls) + } + if w := invoke(h, "POST", `{"name":"two"}`, "alice", "key1"); w.Code != 409 { + t.Fatal(w.Code) + } + if w := invoke(h, "POST", `{"name":"one"}`, "bob", "key1"); w.Code != 201 || calls != 2 { + t.Fatal(w.Code, calls) + } + if w := invoke(h, "POST", `{}`, "", "key1"); w.Code != 401 { + t.Fatal(w.Code) + } + if w := invoke(h, "POST", `{}`, "alice", ""); w.Code != 400 { + t.Fatal(w.Code) + } + now = now.Add(time.Minute) + if w := invoke(h, "POST", `{"name":"one"}`, "alice", "key1"); w.Code != 201 || calls != 3 { + t.Fatal(w.Code, calls) + } +} + +func TestIdempotencyActiveClaimNeverExpires(t *testing.T) { + i, _ := NewIdempotency(IdempotencyConfig{Scope: identity, TTL: time.Nanosecond, MaxInFlight: 2}) + started, release := make(chan struct{}), make(chan struct{}) + var calls atomic.Int32 + h := i.Middleware()(func(c *httpx.Context) error { calls.Add(1); close(started); <-release; return c.Created("saved") }) + done := make(chan *httptest.ResponseRecorder, 1) + go func() { done <- invoke(h, "POST", "{}", "alice", "key") }() + <-started + for n := 0; n < 30; n++ { + if w := invoke(h, "POST", "{}", "alice", "key"); w.Code != 409 { + t.Errorf("duplicate status %d", w.Code) + } + } + close(release) + if w := <-done; w.Code != 201 || calls.Load() != 1 { + t.Fatal(w.Code, calls.Load()) + } +} + +func TestIdempotencyBoundsAndAmbiguousOutcomes(t *testing.T) { + for _, mode := range []string{"panic", "error", "overflow", "cookie", "noStore", "memory"} { + t.Run(mode, func(t *testing.T) { + cfg := IdempotencyConfig{Scope: identity, MaxRequestBytes: 8, MaxResponseBytes: 128, MaxEntries: 1} + if mode == "memory" { + cfg.MaxBytes = 1 + } + i, _ := NewIdempotency(cfg) + calls := 0 + h := Recovery()(i.Middleware()(func(c *httpx.Context) error { + calls++ + switch mode { + case "panic": + panic("failure") + case "error": + return errors.New("failed") + case "overflow": + return c.OK(strings.Repeat("x", 200)) + case "cookie": + c.SetResponseHeader("Set-Cookie", "session=secret") + case "noStore": + c.SetResponseHeader("Cache-Control", "no-store") + } + return c.Created("ok") + })) + if w := invoke(h, "POST", "123456789", "alice", "key"); w.Code != 413 || calls != 0 { + t.Fatal(w.Code, calls) + } + first := invoke(h, "POST", "{}", "alice", "key") + if mode == "panic" || mode == "error" || mode == "overflow" { + if first.Code != 500 { + t.Fatal(first.Code) + } + } else if first.Code != 201 { + t.Fatal(first.Code) + } + if w := invoke(h, "POST", "{}", "alice", "key"); w.Code != 409 || calls != 1 { + t.Fatal(w.Code, calls) + } + if w := invoke(h, "POST", "{}", "alice", "newKey"); w.Code != 503 { + t.Fatal(w.Code) + } + }) + } +} + +func TestIdempotencyInFlightLimitAndTimeoutComposition(t *testing.T) { + i, _ := NewIdempotency(IdempotencyConfig{Scope: identity, MaxInFlight: 1}) + started, release := make(chan struct{}), make(chan struct{}) + h := i.Middleware()(func(c *httpx.Context) error { close(started); <-release; return c.NoContent() }) + done := make(chan struct{}) + go func() { defer close(done); invoke(h, "POST", "{}", "alice", "one") }() + <-started + if w := invoke(h, "POST", "{}", "alice", "two"); w.Code != 503 { + t.Error(w.Code) + } + close(release) + <-done + timed := Timeout(5 * time.Millisecond)(i.Middleware()(func(c *httpx.Context) error { <-c.Context().Done(); return context.DeadlineExceeded })) + if w := invoke(timed, "POST", "{}", "alice", "timeout"); w.Code != 504 { + t.Fatal(w.Code, w.Body.String()) + } + if w := invoke(timed, "POST", "{}", "alice", "timeout"); w.Code != 409 { + t.Fatal(w.Code) + } +} diff --git a/middleware/rateLimit.go b/middleware/rateLimit.go new file mode 100644 index 0000000..1ebe1c1 --- /dev/null +++ b/middleware/rateLimit.go @@ -0,0 +1,109 @@ +package middleware + +import ( + "container/list" + "errors" + "math" + "strconv" + "sync" + "time" + + apierrors "git.yuntuo.hk/k2safe/go-omni/errors" + "git.yuntuo.hk/k2safe/go-omni/httpx" +) + +type RateLimitConfig struct { + Rate float64 `json:"rate"` // Tokens per second, per key. + Burst int `json:"burst"` + MaxKeys int `json:"maxKeys"` + IdleTTL time.Duration `json:"idleTtl"` + Key func(*httpx.Context) string `json:"-"` // Default: trusted ClientIP. +} + +type bucket struct { + key string + tokens float64 + updated time.Time +} +type RateLimiter struct { + mu sync.Mutex + config RateLimitConfig + items map[string]*list.Element + lru *list.List + now func() time.Time +} + +func NewRateLimiter(cfg RateLimitConfig) (*RateLimiter, error) { + if math.IsNaN(cfg.Rate) || math.IsInf(cfg.Rate, 0) || cfg.Rate <= 0 || cfg.Burst < 1 || cfg.MaxKeys < 0 || cfg.IdleTTL < 0 { + return nil, errors.New("middleware: invalid rate limit configuration") + } + refill := float64(cfg.Burst) / cfg.Rate + if refill > (24 * time.Hour).Seconds() { + return nil, errors.New("middleware: bucket refill time must not exceed 24h") + } + minimum := time.Duration(math.Ceil(refill * float64(time.Second))) + if cfg.IdleTTL == 0 { + cfg.IdleTTL = max(5*time.Minute, minimum) + } + if cfg.IdleTTL < minimum { + return nil, errors.New("middleware: idleTtl must allow a full bucket refill") + } + if cfg.MaxKeys == 0 { + cfg.MaxKeys = 10000 + } + if cfg.Key == nil { + cfg.Key = func(c *httpx.Context) string { return c.ClientIP() } + } + return &RateLimiter{config: cfg, items: make(map[string]*list.Element), lru: list.New(), now: time.Now}, nil +} + +func (l *RateLimiter) Middleware() httpx.Middleware { + return func(next httpx.Handler) httpx.Handler { + return func(c *httpx.Context) error { + key := l.config.Key(c) + if key == "" || len(key) > 1024 { + return apierrors.BadRequest("INVALID_RATE_KEY", "invalid rate limit key") + } + delay, full := l.allow(key) + if full { + return apierrors.New(503, "RATE_LIMIT_CAPACITY", "rate limiter capacity reached") + } + if delay > 0 { + c.SetResponseHeader("Retry-After", strconv.Itoa(max(1, int(math.Ceil(delay.Seconds()))))) + return apierrors.New(429, "RATE_LIMITED", "too many requests") + } + return next(c) + } + } +} + +func (l *RateLimiter) allow(key string) (time.Duration, bool) { + l.mu.Lock() + defer l.mu.Unlock() + now := l.now() + e := l.items[key] + if e == nil { + for l.lru.Len() > 0 { + old := l.lru.Back() + if now.Sub(old.Value.(*bucket).updated) < l.config.IdleTTL { + break + } + delete(l.items, old.Value.(*bucket).key) + l.lru.Remove(old) + } + if len(l.items) >= l.config.MaxKeys { + return 0, true + } + e = l.lru.PushFront(&bucket{key: key, tokens: float64(l.config.Burst), updated: now}) + l.items[key] = e + } + b := e.Value.(*bucket) + b.tokens = min(float64(l.config.Burst), b.tokens+max(0, now.Sub(b.updated).Seconds())*l.config.Rate) + b.updated = now + l.lru.MoveToFront(e) + if b.tokens >= 1 { + b.tokens-- + return 0, false + } + return time.Duration(math.Ceil((1 - b.tokens) / l.config.Rate * float64(time.Second))), false +} diff --git a/middleware/timeout.go b/middleware/timeout.go new file mode 100644 index 0000000..02e4407 --- /dev/null +++ b/middleware/timeout.go @@ -0,0 +1,35 @@ +package middleware + +import ( + "context" + "time" + + apierrors "git.yuntuo.hk/k2safe/go-omni/errors" + "git.yuntuo.hk/k2safe/go-omni/httpx" +) + +// Timeout applies a cooperative deadline to handlers and downstream calls. +// It does not spawn a goroutine or race a handler's response. Business code must +// observe Context().Done(); an already committed response cannot be replaced. +func Timeout(duration time.Duration) httpx.Middleware { + if duration <= 0 { + panic("middleware: timeout must be positive") + } + return func(next httpx.Handler) httpx.Handler { + return func(c *httpx.Context) error { + parent := c.Context() + ctx, cancel := context.WithTimeout(parent, duration) + defer cancel() + c.WithContext(ctx) + defer c.WithContext(parent) + if ctx.Err() != nil { + return apierrors.New(504, "REQUEST_TIMEOUT", "request deadline exceeded") + } + err := next(c) + if ctx.Err() == context.DeadlineExceeded && !c.Responded() { + return apierrors.New(504, "REQUEST_TIMEOUT", "request deadline exceeded").WithCause(err) + } + return err + } + } +} diff --git a/scripts/verifyModules.py b/scripts/verifyModules.py new file mode 100644 index 0000000..a910926 --- /dev/null +++ b/scripts/verifyModules.py @@ -0,0 +1,157 @@ +#!/usr/bin/env python3 +"""Verify feature modules and real Redis in sequential resource-limited containers.""" + +import argparse +from datetime import datetime, timezone +import hashlib +import json +from pathlib import Path +import shutil +import signal +import subprocess +import tempfile +import time +import uuid + + +def capture(command): + return subprocess.check_output(command, text=True, timeout=15).strip() + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--image", default="golang:1.25") + parser.add_argument("--redis-image", default="redis:7.4-alpine") + parser.add_argument("--goproxy", default="") + parser.add_argument("--wall-timeout", type=int, default=900) + parser.add_argument("--output", type=Path, default=Path("performance/modulesLinux.json")) + args = parser.parse_args() + if args.wall_timeout < 1: + parser.error("wall timeout must be positive") + docker = shutil.which("docker") + if not docker: + raise SystemExit("docker is required") + root = Path(__file__).resolve().parent.parent + running = capture([docker, "ps", "--format", "{{.ID}} {{.Names}}"]) + if any("go-omni-" in line for line in running.splitlines()): + raise SystemExit("another go-omni validation container is active") + before = {line.split()[0] for line in running.splitlines()} + output_path = args.output.resolve() + output_path.parent.mkdir(parents=True, exist_ok=True) + log_path = output_path.with_suffix(".txt") + prefix = "go-omni-modules-" + uuid.uuid4().hex[:12] + names = [prefix + "-" + stage for stage in ("build", "redis")] + paths = [root / "go.mod", root / "go.sum", root / "httpx/context.go", Path(__file__).resolve()] + for directory in ("cache", "httpc", "task", "middleware", "cmd/go-omni", "examples/modules"): + paths.extend(p for p in (root / directory).rglob("*") if p.is_file() and p.suffix in (".go", ".tmpl")) + report = { + "status": "running", "startedAt": datetime.now(timezone.utc).isoformat(), + "testedCommit": capture(["git", "-C", str(root), "rev-parse", "HEAD"]), + "workingTreeStatus": capture(["git", "-C", str(root), "status", "--porcelain"]), + "sourceSha256": {str(p.relative_to(root)): hashlib.sha256(p.read_bytes()).hexdigest() for p in sorted(set(paths))}, + "containerLimits": {"cpus": 2, "memoryBytes": 2147483648, "memoryAndSwapBytes": 2147483648, "pidsLimit": 256}, + "wallTimeoutSeconds": args.wall_timeout, "imageIds": {}, "commands": [], "cleanedContainers": [], + "notes": [ + "Full repository race/vet/build plus independent generated-module CRUD in the compiler container.", + "Real Redis uses a statically linked non-race test binary, loopback only, network=none and no published ports.", + "At most one validation container; no pre-existing service or database is accessed.", + ], + } + output_path.write_text(json.dumps(report, indent=2) + "\n") + deadline = time.monotonic() + args.wall_timeout + failure = None + + def interrupt(signum, frame): + raise KeyboardInterrupt + + previous = {sig: signal.signal(sig, interrupt) for sig in (signal.SIGTERM, signal.SIGHUP)} + with log_path.open("w") as log: + def execute(command): + remaining = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError("module verification wall timeout") + report["commands"].append(command) + subprocess.run(command, stdout=log, stderr=subprocess.STDOUT, check=True, timeout=remaining) + log.flush() + + def container(name, arguments): + try: + execute([docker, "run", "--rm", "--name", name, + "--cpus", "2", "--memory", "2048m", "--memory-swap", "2048m", "--pids-limit", "256", + "--log-opt", "max-size=10m", "--log-opt", "max-file=1"] + arguments) + finally: + result = subprocess.run([docker, "rm", "--force", name], capture_output=True, text=True, timeout=15) + if result.returncode and "No such container" not in result.stderr: + raise RuntimeError("could not clean up " + name) + report["cleanedContainers"].append(name) + + try: + for image in (args.image, args.redis_image): + result = subprocess.run([docker, "image", "inspect", image], capture_output=True, timeout=15) + if result.returncode: + print("Pulling " + image, flush=True) + execute([docker, "pull", image]) + report["imageIds"][image] = capture([docker, "image", "inspect", image, "--format", "{{.Id}}"]) + with tempfile.TemporaryDirectory(prefix="goOmniModules") as temporary: + # The Redis stage runs as an unprivileged uid and must traverse + # this directory to execute the read-only test binary. + Path(temporary).chmod(0o755) + print("Full Linux race/vet/build and generated-project CRUD; 2 CPUs, 2048 MiB, 256 pids", flush=True) + arguments = ["-v", f"{root}:/workspace:ro", "-v", f"{temporary}:/output", "-w", "/workspace", + "-e", "GOWORK=off", "-e", "GOMAXPROCS=2", "-e", "GO_OMNI_SCAFFOLD_INTEGRATION=1"] + if args.goproxy: + arguments += ["-e", "GOPROXY=" + args.goproxy] + arguments += [args.image, "sh", "-c", + "go version && go test -race -count=1 -timeout=240s ./... && go vet ./... && go build ./... && " + "CGO_ENABLED=0 go test -c -o /output/redis.test ./cache/redis"] + container(names[0], arguments) + print("Real Redis TTL/read/write/delete/failure checks; external network disabled", flush=True) + redis_check = '''set -eu +redis-server --version +redis-server --bind 127.0.0.1 --port 6379 --save '' --appendonly no --maxmemory 32mb --daemonize yes --pidfile /tmp/redis.pid --logfile /tmp/redis.log --dir /tmp +trap 'redis-cli -h 127.0.0.1 shutdown nosave >/dev/null 2>&1 || true' EXIT +attempt=0 +until redis-cli -h 127.0.0.1 ping >/dev/null 2>&1; do + attempt=$((attempt + 1)) + if [ "$attempt" -gt 50 ]; then cat /tmp/redis.log; exit 1; fi + sleep 0.1 +done +/test/redis.test -test.v -test.run '^TestRedisIntegration$' -test.count=1 -test.timeout=30s +''' + container(names[1], ["--network", "none", "--read-only", "--cap-drop", "ALL", "--security-opt", "no-new-privileges", + "--user", "999:999", "--tmpfs", "/tmp:rw,nosuid,noexec,size=64m,mode=1777", + "-v", f"{temporary}:/test:ro", "--entrypoint", "sh", + "-e", "GOMAXPROCS=2", "-e", "GO_OMNI_REDIS_TEST_ADDR=127.0.0.1:6379", + args.redis_image, "-c", redis_check]) + log.flush() + content = log_path.read_text() + if "--- PASS: TestRedisIntegration" not in content: + raise RuntimeError("missing real Redis pass evidence") + report["status"] = "passed" + except (Exception, KeyboardInterrupt) as error: + failure = str(error) or type(error).__name__ + report["status"], report["error"] = "failed", failure + finally: + for sig, handler in previous.items(): + signal.signal(sig, handler) + log.flush() + report["logSha256"] = hashlib.sha256(log_path.read_bytes()).hexdigest() + report["finishedAt"] = datetime.now(timezone.utc).isoformat() + try: + after = set(capture([docker, "ps", "--format", "{{.ID}}"]).split()) + if not before.issubset(after): + raise RuntimeError("a pre-existing container is no longer running") + report["preExistingContainersStillRunning"] = len(before) + if failure is None and set(names) != set(report["cleanedContainers"]): + raise RuntimeError("container cleanup incomplete") + except Exception as error: + failure = str(error) + report["status"], report["error"] = "failed", failure + output_path.write_text(json.dumps(report, indent=2) + "\n") + if failure: + raise SystemExit(f"FAIL: {failure}; see {log_path}") + print(f"PASS feature module acceptance; report: {output_path}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/task/task.go b/task/task.go new file mode 100644 index 0000000..0e6f707 --- /dev/null +++ b/task/task.go @@ -0,0 +1,262 @@ +// Package task provides a bounded in-process worker pool and interval jobs. +// Jobs are not durable and are not coordinated across service replicas. +package task + +import ( + "context" + "errors" + "fmt" + "sync" + "sync/atomic" + "time" + + "git.yuntuo.hk/k2safe/go-omni/logx" +) + +var ErrClosed = errors.New("task: closed") +var ErrQueueFull = errors.New("task: queue full") + +type Func func(context.Context) error +type Config struct { + Workers int `json:"workers"` + QueueSize int `json:"queueSize"` + MaxSchedules int `json:"maxSchedules"` + TaskTimeout time.Duration `json:"taskTimeout"` + ShutdownTimeout time.Duration `json:"shutdownTimeout"` + OnError func(string, error) `json:"-"` +} + +type Future struct { + done chan struct{} + err error +} + +func (f *Future) Wait(ctx context.Context) error { + if ctx == nil { + return errors.New("task: nil context") + } + select { + case <-ctx.Done(): + return ctx.Err() + case <-f.done: + return f.err + } +} +func (f *Future) Done() <-chan struct{} { return f.done } + +type job struct { + ctx context.Context + name string + fn Func + future *Future +} + +type Runner struct { + config Config + mu sync.Mutex + queue chan job + closed bool + schedules map[string]struct{} + stop chan struct{} + done chan struct{} + wg sync.WaitGroup + ctx context.Context + cancel context.CancelFunc + active atomic.Int64 + completed atomic.Uint64 + failed atomic.Uint64 + skipped atomic.Uint64 +} + +type Stats struct { + Active int64 `json:"active"` + Queued int `json:"queued"` + Completed uint64 `json:"completed"` + Failed uint64 `json:"failed"` + Skipped uint64 `json:"skipped"` +} + +func New(cfg Config) (*Runner, error) { + if cfg.Workers < 0 || cfg.QueueSize < 0 || cfg.MaxSchedules < 0 || cfg.TaskTimeout < 0 || cfg.ShutdownTimeout < 0 { + return nil, errors.New("task: invalid configuration") + } + if cfg.Workers == 0 { + cfg.Workers = 4 + } + if cfg.QueueSize == 0 { + cfg.QueueSize = 128 + } + if cfg.MaxSchedules == 0 { + cfg.MaxSchedules = 128 + } + if cfg.TaskTimeout == 0 { + cfg.TaskTimeout = 30 * time.Second + } + if cfg.ShutdownTimeout == 0 { + cfg.ShutdownTimeout = 10 * time.Second + } + ctx, cancel := context.WithCancel(context.Background()) + r := &Runner{config: cfg, queue: make(chan job, cfg.QueueSize), schedules: make(map[string]struct{}), stop: make(chan struct{}), done: make(chan struct{}), ctx: ctx, cancel: cancel} + r.wg.Add(cfg.Workers) + for i := 0; i < cfg.Workers; i++ { + go r.worker() + } + return r, nil +} + +// Submit never blocks on a full queue. The job retains the caller's context; +// use a service-owned context for work intended to outlive an HTTP request. +func (r *Runner) Submit(ctx context.Context, name string, fn Func) (*Future, error) { + if ctx == nil || fn == nil || name == "" || len(name) > 256 { + return nil, errors.New("task: context, function and bounded name are required") + } + if err := ctx.Err(); err != nil { + return nil, err + } + r.mu.Lock() + defer r.mu.Unlock() + if r.closed { + return nil, ErrClosed + } + f := &Future{done: make(chan struct{})} + select { + case r.queue <- job{ctx: ctx, name: name, fn: fn, future: f}: + return f, nil + default: + return nil, ErrQueueFull + } +} + +// Every schedules at a fixed interval, with no initial execution. A tick is +// skipped if this schedule's previous job is queued/running or the pool is full. +func (r *Runner) Every(name string, interval time.Duration, fn Func) error { + if name == "" || len(name) > 256 || interval <= 0 || fn == nil { + return errors.New("task: invalid schedule") + } + r.mu.Lock() + defer r.mu.Unlock() + if r.closed { + return ErrClosed + } + if _, ok := r.schedules[name]; ok { + return errors.New("task: duplicate schedule name") + } + if len(r.schedules) >= r.config.MaxSchedules { + return errors.New("task: schedule capacity reached") + } + r.schedules[name] = struct{}{} + r.wg.Add(1) + go func() { + defer r.wg.Done() + ticker := time.NewTicker(interval) + defer ticker.Stop() + var previous *Future + for { + select { + case <-r.stop: + return + case <-ticker.C: + if previous != nil { + select { + case <-previous.Done(): + previous = nil + default: + r.skipped.Add(1) + continue + } + } + var err error + previous, err = r.Submit(r.ctx, name, fn) + if err != nil { + r.skipped.Add(1) + } + } + } + }() + return nil +} + +func (r *Runner) worker() { + defer r.wg.Done() + for item := range r.queue { + r.active.Add(1) + ctx, cancel := context.WithTimeout(item.ctx, r.config.TaskTimeout) + stop := context.AfterFunc(r.ctx, cancel) + if r.ctx.Err() != nil { + cancel() + } + err := run(ctx, item.fn) + // AfterFunc cancellation may still be queued when the job returns. + // Check the pool directly before publishing a successful outcome. + if err == nil { + err = r.ctx.Err() + } + stop() + cancel() + if err != nil { + r.failed.Add(1) + func() { + defer func() { _ = recover() }() + if r.config.OnError != nil { + r.config.OnError(item.name, err) + } else { + logx.ErrorContext(item.ctx, "task.failed", "task", item.name, "error", err) + } + }() + } + r.active.Add(-1) + r.completed.Add(1) + item.future.err = err + close(item.future.done) + } +} + +func run(ctx context.Context, fn Func) (err error) { + defer func() { + if value := recover(); value != nil { + err = fmt.Errorf("task: panic: %v", value) + } + }() + if err = ctx.Err(); err != nil { + return err + } + err = fn(ctx) + if err == nil { + err = ctx.Err() + } + return err +} + +// Shutdown stops schedules and submission, then drains accepted jobs. Deadline +// expiry cancels running/queued jobs and returns without waiting for code that +// ignores cancellation. Such code cannot be forcibly stopped by Go. +func (r *Runner) Shutdown(ctx context.Context) error { + if ctx == nil { + return errors.New("task: nil context") + } + r.mu.Lock() + if !r.closed { + r.closed = true + close(r.stop) + close(r.queue) + go func() { r.wg.Wait(); r.cancel(); close(r.done) }() + } + r.mu.Unlock() + select { + case <-r.done: + return nil + case <-ctx.Done(): + r.cancel() + return ctx.Err() + } +} + +func (r *Runner) Close() error { + ctx, cancel := context.WithTimeout(context.Background(), r.config.ShutdownTimeout) + defer cancel() + return r.Shutdown(ctx) +} + +func (r *Runner) Stats() Stats { + return Stats{Active: r.active.Load(), Queued: len(r.queue), Completed: r.completed.Load(), Failed: r.failed.Load(), Skipped: r.skipped.Load()} +} diff --git a/task/task_test.go b/task/task_test.go new file mode 100644 index 0000000..3418e22 --- /dev/null +++ b/task/task_test.go @@ -0,0 +1,102 @@ +package task + +import ( + "context" + "errors" + "sync/atomic" + "testing" + "time" +) + +func TestCapacityDrainPanicAndClose(t *testing.T) { + r, err := New(Config{Workers: 1, QueueSize: 1, OnError: func(string, error) { panic("callback failure") }}) + if err != nil { + t.Fatal(err) + } + defer r.Close() + started, release := make(chan struct{}), make(chan struct{}) + first, _ := r.Submit(context.Background(), "first", func(context.Context) error { close(started); <-release; return nil }) + <-started + second, _ := r.Submit(context.Background(), "panic", func(context.Context) error { panic("job failure") }) + if _, err := r.Submit(context.Background(), "overflow", func(context.Context) error { return nil }); !errors.Is(err, ErrQueueFull) { + t.Fatal(err) + } + done := make(chan error, 1) + go func() { done <- r.Close() }() + close(release) + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + if err := first.Wait(ctx); err != nil { + t.Fatal(err) + } + if err := second.Wait(ctx); err == nil { + t.Fatal("panic lost") + } + if err := <-done; err != nil { + t.Fatal(err) + } + if s := r.Stats(); s.Completed != 2 || s.Failed != 1 || s.Active != 0 || s.Queued != 0 { + t.Fatal(s) + } + if _, err := r.Submit(ctx, "late", func(context.Context) error { return nil }); !errors.Is(err, ErrClosed) { + t.Fatal(err) + } +} + +func TestShutdownDeadlineAndQueuedCancellation(t *testing.T) { + r, _ := New(Config{Workers: 1, QueueSize: 1, ShutdownTimeout: 20 * time.Millisecond, OnError: func(string, error) {}}) + started, release := make(chan struct{}), make(chan struct{}) + var queued atomic.Bool + first, _ := r.Submit(context.Background(), "ignoresCancel", func(context.Context) error { close(started); <-release; return nil }) + <-started + second, _ := r.Submit(context.Background(), "queued", func(context.Context) error { queued.Store(true); return nil }) + if err := r.Close(); !errors.Is(err, context.DeadlineExceeded) { + t.Fatal(err) + } + close(release) + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + if err := first.Wait(ctx); !errors.Is(err, context.Canceled) { + t.Fatal(err) + } + if err := second.Wait(ctx); !errors.Is(err, context.Canceled) { + t.Fatal(err) + } + if queued.Load() { + t.Fatal("ran queued work after forced shutdown") + } + if err := r.Shutdown(ctx); err != nil { + t.Fatal(err) + } +} + +func TestIntervalsDoNotOverlapAndTaskDeadline(t *testing.T) { + r, _ := New(Config{Workers: 2, TaskTimeout: 25 * time.Millisecond, OnError: func(string, error) {}}) + defer r.Close() + var active, maxActive atomic.Int32 + job := func(ctx context.Context) error { + n := active.Add(1) + if n > maxActive.Load() { + maxActive.Store(n) + } + defer active.Add(-1) + <-ctx.Done() + return ctx.Err() + } + if err := r.Every("periodic", 3*time.Millisecond, job); err != nil { + t.Fatal(err) + } + if err := r.Every("periodic", time.Second, job); err == nil { + t.Fatal("duplicate schedule accepted") + } + deadline := time.Now().Add(time.Second) + for r.Stats().Completed < 2 && time.Now().Before(deadline) { + time.Sleep(2 * time.Millisecond) + } + if err := r.Close(); err != nil { + t.Fatal(err) + } + if maxActive.Load() != 1 || r.Stats().Skipped == 0 || r.Stats().Failed < 2 { + t.Fatal(maxActive.Load(), r.Stats()) + } +}