diff --git a/cmd/cachemount/main_linux.go b/cmd/cachemount/main_linux.go new file mode 100644 index 0000000..2bcb862 --- /dev/null +++ b/cmd/cachemount/main_linux.go @@ -0,0 +1,177 @@ +// cachemount presents existing DingCache files as a read-only model directory. +package main + +import ( + "encoding/json" + "flag" + "fmt" + "github.com/hanwen/go-fuse/v2/fuse" + "log" + "os" + "os/signal" + "path/filepath" + "strings" + "syscall" + + "dingospeed/pkg/cachemount" + "dingospeed/pkg/hfprojection" +) + +func main() { + if err := run(); err != nil { + log.Fatal(err) + } +} +func run() error { + treeRoot := flag.String("tree-root", "", "discover HF and hosted models under this local cache root") + treeConfig := flag.String("tree-config", "", "JSON array of selected cache_root/namespace/repo/revision objects") + reportPath := flag.String("report", "", "write pinned version/file mapping JSON outside the mount") + manifest := flag.String("manifest", "", "JSON object mapping model-relative paths to absolute OLAH cache paths") + root := flag.String("cache-root", "", "existing HF cache root (alternative to -manifest)") + repo := flag.String("repo", "", "HF repository, e.g. owner/model") + revision := flag.String("revision", "main", "cached revision to pin") + flag.Parse() + if *treeRoot != "" || *treeConfig != "" { + if flag.NArg() != 1 || (*treeRoot != "" && *treeConfig != "") || *manifest != "" || *root != "" { + return fmt.Errorf("use exactly one of -tree-root or -tree-config and one mountpoint") + } + var selections []cachemount.Selection + if *treeRoot != "" { + var err error + selections, err = cachemount.Discover(*treeRoot) + if err != nil { + return err + } + } else { + b, err := os.ReadFile(*treeConfig) + if err != nil { + return err + } + if err = json.Unmarshal(b, &selections); err != nil { + return err + } + } + if len(selections) == 0 { + return fmt.Errorf("no model versions found or selected") + } + mountpoint, err := filepath.Abs(flag.Arg(0)) + if err != nil { + return err + } + if *reportPath != "" { + reportAbs, err := filepath.Abs(*reportPath) + if err != nil { + return err + } + rel, err := filepath.Rel(mountpoint, reportAbs) + if err != nil { + return err + } + if rel == "." || (!filepath.IsAbs(rel) && rel != ".." && !strings.HasPrefix(rel, ".."+string(filepath.Separator))) { + return fmt.Errorf("report must be outside mountpoint") + } + } + snapshot, versions, err := cachemount.PrepareTree(selections, mountpoint) + if err != nil { + return err + } + defer snapshot.Close() + server, err := cachemount.Mount(snapshot, mountpoint) + if err != nil { + return err + } + b, err := json.MarshalIndent(versions, "", " ") + if err != nil { + server.Unmount() + return err + } + if *reportPath != "" { + f, err := os.OpenFile(*reportPath, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0600) + if err != nil { + server.Unmount() + return err + } + _, err = f.Write(b) + closeErr := f.Close() + if err != nil { + server.Unmount() + return err + } + if closeErr != nil { + server.Unmount() + return closeErr + } + } + for _, v := range versions { + log.Printf("%s %s (%d files) %s", v.Status, v.SourceFilePath, len(v.Files), v.Error) + } + return serve(server) + } + if flag.NArg() != 1 || (*manifest == "") == (*root == "") { + return fmt.Errorf("usage: cachemount (-manifest files.json | -cache-root ROOT -repo OWNER/MODEL [-revision REV]) MOUNTPOINT") + } + sources := map[string]string{} + expected := map[string]int64{} + if *manifest != "" { + b, err := os.ReadFile(*manifest) + if err != nil { + return err + } + if err = json.Unmarshal(b, &sources); err != nil { + return err + } + } else { + m, paths, err := (hfprojection.Reader{Root: *root}).MountSources("models", *repo, *revision) + if err != nil { + return err + } + sources = paths + for _, f := range m.Files { + expected[f.Path] = f.Size + } + log.Printf("pinning HF commit %s", m.Commit) + } + for name, p := range sources { + if !filepath.IsAbs(p) { + return fmt.Errorf("%s: backing path must be absolute", name) + } + } + snapshot, err := cachemount.Pin(sources) + if err != nil { + return err + } + defer snapshot.Close() + for name, size := range expected { + if size >= 0 && snapshot.Files[name].Size != size { + return fmt.Errorf("%s: cache size differs from repository manifest", name) + } + } + server, err := cachemount.Mount(snapshot, flag.Arg(0)) + if err != nil { + return err + } + log.Printf("mounted %d files read-only at %s; no payload copied", len(sources), flag.Arg(0)) + return serve(server) +} + +func serve(server *fuse.Server) error { + sig := make(chan os.Signal, 1) + signal.Notify(sig, syscall.SIGINT, syscall.SIGTERM) + defer signal.Stop(sig) + done := make(chan struct{}) + defer close(done) + go func() { + for { + select { + case <-sig: + if err := server.Unmount(); err != nil { + log.Printf("unmount failed (stop readers and run fusermount3 -u): %v", err) + } + case <-done: + return + } + } + }() + server.Wait() + return nil +} diff --git a/cmd/main.go b/cmd/main.go index c5cbc6b..3c7acea 100644 --- a/cmd/main.go +++ b/cmd/main.go @@ -43,7 +43,7 @@ func init() { func newApp(s *server.HTTPServer, uploadServer *server.UploadServer, uploadCleanupServer *server.UploadCleanupServer, schedulerServer *server.SchedulerServer) *app.App { app := app.New(app.ID(id), app.Name(Name), app.Version(Version), - app.Server(s, uploadServer, uploadCleanupServer, schedulerServer)) + app.Server(s, uploadServer, uploadCleanupServer, schedulerServer, server.NewStorageProbeServer(config.SysConfig))) return app } diff --git a/cmd/wire_gen.go b/cmd/wire_gen.go index 184a28c..86b8020 100644 --- a/cmd/wire_gen.go +++ b/cmd/wire_gen.go @@ -54,7 +54,7 @@ func wireApp(configConfig *config.Config) (*app.App, func(), error) { uploadRouter := router.NewUploadRouter(uploadEcho, uploadHandler, cacheAdminHandler) uploadServer := server.NewUploadServer(configConfig, uploadRouter) uploadCleanupServer := server.NewUploadCleanupServer(uploadDao) - schedulerService := service.NewSchedulerService(schedulerDao) + schedulerService := service.NewSchedulerService(schedulerDao, metaService) schedulerServer := server.NewSchedulerServer(schedulerService, sysService, localOperationService) appApp := newApp(httpServer, uploadServer, uploadCleanupServer, schedulerServer) return appApp, func() { diff --git a/config/config-backup.yaml b/config/config-backup.yaml index f9a215f..606e8c1 100644 --- a/config/config-backup.yaml +++ b/config/config-backup.yaml @@ -63,7 +63,7 @@ diskClean: enabled: true cacheSizeLimit: 41781441855488 #38T cacheCleanStrategy: "LRU" #LRU,FIFO,LARGE_FIRST - collectTimePeriod: 1 #定期检测磁盘使用量时间周期，单位小时（H） + collectTimePeriod: 24 #定期检测磁盘使用量时间周期，单位小时（H） instanceID: mas # 唯一标识符 dynamicProxy: diff --git a/config/config-offline.yaml b/config/config-offline.yaml index 27f117e..6cfd987 100644 --- a/config/config-offline.yaml +++ b/config/config-offline.yaml @@ -57,7 +57,7 @@ diskClean: enabled: true cacheSizeLimit: 41781441855488 #38T cacheCleanStrategy: "LRU" #LRU,FIFO,LARGE_FIRST - collectTimePeriod: 1 #定期检测磁盘使用量时间周期，单位小时（H） + collectTimePeriod: 24 #定期检测磁盘使用量时间周期，单位小时（H） instanceID: mas # 唯一标识符 dynamicProxy: diff --git a/config/config.yaml b/config/config.yaml index 85f56b0..3b03521 100644 --- a/config/config.yaml +++ b/config/config.yaml @@ -79,7 +79,7 @@ diskClean: enabled: false #是否启用磁盘清理 cacheSizeLimit: 41781441855488 #38T 41781441855488 cacheCleanStrategy: "LRU" #LRU,FIFO,LARGE_FIRST - collectTimePeriod: 1 #定期检测磁盘使用量时间周期，单位小时（H） + collectTimePeriod: 24 #定期检测磁盘使用量时间周期，单位小时（H） dynamicProxy: enabled: false #是否启用动态代理，当hfNetLoc配置的地址访问异常时，会自动切换到bpHfNetLoc。 @@ -114,7 +114,7 @@ proxyPool: # url: http://10.220.70.213:8121 # weight: 1 -modelscope: +modelscope: # 可省略整个配置段；以下为程序默认值，按需覆盖 officialBaseURL: https://www.modelscope.cn # ModelScope官方基础地址 chunkSize: 8388608 # 8MB分块，16*1024*1024的数值结果 maxRetry: 5 # 超时重试次数 diff --git a/config/prometheus/cache_task_alerts.yml b/config/prometheus/cache_task_alerts.yml new file mode 100644 index 0000000..23c8896 --- /dev/null +++ b/config/prometheus/cache_task_alerts.yml @@ -0,0 +1,22 @@ +groups: + - name: cache-task-lifecycle + rules: + - alert: DingoCacheTaskStoppingStalled + expr: dingospeed_cache_transition_oldest_seconds{state=~"pausing|canceling"} > 30 + for: 1m + labels: + severity: warning + annotations: + summary: Cache task has not completed its stop barrier + - alert: DingoCacheActivityAfterStop + expr: increase(dingospeed_cache_activity_after_stop_total[5m]) > 0 + labels: + severity: critical + annotations: + summary: Cache activity occurred after confirmed stop + - alert: DingoCacheSnapshotPersistenceFailure + expr: increase(dingospeed_cache_snapshot_write_failures_total[5m]) > 0 + labels: + severity: warning + annotations: + summary: Cache lifecycle state could not be saved diff --git a/config/prometheus/dependency_alerts.yml b/config/prometheus/dependency_alerts.yml new file mode 100644 index 0000000..cc7f965 --- /dev/null +++ b/config/prometheus/dependency_alerts.yml @@ -0,0 +1,19 @@ +groups: + - name: dingospeed-dependency-observation + rules: + - alert: DingoSpeedDependencyDegraded + expr: dingospeed_dependency_state == 3 + for: 30s + labels: + severity: warning + annotations: + summary: 'Dependency capability affected: {{ $labels.dependency }}' + description: 'Instance {{ $labels.instance }} has sustained failed operations affecting {{ $labels.impact }}. This does not establish whole-node or mount failure; do not automatically restart.' + - alert: DingoSpeedDependencyUnresolved + expr: dingospeed_dependency_unresolved == 1 + for: 2m + labels: + severity: warning + annotations: + summary: 'Dependency failure lacks recovery evidence: {{ $labels.dependency }}' + description: 'Instance {{ $labels.instance }} has an unresolved observation affecting {{ $labels.impact }}. Repeated failures or absence of further traffic may explain this; inspect evidence before changing request policy.' diff --git a/config/prometheus/storage_probe_alerts.yml b/config/prometheus/storage_probe_alerts.yml new file mode 100644 index 0000000..ff1e2d6 --- /dev/null +++ b/config/prometheus/storage_probe_alerts.yml @@ -0,0 +1,19 @@ +groups: + - name: dingospeed-storage-probes + rules: + - alert: DingoSpeedStorageProbeBlocked + expr: dingospeed_storage_probe_blocked == 1 + for: 30s + labels: + severity: warning + annotations: + summary: 'Storage probe still blocked: {{ $labels.operation }}' + description: 'Instance {{ $labels.instance }} has a timed-out dedicated probe still running. Its slot remains occupied; no replacement is launched. This does not establish whole-mount failure.' + - alert: DingoSpeedStorageProbeUnresolved + expr: dingospeed_storage_probe_unresolved == 1 + for: 2m + labels: + severity: warning + annotations: + summary: 'Storage probe lacks recovery evidence: {{ $labels.operation }}' + description: 'Instance {{ $labels.instance }} has unresolved dedicated probe evidence. Compare business dependency metrics; probe success cannot clear business file errors.' diff --git a/config/prometheus/tests/storage_probe.test.yml b/config/prometheus/tests/storage_probe.test.yml new file mode 100644 index 0000000..06fdc1c --- /dev/null +++ b/config/prometheus/tests/storage_probe.test.yml @@ -0,0 +1,54 @@ +rule_files: + - ../storage_probe_alerts.yml +evaluation_interval: 5s +tests: + - name: blocked_read_late_return_and_verified_recovery + interval: 5s + input_series: + - series: 'dingospeed_storage_probe_blocked{job="dingospeed",instance="node-a",operation="read"}' + values: '1x29 0x7' + - series: 'dingospeed_storage_probe_unresolved{job="dingospeed",instance="node-a",operation="read"}' + values: '1x35 0' + - series: 'dingospeed_storage_probe_blocked{job="dingospeed",instance="node-a",operation="write"}' + values: '0x36' + - series: 'dingospeed_storage_probe_unresolved{job="dingospeed",instance="node-a",operation="write"}' + values: '0x36' + - series: 'dingospeed_storage_probe_blocked{job="dingospeed",instance="node-b",operation="read"}' + values: '0x36' + - series: 'dingospeed_storage_probe_unresolved{job="dingospeed",instance="node-b",operation="read"}' + values: '0x36' + alert_rule_test: + - eval_time: 25s + alertname: DingoSpeedStorageProbeBlocked + exp_alerts: [] + - eval_time: 30s + alertname: DingoSpeedStorageProbeBlocked + exp_alerts: + - &blocked + exp_labels: {job: dingospeed, instance: node-a, operation: read, severity: warning} + exp_annotations: + summary: 'Storage probe still blocked: read' + description: 'Instance node-a has a timed-out dedicated probe still running. Its slot remains occupied; no replacement is launched. This does not establish whole-mount failure.' + - eval_time: 145s + alertname: DingoSpeedStorageProbeBlocked + exp_alerts: [*blocked] + - eval_time: 150s + alertname: DingoSpeedStorageProbeBlocked + exp_alerts: [] + - eval_time: 115s + alertname: DingoSpeedStorageProbeUnresolved + exp_alerts: [] + - eval_time: 120s + alertname: DingoSpeedStorageProbeUnresolved + exp_alerts: + - &unresolved + exp_labels: {job: dingospeed, instance: node-a, operation: read, severity: warning} + exp_annotations: + summary: 'Storage probe lacks recovery evidence: read' + description: 'Instance node-a has unresolved dedicated probe evidence. Compare business dependency metrics; probe success cannot clear business file errors.' + - eval_time: 150s + alertname: DingoSpeedStorageProbeUnresolved + exp_alerts: [*unresolved] + - eval_time: 180s + alertname: DingoSpeedStorageProbeUnresolved + exp_alerts: [] diff --git a/deploy/cachemount-sidecar/Dockerfile b/deploy/cachemount-sidecar/Dockerfile new file mode 100644 index 0000000..f21cd44 --- /dev/null +++ b/deploy/cachemount-sidecar/Dockerfile @@ -0,0 +1,18 @@ +# Build context: this directory (source snapshot included). +FROM golang:1.24 AS build +WORKDIR /src +COPY source/ ./ +RUN CGO_ENABLED=0 go build -trimpath -o /out/cachemount ./cmd/cachemount + +FROM debian:bookworm-slim +RUN apt-get update && apt-get install -y --no-install-recommends python3 fuse3 ca-certificates \ + && rm -rf /var/lib/apt/lists/* \ + && groupadd -g 65532 nonroot \ + && useradd -u 65532 -g 65532 -d /home/nonroot nonroot \ + && mkdir -p /run/cachemount /home/nonroot \ + && chown -R 65532:65532 /run/cachemount /home/nonroot +COPY --from=build /out/cachemount /usr/local/bin/cachemount +COPY source/sidecar.py /opt/cachemount/sidecar.py +USER 65532:65532 +ENV PYTHONDONTWRITEBYTECODE=1 PYTHONUNBUFFERED=1 +ENTRYPOINT ["python3", "/opt/cachemount/sidecar.py"] diff --git a/deploy/cachemount-sidecar/README.md b/deploy/cachemount-sidecar/README.md new file mode 100644 index 0000000..ca1e683 --- /dev/null +++ b/deploy/cachemount-sidecar/README.md @@ -0,0 +1,81 @@ +# cachemount Sidecar 启动包 + +本服务只做一件事：把原始缓存仓库展示成正常模型文件目录。它不启动 Spinfield，也不选择或启动推理服务。 + +## 交付内容 + +- `Dockerfile`：构建 cachemount 镜像。 +- `config.env`：只配置原始缓存根目录和模型正文根目录。 +- `README.md`：本文。 +- `source/`：构建必需的程序源码和许可证，无需修改。 + +## 配置与构建 + +修改 `config.env` 中两个容器内路径即可，目录位置完全由接收方决定： + +```env +DINGO_CACHE_ROOT=/cache-input +DINGO_MODEL_MOUNT=/model-view/models +``` + +输入根目录必须包含完整仓库元数据与文件（例如 api/、files/），不是只有 blob 的目录。输入与输出不能重叠。正文挂载在专用共享卷的空子目录，不能覆盖已有数据。 + +在本目录构建镜像： + +```sh +docker build -t 你们的镜像地址:版本 . +``` + +镜像地址写在接收方自己的 Pod 配置中。把 config.env 的两项作为 sidecar 环境变量注入即可（例如用 ConfigMap 的 envFrom；Docker 可用 --env-file config.env）。文件不会自动被 Kubernetes 加载。 + +## 消费程序如何读取 + +启动后，服务自动发现原始仓库中的 HF 和上传仓库版本，并建立目录树，不需要配置 namespace、repo 或 revision。 + +```text +正文根目录/ + namespace/ + 仓库名/ + revision/ + config.json + 模型权重文件 + 其他仓库文件 +``` + +接收方自己选择仓库与版本，将以下路径传给推理程序： + +```text +source_file_path = 正文根目录/namespace/仓库名/revision +``` + +例如输出根目录选为 `/dingofs/data2/userdata/dingo-models`，就读取： + +```text +/dingofs/data2/userdata/dingo-models/datacanvas/hf-tiny-gpt2/main +``` + +仓库名可以包含多级路径（如 sshleifer/tiny-gpt2），应完整保留。文件名与版本清单一致，读取返回原始正文，不包含 dingocache 头。目录只读；模型本身的格式和依赖仍需被推理框架支持。 + +## 部署方需做的挂载接线 + +由接收方把本镜像加入每个读权重 Pod 的 sidecar，无需提供推理镜像、启动命令或模型选择给本服务。 + +1. 把原始仓库卷只读挂到 DINGO_CACHE_ROOT；确保 Pod 调度到任何目标节点都能读取这份数据。 +2. 准备 Pod 专用共享卷（推荐内存型 emptyDir），例如在 sidecar 挂到 `/model-view`，cachemount 在它的子目录 `/model-view/models` 建立 FUSE 挂载。这个卷不存储权重副本。 +3. 同一卷挂给推理容器。sidecar 设置 mountPropagation: Bidirectional，推理容器设置 HostToContainer。仅仅共享卷或同 Pod 不会自动共享 FUSE 子挂载。 +4. sidecar 需要 /dev/fuse 和 privileged；两个容器使用相同 UID（镜像默认 65532），共享卷父目录需允许该 UID 创建目录。当前不启用 allow_other；变更 UID 须同步镜像用户记录与权限。 +5. 先等挂载就绪，再加载模型。sidecar 就绪检查命令为 `python3 /opt/cachemount/sidecar.py check`；推理容器也要确认其模型路径确实可读。 + +两边共享卷可以挂在不同路径。消费程序使用的是推理容器内对应的正文路径，不一定与 sidecar 内的 DINGO_MODEL_MOUNT 字符串相同。挂载父目录应专用，不要覆盖其他正在使用的目录。 + +集群需允许上述 FUSE 权限和挂载传播。可用原生 sidecar 加启动探针安排启动顺序；具体 Pod/Deployment/worker 模板由接收方维护。 + +## 使用边界与验收 + +- 目录清单在启动时确定。新增仓库或更新版本后重建 Pod；当前不持续监听目录变化。 +- 缺块或损坏版本不可读，日志会报告原因；就绪不代表每个版本都完整。没有可用版本时服务启动失败。 +- cachemount 必须持续运行。发生故障时重建整个 Pod 并重新加载模型；sidecar 重启不保证推理进程已有的文件描述符或 mmap 恢复。 +- 挂载期间原始缓存不能被原地修改或截断；当前没有与 Speed 写入/清理租约联动。该目录会暴露输入仓库中的可用版本，部署方应做好访问授权与卷隔离。 +- 正式验收必须从推理容器中列目录、读取配置并加载模型。不能仅用 sidecar 自己能读作为跨容器通过的依据。 + +本包仍为未压缩目录。此前已验证独立容器的 FUSE 正文读取；Kubernetes 挂载传播、DingoFS 和实际推理部署需要接收方在目标环境验收。 diff --git a/deploy/cachemount-sidecar/config.env b/deploy/cachemount-sidecar/config.env new file mode 100644 index 0000000..f21a082 --- /dev/null +++ b/deploy/cachemount-sidecar/config.env @@ -0,0 +1,7 @@ +# 两个路径都是 cachemount 容器内的路径，由部署方自行决定。 + +# 原始缓存仓库根目录，包含 api/、files/ 等仓库数据 +DINGO_CACHE_ROOT=/cache-input + +# 正常模型文件的展示根目录，放在专用共享卷的空子目录中 +DINGO_MODEL_MOUNT=/model-view/models diff --git a/deploy/cachemount-sidecar/source/LICENSE b/deploy/cachemount-sidecar/source/LICENSE new file mode 100644 index 0000000..c26ac57 --- /dev/null +++ b/deploy/cachemount-sidecar/source/LICENSE @@ -0,0 +1,13 @@ + Copyright (c) 2025 dingodb.com, Inc. All Rights Reserved + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http:www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. \ No newline at end of file diff --git a/deploy/cachemount-sidecar/source/cmd/cachemount/main_linux.go b/deploy/cachemount-sidecar/source/cmd/cachemount/main_linux.go new file mode 100644 index 0000000..2bcb862 --- /dev/null +++ b/deploy/cachemount-sidecar/source/cmd/cachemount/main_linux.go @@ -0,0 +1,177 @@ +// cachemount presents existing DingCache files as a read-only model directory. +package main + +import ( + "encoding/json" + "flag" + "fmt" + "github.com/hanwen/go-fuse/v2/fuse" + "log" + "os" + "os/signal" + "path/filepath" + "strings" + "syscall" + + "dingospeed/pkg/cachemount" + "dingospeed/pkg/hfprojection" +) + +func main() { + if err := run(); err != nil { + log.Fatal(err) + } +} +func run() error { + treeRoot := flag.String("tree-root", "", "discover HF and hosted models under this local cache root") + treeConfig := flag.String("tree-config", "", "JSON array of selected cache_root/namespace/repo/revision objects") + reportPath := flag.String("report", "", "write pinned version/file mapping JSON outside the mount") + manifest := flag.String("manifest", "", "JSON object mapping model-relative paths to absolute OLAH cache paths") + root := flag.String("cache-root", "", "existing HF cache root (alternative to -manifest)") + repo := flag.String("repo", "", "HF repository, e.g. owner/model") + revision := flag.String("revision", "main", "cached revision to pin") + flag.Parse() + if *treeRoot != "" || *treeConfig != "" { + if flag.NArg() != 1 || (*treeRoot != "" && *treeConfig != "") || *manifest != "" || *root != "" { + return fmt.Errorf("use exactly one of -tree-root or -tree-config and one mountpoint") + } + var selections []cachemount.Selection + if *treeRoot != "" { + var err error + selections, err = cachemount.Discover(*treeRoot) + if err != nil { + return err + } + } else { + b, err := os.ReadFile(*treeConfig) + if err != nil { + return err + } + if err = json.Unmarshal(b, &selections); err != nil { + return err + } + } + if len(selections) == 0 { + return fmt.Errorf("no model versions found or selected") + } + mountpoint, err := filepath.Abs(flag.Arg(0)) + if err != nil { + return err + } + if *reportPath != "" { + reportAbs, err := filepath.Abs(*reportPath) + if err != nil { + return err + } + rel, err := filepath.Rel(mountpoint, reportAbs) + if err != nil { + return err + } + if rel == "." || (!filepath.IsAbs(rel) && rel != ".." && !strings.HasPrefix(rel, ".."+string(filepath.Separator))) { + return fmt.Errorf("report must be outside mountpoint") + } + } + snapshot, versions, err := cachemount.PrepareTree(selections, mountpoint) + if err != nil { + return err + } + defer snapshot.Close() + server, err := cachemount.Mount(snapshot, mountpoint) + if err != nil { + return err + } + b, err := json.MarshalIndent(versions, "", " ") + if err != nil { + server.Unmount() + return err + } + if *reportPath != "" { + f, err := os.OpenFile(*reportPath, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0600) + if err != nil { + server.Unmount() + return err + } + _, err = f.Write(b) + closeErr := f.Close() + if err != nil { + server.Unmount() + return err + } + if closeErr != nil { + server.Unmount() + return closeErr + } + } + for _, v := range versions { + log.Printf("%s %s (%d files) %s", v.Status, v.SourceFilePath, len(v.Files), v.Error) + } + return serve(server) + } + if flag.NArg() != 1 || (*manifest == "") == (*root == "") { + return fmt.Errorf("usage: cachemount (-manifest files.json | -cache-root ROOT -repo OWNER/MODEL [-revision REV]) MOUNTPOINT") + } + sources := map[string]string{} + expected := map[string]int64{} + if *manifest != "" { + b, err := os.ReadFile(*manifest) + if err != nil { + return err + } + if err = json.Unmarshal(b, &sources); err != nil { + return err + } + } else { + m, paths, err := (hfprojection.Reader{Root: *root}).MountSources("models", *repo, *revision) + if err != nil { + return err + } + sources = paths + for _, f := range m.Files { + expected[f.Path] = f.Size + } + log.Printf("pinning HF commit %s", m.Commit) + } + for name, p := range sources { + if !filepath.IsAbs(p) { + return fmt.Errorf("%s: backing path must be absolute", name) + } + } + snapshot, err := cachemount.Pin(sources) + if err != nil { + return err + } + defer snapshot.Close() + for name, size := range expected { + if size >= 0 && snapshot.Files[name].Size != size { + return fmt.Errorf("%s: cache size differs from repository manifest", name) + } + } + server, err := cachemount.Mount(snapshot, flag.Arg(0)) + if err != nil { + return err + } + log.Printf("mounted %d files read-only at %s; no payload copied", len(sources), flag.Arg(0)) + return serve(server) +} + +func serve(server *fuse.Server) error { + sig := make(chan os.Signal, 1) + signal.Notify(sig, syscall.SIGINT, syscall.SIGTERM) + defer signal.Stop(sig) + done := make(chan struct{}) + defer close(done) + go func() { + for { + select { + case <-sig: + if err := server.Unmount(); err != nil { + log.Printf("unmount failed (stop readers and run fusermount3 -u): %v", err) + } + case <-done: + return + } + } + }() + server.Wait() + return nil +} diff --git a/deploy/cachemount-sidecar/source/go.mod b/deploy/cachemount-sidecar/source/go.mod new file mode 100644 index 0000000..be88a61 --- /dev/null +++ b/deploy/cachemount-sidecar/source/go.mod @@ -0,0 +1,66 @@ +module dingospeed + +go 1.23.0 + +toolchain go1.23.1 + +require ( + github.com/andybalholm/brotli v1.1.1 + github.com/avast/retry-go v3.0.0+incompatible + github.com/bytedance/sonic v1.15.2 + github.com/go-playground/validator/v10 v10.27.0 + github.com/go-redsync/redsync/v4 v4.13.0 + github.com/gofrs/flock v0.8.1 + github.com/google/uuid v1.6.0 + github.com/google/wire v0.6.0 + github.com/hanwen/go-fuse/v2 v2.8.0 + github.com/klauspost/compress v1.18.0 + github.com/labstack/echo/v4 v4.13.3 + github.com/labstack/gommon v0.4.2 + github.com/patrickmn/go-cache v2.1.0+incompatible + github.com/prometheus/client_golang v1.22.0 + github.com/robfig/cron/v3 v3.0.1 + github.com/shirou/gopsutil v3.21.11+incompatible + go.uber.org/zap v1.24.0 + golang.org/x/sync v0.13.0 + golang.org/x/sys v0.35.0 + google.golang.org/grpc v1.73.0 + google.golang.org/protobuf v1.36.6 + gopkg.in/natefinch/lumberjack.v2 v2.2.1 + gopkg.in/yaml.v3 v3.0.1 +) + +require ( + github.com/beorn7/perks v1.0.1 // indirect + github.com/bytedance/gopkg v0.1.3 // indirect + github.com/bytedance/sonic/loader v0.5.1 // indirect + github.com/cespare/xxhash/v2 v2.3.0 // indirect + github.com/cloudwego/base64x v0.1.6 // indirect + github.com/gabriel-vasile/mimetype v1.4.8 // indirect + github.com/go-ole/go-ole v1.2.6 // indirect + github.com/go-playground/locales v0.14.1 // indirect + github.com/go-playground/universal-translator v0.18.1 // indirect + github.com/hashicorp/errwrap v1.1.0 // indirect + github.com/hashicorp/go-multierror v1.1.1 // indirect + github.com/klauspost/cpuid/v2 v2.2.10 // indirect + github.com/kr/text v0.2.0 // indirect + github.com/leodido/go-urn v1.4.0 // indirect + github.com/mattn/go-colorable v0.1.14 // indirect + github.com/mattn/go-isatty v0.0.20 // indirect + github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect + github.com/pkg/errors v0.9.1 // indirect + github.com/prometheus/client_model v0.6.1 // indirect + github.com/prometheus/common v0.62.0 // indirect + github.com/prometheus/procfs v0.15.1 // indirect + github.com/twitchyliquid64/golang-asm v0.15.1 // indirect + github.com/valyala/bytebufferpool v1.0.0 // indirect + github.com/valyala/fasttemplate v1.2.2 // indirect + github.com/yusufpapurcu/wmi v1.2.4 // indirect + go.uber.org/atomic v1.10.0 // indirect + go.uber.org/multierr v1.9.0 // indirect + golang.org/x/arch v0.15.0 // indirect + golang.org/x/crypto v0.37.0 // indirect + golang.org/x/net v0.39.0 // indirect + golang.org/x/text v0.24.0 // indirect + google.golang.org/genproto/googleapis/rpc v0.0.0-20250324211829-b45e905df463 // indirect +) diff --git a/deploy/cachemount-sidecar/source/go.sum b/deploy/cachemount-sidecar/source/go.sum new file mode 100644 index 0000000..06d1855 --- /dev/null +++ b/deploy/cachemount-sidecar/source/go.sum @@ -0,0 +1,249 @@ +github.com/andybalholm/brotli v1.1.1 h1:PR2pgnyFznKEugtsUo0xLdDop5SKXd5Qf5ysW+7XdTA= +github.com/andybalholm/brotli v1.1.1/go.mod h1:05ib4cKhjx3OQYUY22hTVd34Bc8upXjOLL2rKwwZBoA= +github.com/avast/retry-go v3.0.0+incompatible h1:4SOWQ7Qs+oroOTQOYnAHqelpCO0biHSxpiH9JdtuBj0= +github.com/avast/retry-go v3.0.0+incompatible/go.mod h1:XtSnn+n/sHqQIpZ10K1qAevBhOOCWBLXXy3hyiqqBrY= +github.com/benbjohnson/clock v1.1.0 h1:Q92kusRqC1XV2MjkWETPvjJVqKetz1OzxZB7mHJLju8= +github.com/benbjohnson/clock v1.1.0/go.mod h1:J11/hYXuz8f4ySSvYwY0FKfm+ezbsZBKZxNJlLklBHA= +github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM= +github.com/beorn7/perks v1.0.1/go.mod h1:G2ZrVWU2WbWT9wwq4/hrbKbnv/1ERSJQ0ibhJ6rlkpw= +github.com/bytedance/gopkg v0.1.3 h1:TPBSwH8RsouGCBcMBktLt1AymVo2TVsBVCY4b6TnZ/M= +github.com/bytedance/gopkg v0.1.3/go.mod h1:576VvJ+eJgyCzdjS+c4+77QF3p7ubbtiKARP3TxducM= +github.com/bytedance/sonic v1.13.2 h1:8/H1FempDZqC4VqjptGo14QQlJx8VdZJegxs6wwfqpQ= +github.com/bytedance/sonic v1.13.2/go.mod h1:o68xyaF9u2gvVBuGHPlUVCy+ZfmNNO5ETf1+KgkJhz4= +github.com/bytedance/sonic v1.15.2 h1:90H+rcF/FwLXwfB1cudOLq/je83n683Utf4Cbp0xHCo= +github.com/bytedance/sonic v1.15.2/go.mod h1:mT2NbXunuaEbnZ+mRIX/vYqKISmgEuHFDI4UzmKx2SA= +github.com/bytedance/sonic/loader v0.1.1/go.mod h1:ncP89zfokxS5LZrJxl5z0UJcsk4M4yY2JpfqGeCtNLU= +github.com/bytedance/sonic/loader v0.2.4 h1:ZWCw4stuXUsn1/+zQDqeE7JKP+QO47tz7QCNan80NzY= +github.com/bytedance/sonic/loader v0.2.4/go.mod h1:N8A3vUdtUebEY2/VQC0MyhYeKUFosQU6FxH2JmUe6VI= +github.com/bytedance/sonic/loader v0.5.1 h1:Ygpfa9zwRCCKSlrp5bBP/b/Xzc3VxsAW+5NIYXrOOpI= +github.com/bytedance/sonic/loader v0.5.1/go.mod h1:AR4NYCk5DdzZizZ5djGqQ92eEhCCcdf5x77udYiSJRo= +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/cloudwego/base64x v0.1.5 h1:XPciSp1xaq2VCSt6lF0phncD4koWyULpl5bUxbfCyP4= +github.com/cloudwego/base64x v0.1.5/go.mod h1:0zlkT4Wn5C6NdauXdJRhSKRlJvmclQ1hhJgA0rcu/8w= +github.com/cloudwego/base64x v0.1.6 h1:t11wG9AECkCDk5fMSoxmufanudBtJ+/HemLstXDLI2M= +github.com/cloudwego/base64x v0.1.6/go.mod h1:OFcloc187FXDaYHvrNIjxSe8ncn0OOM8gEHfghB2IPU= +github.com/cloudwego/iasm v0.2.0/go.mod h1:8rXZaNYT2n95jn+zTI1sDr+IgcD2GVs0nlbbQPiEFhY= +github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E= +github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +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.8 h1:FfZ3gj38NjllZIeJAmMhr+qKL8Wu+nOoI3GqacKw1NM= +github.com/gabriel-vasile/mimetype v1.4.8/go.mod h1:ByKUIKGjh1ODkGM1asKUbQZOLGrPjydw3hYPU2YU9t8= +github.com/go-logr/logr v1.4.2 h1:6pFjapn8bFcIbiKo3XT4j/BhANplGihG6tvd+8rYgrY= +github.com/go-logr/logr v1.4.2/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY= +github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag= +github.com/go-logr/stdr v1.2.2/go.mod h1:mMo/vtBO5dYbehREoey6XUKy/eSumjCCveDpRre4VKE= +github.com/go-ole/go-ole v1.2.6 h1:/Fpf6oFPoeFik9ty7siob0G6Ke8QvQEuVcuChpwXzpY= +github.com/go-ole/go-ole v1.2.6/go.mod h1:pprOEPIfldk/42T2oK7lQ4v4JSDwmV0As9GaiUsvbm0= +github.com/go-playground/assert/v2 v2.2.0 h1:JvknZsQTYeFEAhQwI4qEt9cyV5ONwRHC+lYKSsYSR8s= +github.com/go-playground/assert/v2 v2.2.0/go.mod h1:VDjEfimB/XKnb+ZQfWdccd7VUvScMdVu0Titje2rxJ4= +github.com/go-playground/locales v0.14.1 h1:EWaQ/wswjilfKLTECiXz7Rh+3BjFhfDFKv/oXslEjJA= +github.com/go-playground/locales v0.14.1/go.mod h1:hxrqLVvrK65+Rwrd5Fc6F2O76J/NuW9t0sjnWqG1slY= +github.com/go-playground/universal-translator v0.18.1 h1:Bcnm0ZwsGyWbCzImXv+pAJnYK9S473LQFuzCbDbfSFY= +github.com/go-playground/universal-translator v0.18.1/go.mod h1:xekY+UJKNuX9WP91TpwSH2VMlDf28Uj24BCp08ZFTUY= +github.com/go-playground/validator/v10 v10.27.0 h1:w8+XrWVMhGkxOaaowyKH35gFydVHOvC0/uWoy2Fzwn4= +github.com/go-playground/validator/v10 v10.27.0/go.mod h1:I5QpIEbmr8On7W0TktmJAumgzX4CA1XNl4ZmDuVHKKo= +github.com/go-redis/redis v6.15.9+incompatible h1:K0pv1D7EQUjfyoMql+r/jZqCLizCGKFlFgcHWWmHQjg= +github.com/go-redis/redis v6.15.9+incompatible/go.mod h1:NAIEuMOZ/fxfXJIrKDQDz8wamY7mA7PouImQ2Jvg6kA= +github.com/go-redis/redis/v7 v7.4.1 h1:PASvf36gyUpr2zdOUS/9Zqc80GbM+9BDyiJSJDDOrTI= +github.com/go-redis/redis/v7 v7.4.1/go.mod h1:JDNMw23GTyLNC4GZu9njt15ctBQVn7xjRfnwdHj/Dcg= +github.com/go-redis/redis/v8 v8.11.5 h1:AcZZR7igkdvfVmQTPnu9WE37LRrO/YrBH5zWyjDC0oI= +github.com/go-redis/redis/v8 v8.11.5/go.mod h1:gREzHqY1hg6oD9ngVRbLStwAWKhA0FEgq8Jd4h5lpwo= +github.com/go-redsync/redsync/v4 v4.13.0 h1:49X6GJfnbLGaIpBBREM/zA4uIMDXKAh1NDkvQ1EkZKA= +github.com/go-redsync/redsync/v4 v4.13.0/go.mod h1:HMW4Q224GZQz6x1Xc7040Yfgacukdzu7ifTDAKiyErQ= +github.com/gofrs/flock v0.8.1 h1:+gYjHKf32LDeiEEFhQaotPbLuUXjY5ZqxKgXy7n59aw= +github.com/gofrs/flock v0.8.1/go.mod h1:F1TvTiK9OcQqauNUHlbJvyl9Qa1QvF/gOUDKA14jxHU= +github.com/gofrs/flock v0.12.1 h1:MTLVXXHf8ekldpJk3AKicLij9MdwOWkZ+a/jHHZby9E= +github.com/gofrs/flock v0.12.1/go.mod h1:9zxTsyu5xtJ9DK+1tFZyibEV7y3uwDxPPfbxeeHCoD0= +github.com/golang/protobuf v1.5.4 h1:i7eJL8qZTpSEXOPTxNKhASYpMn+8e5Q6AdndVa1dWek= +github.com/golang/protobuf v1.5.4/go.mod h1:lnTiLA8Wa4RWRcIUkrtSVa5nRhsEGBg48fD6rSs7xps= +github.com/gomodule/redigo v1.8.9 h1:Sl3u+2BI/kk+VEatbj0scLdrFhjPmbxOc1myhDP41ws= +github.com/gomodule/redigo v1.8.9/go.mod h1:7ArFNvsTjH8GMMzB4uy1snslv2BwmginuMs06a1uzZE= +github.com/google/go-cmp v0.2.0/go.mod h1:oXzfMopK8JAjlY9xF4vHSVASa0yLyX7SntLO5aqRK0M= +github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8= +github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU= +github.com/google/subcommands v1.2.0/go.mod h1:ZjhPrFU+Olkh9WazFPsl27BQ4UPiG37m3yTrtFlrHVk= +github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= +github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/google/wire v0.6.0 h1:HBkoIh4BdSxoyo9PveV8giw7ZsaBOvzWKfcg/6MrVwI= +github.com/google/wire v0.6.0/go.mod h1:F4QhpQ9EDIdJ1Mbop/NZBRB+5yrR6qg3BnctaoUk6NA= +github.com/hanwen/go-fuse/v2 v2.8.0 h1:wV8rG7rmCz8XHSOwBZhG5YcVqcYjkzivjmbaMafPlAs= +github.com/hanwen/go-fuse/v2 v2.8.0/go.mod h1:yE6D2PqWwm3CbYRxFXV9xUd8Md5d6NG0WBs5spCswmI= +github.com/hashicorp/errwrap v1.0.0/go.mod h1:YH+1FKiLXxHSkmPseP+kNlulaMuP3n2brvKWEqk/Jc4= +github.com/hashicorp/errwrap v1.1.0 h1:OxrOeh75EUXMY8TBjag2fzXGZ40LB6IKw45YeGUDY2I= +github.com/hashicorp/errwrap v1.1.0/go.mod h1:YH+1FKiLXxHSkmPseP+kNlulaMuP3n2brvKWEqk/Jc4= +github.com/hashicorp/go-multierror v1.1.1 h1:H5DkEtf6CXdFp0N0Em5UCwQpXMWke8IA0+lD48awMYo= +github.com/hashicorp/go-multierror v1.1.1/go.mod h1:iw975J/qwKPdAO1clOe2L8331t/9/fmwbPZ6JB6eMoM= +github.com/klauspost/compress v1.18.0 h1:c/Cqfb0r+Yi+JtIEq73FWXVkRonBlf0CRNYc8Zttxdo= +github.com/klauspost/compress v1.18.0/go.mod h1:2Pp+KzxcywXVXMr50+X0Q/Lsb43OQHYWRCY2AiWywWQ= +github.com/klauspost/cpuid/v2 v2.0.9/go.mod h1:FInQzS24/EEf25PyTYn52gqo7WaD8xa0213Md/qVLRg= +github.com/klauspost/cpuid/v2 v2.2.10 h1:tBs3QSyvjDyFTq3uoc/9xFpCuOsJQFNPiAhYdw2skhE= +github.com/klauspost/cpuid/v2 v2.2.10/go.mod h1:hqwkgyIinND0mEev00jJYCxPNVRVXFQeu1XKlok6oO0= +github.com/knz/go-libedit v1.10.1/go.mod h1:MZTVkCWyz0oBc7JOWP3wNAzd002ZbM/5hgShxwh4x8M= +github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE= +github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk= +github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= +github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= +github.com/kylelemons/godebug v1.1.0 h1:RPNrshWIDI6G2gRW9EHilWtl7Z6Sb1BR0xunSBf0SNc= +github.com/kylelemons/godebug v1.1.0/go.mod h1:9/0rRGxNHcop5bhtWyNeEfOS8JIWk580+fNqagV/RAw= +github.com/labstack/echo/v4 v4.13.3 h1:pwhpCPrTl5qry5HRdM5FwdXnhXSLSY+WE+YQSeCaafY= +github.com/labstack/echo/v4 v4.13.3/go.mod h1:o90YNEeQWjDozo584l7AwhJMHN0bOC4tAfg+Xox9q5g= +github.com/labstack/gommon v0.4.2 h1:F8qTUNXgG1+6WQmqoUWnz8WiEU60mXVVw0P4ht1WRA0= +github.com/labstack/gommon v0.4.2/go.mod h1:QlUFxVM+SNXhDL/Z7YhocGIBYOiwB0mXm1+1bAPHPyU= +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/mattn/go-colorable v0.1.14 h1:9A9LHSqF/7dyVVX6g0U9cwm9pG3kP9gSzcuIPHPsaIE= +github.com/mattn/go-colorable v0.1.14/go.mod h1:6LmQG8QLFO4G5z1gPvYEzlUgJ2wF+stgPZH1UqBm1s8= +github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY= +github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y= +github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 h1:C3w9PqII01/Oq1c1nUAm88MOHcQC9l5mIlSMApZMrHA= +github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822/go.mod h1:+n7T8mK8HuQTcFwEeznm/DIxMOiR9yIdICNftLE1DvQ= +github.com/patrickmn/go-cache v2.1.0+incompatible h1:HRMgzkcYKYpi3C8ajMPV8OFXaaRUnok+kx1WdO15EQc= +github.com/patrickmn/go-cache v2.1.0+incompatible/go.mod h1:3Qf8kWWT7OJRJbdiICTKqZju1ZixQ/KpMGzzAfe6+WQ= +github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4= +github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= +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/prometheus/client_golang v1.22.0 h1:rb93p9lokFEsctTys46VnV1kLCDpVZ0a/Y92Vm0Zc6Q= +github.com/prometheus/client_golang v1.22.0/go.mod h1:R7ljNsLXhuQXYZYtw6GAE9AZg8Y7vEW5scdCXrWRXC0= +github.com/prometheus/client_model v0.6.1 h1:ZKSh/rekM+n3CeS952MLRAdFwIKqeY8b62p8ais2e9E= +github.com/prometheus/client_model v0.6.1/go.mod h1:OrxVMOVHjw3lKMa8+x6HeMGkHMQyHDk9E3jmP2AmGiY= +github.com/prometheus/common v0.62.0 h1:xasJaQlnWAeyHdUBeGjXmutelfJHWMRr+Fg4QszZ2Io= +github.com/prometheus/common v0.62.0/go.mod h1:vyBcEuLSvWos9B1+CyL7JZ2up+uFzXhkqml0W5zIY1I= +github.com/prometheus/procfs v0.15.1 h1:YagwOFzUgYfKKHX6Dr+sHT7km/hxC76UB0learggepc= +github.com/prometheus/procfs v0.15.1/go.mod h1:fB45yRUv8NstnjriLhBQLuOUt+WW4BsoGhij/e3PBqk= +github.com/redis/go-redis/v9 v9.5.1 h1:H1X4D3yHPaYrkL5X06Wh6xNVM/pX0Ft4RV0vMGvLBh8= +github.com/redis/go-redis/v9 v9.5.1/go.mod h1:hdY0cQFCN4fnSYT6TkisLufl/4W5UIXyv0b/CLO2V2M= +github.com/redis/rueidis v1.0.19 h1:s65oWtotzlIFN8eMPhyYwxlwLR1lUdhza2KtWprKYSo= +github.com/redis/rueidis v1.0.19/go.mod h1:8B+r5wdnjwK3lTFml5VtxjzGOQAC+5UmujoD12pDrEo= +github.com/robfig/cron/v3 v3.0.1 h1:WdRxkvbJztn8LMz/QEvLN5sBU+xKpSqwwUO1Pjr4qDs= +github.com/robfig/cron/v3 v3.0.1/go.mod h1:eQICP3HwyT7UooqI/z+Ov+PtYAWygg1TEWWzGIFLtro= +github.com/rogpeppe/go-internal v1.10.0 h1:TMyTOH3F/DB16zRVcYyreMH6GnZZrwQVAoYjRBZyWFQ= +github.com/rogpeppe/go-internal v1.10.0/go.mod h1:UQnix2H7Ngw/k4C5ijL5+65zddjncjaFoBhdsK/akog= +github.com/shirou/gopsutil v3.21.11+incompatible h1:+1+c1VGhc88SSonWP6foOcLhvnKlUeu/erjjvaPEYiI= +github.com/shirou/gopsutil v3.21.11+incompatible/go.mod h1:5b4v6he4MtMOwMlS0TUMTu2PcXUg8+E1lC7eC3UO/RA= +github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= +github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw= +github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo= +github.com/stretchr/objx v0.5.2/go.mod h1:FRsXN1f5AsAjCGJKqEizvkpNtU+EGNCLh3NxZ/8L+MA= +github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= +github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= +github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU= +github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4= +github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo= +github.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOfJA= +github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY= +github.com/stvp/tempredis v0.0.0-20181119212430-b82af8480203 h1:QVqDTf3h2WHt08YuiTGPZLls0Wq99X9bWd0Q5ZSBesM= +github.com/stvp/tempredis v0.0.0-20181119212430-b82af8480203/go.mod h1:oqN97ltKNihBbwlX8dLpwxCl3+HnXKV/R0e+sRLd9C8= +github.com/twitchyliquid64/golang-asm v0.15.1 h1:SU5vSMR7hnwNxj24w34ZyCi/FmDZTkS4MhqMhdFk5YI= +github.com/twitchyliquid64/golang-asm v0.15.1/go.mod h1:a1lVb/DtPvCB8fslRZhAngC2+aY1QWCk3Cedj/Gdt08= +github.com/valyala/bytebufferpool v1.0.0 h1:GqA5TC/0021Y/b9FG4Oi9Mr3q7XYx6KllzawFIhcdPw= +github.com/valyala/bytebufferpool v1.0.0/go.mod h1:6bBcMArwyJ5K/AmCkWv1jt77kVWyCJ6HpOuEn7z0Csc= +github.com/valyala/fasttemplate v1.2.2 h1:lxLXG0uE3Qnshl9QyaK6XJxMXlQZELvChBOCmQD0Loo= +github.com/valyala/fasttemplate v1.2.2/go.mod h1:KHLXt3tVN2HBp8eijSv/kGJopbvo7S+qRAEEKiv+SiQ= +github.com/xyproto/randomstring v1.0.5 h1:YtlWPoRdgMu3NZtP45drfy1GKoojuR7hmRcnhZqKjWU= +github.com/xyproto/randomstring v1.0.5/go.mod h1:rgmS5DeNXLivK7YprL0pY+lTuhNQW3iGxZ18UQApw/E= +github.com/yuin/goldmark v1.4.13/go.mod h1:6yULJ656Px+3vBD8DxQVa3kxgyrAnzto9xy5taEt/CY= +github.com/yusufpapurcu/wmi v1.2.4 h1:zFUKzehAFReQwLys1b/iSMl+JQGSCSjtVqQn9bBrPo0= +github.com/yusufpapurcu/wmi v1.2.4/go.mod h1:SBZ9tNy3G9/m5Oi98Zks0QjeHVDvuK0qfxQmPyzfmi0= +go.opentelemetry.io/auto/sdk v1.1.0 h1:cH53jehLUN6UFLY71z+NDOiNJqDdPRaXzTel0sJySYA= +go.opentelemetry.io/auto/sdk v1.1.0/go.mod h1:3wSPjt5PWp2RhlCcmmOial7AvC4DQqZb7a7wCow3W8A= +go.opentelemetry.io/otel v1.35.0 h1:xKWKPxrxB6OtMCbmMY021CqC45J+3Onta9MqjhnusiQ= +go.opentelemetry.io/otel v1.35.0/go.mod h1:UEqy8Zp11hpkUrL73gSlELM0DupHoiq72dR+Zqel/+Y= +go.opentelemetry.io/otel/metric v1.35.0 h1:0znxYu2SNyuMSQT4Y9WDWej0VpcsxkuklLa4/siN90M= +go.opentelemetry.io/otel/metric v1.35.0/go.mod h1:nKVFgxBZ2fReX6IlyW28MgZojkoAkJGaE8CpgeAU3oE= +go.opentelemetry.io/otel/sdk v1.35.0 h1:iPctf8iprVySXSKJffSS79eOjl9pvxV9ZqOWT0QejKY= +go.opentelemetry.io/otel/sdk v1.35.0/go.mod h1:+ga1bZliga3DxJ3CQGg3updiaAJoNECOgJREo9KHGQg= +go.opentelemetry.io/otel/sdk/metric v1.35.0 h1:1RriWBmCKgkeHEhM7a2uMjMUfP7MsOF5JpUCaEqEI9o= +go.opentelemetry.io/otel/sdk/metric v1.35.0/go.mod h1:is6XYCUMpcKi+ZsOvfluY5YstFnhW0BidkR+gL+qN+w= +go.opentelemetry.io/otel/trace v1.35.0 h1:dPpEfJu1sDIqruz7BHFG3c7528f6ddfSWfFDVt/xgMs= +go.opentelemetry.io/otel/trace v1.35.0/go.mod h1:WUk7DtFp1Aw2MkvqGdwiXYDZZNvA/1J8o6xRXLrIkyc= +go.uber.org/atomic v1.10.0 h1:9qC72Qh0+3MqyJbAn8YU5xVq1frD8bn3JtD2oXtafVQ= +go.uber.org/atomic v1.10.0/go.mod h1:LUxbIzbOniOlMKjJjyPfpl4v+PKK2cNJn91OQbhoJI0= +go.uber.org/goleak v1.1.11 h1:wy28qYRKZgnJTxGxvye5/wgWr1EKjmUDGYox5mGlRlI= +go.uber.org/goleak v1.1.11/go.mod h1:cwTWslyiVhfpKIDGSZEM2HlOvcqm+tG4zioyIeLoqMQ= +go.uber.org/multierr v1.9.0 h1:7fIwc/ZtS0q++VgcfqFDxSBZVv/Xo49/SYnDFupUwlI= +go.uber.org/multierr v1.9.0/go.mod h1:X2jQV1h+kxSjClGpnseKVIxpmcjrj7MNnI0bnlfKTVQ= +go.uber.org/zap v1.24.0 h1:FiJd5l1UOLj0wCgbSE0rwwXHzEdAZS6hiiSnxJN/D60= +go.uber.org/zap v1.24.0/go.mod h1:2kMP+WWQ8aoFoedH3T2sq6iJ2yDWpHbP0f6MQbS9Gkg= +golang.org/x/arch v0.15.0 h1:QtOrQd0bTUnhNVNndMpLHNWrDmYzZ2KDqSrEymqInZw= +golang.org/x/arch v0.15.0/go.mod h1:JmwW7aLIoRUKgaTzhkiEFxvcEiQGyOg9BMonBJUS7EE= +golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w= +golang.org/x/crypto v0.0.0-20210921155107-089bfa567519/go.mod h1:GvvjBRRGRdwPK5ydBHafDWAxML/pGHZbMvKqRZ5+Abc= +golang.org/x/crypto v0.13.0/go.mod h1:y6Z2r+Rw4iayiXXAIxJIDAJ1zMW4yaTpebo8fPOliYc= +golang.org/x/crypto v0.18.0/go.mod h1:R0j02AL6hcrfOiy9T4ZYp/rcWeMxM3L6QYxlOuEG1mg= +golang.org/x/crypto v0.37.0 h1:kJNSjF/Xp7kU0iB2Z+9viTPMW4EqqsrywMXLJOOsXSE= +golang.org/x/crypto v0.37.0/go.mod h1:vg+k43peMZ0pUMhYmVAWysMK35e6ioLh3wB8ZCAfbVc= +golang.org/x/mod v0.6.0-dev.0.20220419223038-86c51ed26bb4/go.mod h1:jJ57K6gSWd91VN4djpZkiMVwK6gcyfeH4XE8wZrZaV4= +golang.org/x/mod v0.8.0/go.mod h1:iBbtSCu2XBx23ZKBPSOrRkjjQPZFPuis4dIYUhu/chs= +golang.org/x/mod v0.12.0/go.mod h1:iBbtSCu2XBx23ZKBPSOrRkjjQPZFPuis4dIYUhu/chs= +golang.org/x/mod v0.14.0/go.mod h1:hTbmBsO62+eylJbnUtE2MGJUyE7QWk4xUqPFrRgJ+7c= +golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s= +golang.org/x/net v0.0.0-20210226172049-e18ecbb05110/go.mod h1:m0MpNAwzfU5UDzcl9v0D8zg8gWTRqZa9RBIspLL5mdg= +golang.org/x/net v0.0.0-20220722155237-a158d28d115b/go.mod h1:XRhObCWvk6IyKnWLug+ECip1KBveYUHfp+8e9klMJ9c= +golang.org/x/net v0.6.0/go.mod h1:2Tu9+aMcznHK/AK1HMvgo6xiTLG5rD5rZLDS+rp2Bjs= +golang.org/x/net v0.10.0/go.mod h1:0qNGK6F8kojg2nk9dLZ2mShWaEBan6FAoqfSigmmuDg= +golang.org/x/net v0.15.0/go.mod h1:idbUs1IY1+zTqbi8yxTbhexhEEk5ur9LInksu6HrEpk= +golang.org/x/net v0.20.0/go.mod h1:z8BVo6PvndSri0LbOE3hAn0apkU+1YvI6E70E9jsnvY= +golang.org/x/net v0.39.0 h1:ZCu7HMWDxpXpaiKdhzIfaltL9Lp31x/3fCP11bc6/fY= +golang.org/x/net v0.39.0/go.mod h1:X7NRbYVEA+ewNkCNyJ513WmMdQ3BineSwVtN2zD/d+E= +golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= +golang.org/x/sync v0.0.0-20220722155255-886fb9371eb4/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= +golang.org/x/sync v0.1.0/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= +golang.org/x/sync v0.3.0/go.mod h1:FU7BRWz2tNW+3quACPkgCx/L+uEAv1htQ0V83Z9Rj+Y= +golang.org/x/sync v0.6.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk= +golang.org/x/sync v0.13.0 h1:AauUjRAJ9OSnvULf/ARrrVywoJDy0YS2AwQ98I37610= +golang.org/x/sync v0.13.0/go.mod h1:1dzgHSNfp02xaA81J2MS99Qcpr2w7fw1gpm99rleRqA= +golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= +golang.org/x/sys v0.0.0-20190916202348-b4ddaad3f8a3/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/sys v0.0.0-20201119102817-f84b799fce68/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= +golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.0.0-20220520151302-bc2c85ada10a/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.0.0-20220722155257-8c9f86f7a55f/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.5.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.8.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.12.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.16.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= +golang.org/x/sys v0.35.0 h1:vz1N37gP5bs89s7He8XuIYXpyY0+QlsKmzipCbUtyxI= +golang.org/x/sys v0.35.0/go.mod h1:BJP2sWEmIv4KK5OTEluFJCKSidICx8ciO85XgH3Ak8k= +golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo= +golang.org/x/term v0.0.0-20210927222741-03fcf44c2211/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8= +golang.org/x/term v0.5.0/go.mod h1:jMB1sMXY+tzblOD4FWmEbocvup2/aLOaQEp7JmGp78k= +golang.org/x/term v0.8.0/go.mod h1:xPskH00ivmX89bAKVGSKKtLOWNx2+17Eiy94tnKShWo= +golang.org/x/term v0.12.0/go.mod h1:owVbMEjm3cBLCHdkQu9b1opXd4ETQWc3BhuQGKgXgvU= +golang.org/x/term v0.16.0/go.mod h1:yn7UURbUtPyrVJPGPq404EukNFxcm/foM+bV/bfcDsY= +golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ= +golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ= +golang.org/x/text v0.3.7/go.mod h1:u+2+/6zg+i71rQMx5EYifcz6MCKuco9NR6JIITiCfzQ= +golang.org/x/text v0.7.0/go.mod h1:mrYo+phRRbMaCq/xk9113O4dZlRixOauAjOtrjsXDZ8= +golang.org/x/text v0.9.0/go.mod h1:e1OnstbJyHTd6l/uOt8jFFHp6TRDWZR/bV3emEE/zU8= +golang.org/x/text v0.13.0/go.mod h1:TvPlkZtksWOMsz7fbANvkp4WM8x/WCo/om8BMLbz+aE= +golang.org/x/text v0.14.0/go.mod h1:18ZOQIKpY8NJVqYksKHtTdi31H5itFRjB5/qKTNYzSU= +golang.org/x/text v0.24.0 h1:dd5Bzh4yt5KYA8f9CJHCP4FB4D51c2c6JvN37xJJkJ0= +golang.org/x/text v0.24.0/go.mod h1:L8rBsPeo2pSS+xqN0d5u2ikmjtmoJbDBT1b7nHvFCdU= +golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= +golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo= +golang.org/x/tools v0.1.12/go.mod h1:hNGJHUnrk76NpqgfD5Aqm5Crs+Hm0VOH/i9J2+nxYbc= +golang.org/x/tools v0.6.0/go.mod h1:Xwgl3UAJ/d3gWutnCtw505GrjyAbvKui8lOU390QaIU= +golang.org/x/tools v0.13.0/go.mod h1:HvlwmtVNQAhOuCjW7xxvovg8wbNq7LwfXh/k7wXUl58= +golang.org/x/tools v0.17.0/go.mod h1:xsh6VxdV005rRVaS6SSAf9oiAqljS7UZUacMZ8Bnsps= +golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= +google.golang.org/genproto/googleapis/rpc v0.0.0-20250324211829-b45e905df463 h1:e0AIkUUhxyBKh6ssZNrAMeqhA7RKUj42346d1y02i2g= +google.golang.org/genproto/googleapis/rpc v0.0.0-20250324211829-b45e905df463/go.mod h1:qQ0YXyHHx3XkvlzUtpXDkS29lDSafHMZBAZDc03LQ3A= +google.golang.org/grpc v1.73.0 h1:VIWSmpI2MegBtTuFt5/JWy2oXxtjJ/e89Z70ImfD2ok= +google.golang.org/grpc v1.73.0/go.mod h1:50sbHOUqWoCQGI8V2HQLJM0B+LMlIUjNSZmow7EVBQc= +google.golang.org/protobuf v1.36.6 h1:z1NpPI8ku2WgiWnf+t9wTPsn6eP1L7ksHUlkfLvd9xY= +google.golang.org/protobuf v1.36.6/go.mod h1:jduwjTPXsFjZGTmRluh+L6NjiWu7pchiJ2/5YcXBHnY= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk= +gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q= +gopkg.in/natefinch/lumberjack.v2 v2.2.1 h1:bBRl1b0OH9s/DuPhuXpNl+VtCaJXFZ5/uEFST95x9zc= +gopkg.in/natefinch/lumberjack.v2 v2.2.1/go.mod h1:YD8tP3GAjkrDg1eZH7EGmyESg/lsYskCTPBJVb9jqSc= +gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= +gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= +gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= +nullprogram.com/x/optparse v1.0.0/go.mod h1:KdyPE+Igbe0jQUrVfMqDMeJQIJZEuyV7pjYmp6pbG50= diff --git a/deploy/cachemount-sidecar/source/pkg/cachemount/mount_linux.go b/deploy/cachemount-sidecar/source/pkg/cachemount/mount_linux.go new file mode 100644 index 0000000..6ae69d0 --- /dev/null +++ b/deploy/cachemount-sidecar/source/pkg/cachemount/mount_linux.go @@ -0,0 +1,112 @@ +package cachemount + +import ( + "context" + "fmt" + "io" + "os" + "strings" + "syscall" + + "github.com/hanwen/go-fuse/v2/fs" + "github.com/hanwen/go-fuse/v2/fuse" +) + +type fileNode struct { + fs.Inode + payload *Payload +} + +func (n *fileNode) Getattr(ctx context.Context, f fs.FileHandle, out *fuse.AttrOut) syscall.Errno { + out.Mode = syscall.S_IFREG | 0444 + out.Size = uint64(n.payload.Size) + return 0 +} +func (n *fileNode) Open(ctx context.Context, flags uint32) (fs.FileHandle, uint32, syscall.Errno) { + if flags&uint32(syscall.O_ACCMODE|syscall.O_TRUNC|syscall.O_APPEND) != 0 { + return nil, 0, syscall.EROFS + } + return nil, fuse.FOPEN_KEEP_CACHE, 0 +} +func (n *fileNode) Read(ctx context.Context, f fs.FileHandle, dest []byte, off int64) (fuse.ReadResult, syscall.Errno) { + count, err := n.payload.ReadAt(dest, off) + // EOF within the declared range is truncation, never a successful short read. + want := int64(len(dest)) + if off >= n.payload.Size { + want = 0 + } else if want > n.payload.Size-off { + want = n.payload.Size - off + } + if off < 0 || (err != nil && err != io.EOF) || int64(count) != want { + return nil, syscall.EIO + } + return fuse.ReadResultData(dest[:count]), 0 +} + +type rootNode struct { + fs.Inode + snapshot *Snapshot +} + +type unavailableDir struct{ fs.Inode } + +func (n *unavailableDir) Readdir(ctx context.Context) (fs.DirStream, syscall.Errno) { + return nil, syscall.EIO +} +func (n *unavailableDir) Lookup(ctx context.Context, name string, out *fuse.EntryOut) (*fs.Inode, syscall.Errno) { + return nil, syscall.EIO +} + +func ensureDir(ctx context.Context, root *fs.Inode, parts []string) *fs.Inode { + dir := root + for _, part := range parts { + next := dir.GetChild(part) + if next == nil { + next = dir.NewPersistentInode(ctx, &fs.Inode{}, fs.StableAttr{Mode: syscall.S_IFDIR}) + dir.AddChild(part, next, false) + } + dir = next + } + return dir +} + +func (r *rootNode) OnAdd(ctx context.Context) { + for name, reason := range r.snapshot.Directories { + parts := strings.Split(name, "/") + parent := ensureDir(ctx, &r.Inode, parts[:len(parts)-1]) + if reason != "" { + parent.AddChild(parts[len(parts)-1], parent.NewPersistentInode(ctx, &unavailableDir{}, fs.StableAttr{Mode: syscall.S_IFDIR}), false) + } else { + ensureDir(ctx, parent, parts[len(parts)-1:]) + } + } + for name, p := range r.snapshot.Files { + parts := strings.Split(name, "/") + dir := &r.Inode + for _, part := range parts[:len(parts)-1] { + next := dir.GetChild(part) + if next == nil { + next = dir.NewPersistentInode(ctx, &fs.Inode{}, fs.StableAttr{Mode: syscall.S_IFDIR}) + dir.AddChild(part, next, false) + } + dir = next + } + n := dir.NewPersistentInode(ctx, &fileNode{payload: p}, fs.StableAttr{Mode: syscall.S_IFREG}) + dir.AddChild(parts[len(parts)-1], n, false) + } +} + +// Mount requires an existing empty directory, never hides existing user files. +// The caller owns Snapshot until the returned server has finished serving. +func Mount(s *Snapshot, mountpoint string) (*fuse.Server, error) { + entries, err := os.ReadDir(mountpoint) + if err != nil { + return nil, err + } + if len(entries) != 0 { + return nil, fmt.Errorf("mountpoint must be empty") + } + return fs.Mount(mountpoint, &rootNode{snapshot: s}, &fs.Options{ + MountOptions: fuse.MountOptions{Options: []string{"ro", "default_permissions"}, FsName: "dingocache", Name: "dingocache"}, + }) +} diff --git a/deploy/cachemount-sidecar/source/pkg/cachemount/payload.go b/deploy/cachemount-sidecar/source/pkg/cachemount/payload.go new file mode 100644 index 0000000..eda6fef --- /dev/null +++ b/deploy/cachemount-sidecar/source/pkg/cachemount/payload.go @@ -0,0 +1,139 @@ +// Package cachemount exposes complete DingCache payloads without copying or +// modifying cache files. All backing descriptors remain open until unmount. +package cachemount + +import ( + "encoding/binary" + "fmt" + "io" + "math" + "os" + "path" + "strings" +) + +type Payload struct { + File *os.File + Offset int64 + Size int64 +} + +// Open accepts OLAH v8 only. A damaged container never falls back to raw data. +func Open(name string) (*Payload, error) { + f, err := os.Open(name) + if err != nil { + return nil, err + } + p, err := inspect(f) + if err != nil { + f.Close() + return nil, fmt.Errorf("%s: %w", name, err) + } + return p, nil +} + +func inspect(f *os.File) (*Payload, error) { + st, err := f.Stat() + if err != nil { + return nil, err + } + if !st.Mode().IsRegular() { + return nil, fmt.Errorf("not a regular cache file") + } + var h [36]byte + if _, err := f.ReadAt(h[:], 0); err != nil { + return nil, err + } + v := binary.LittleEndian.Uint64(h[4:12]) + block := binary.LittleEndian.Uint64(h[12:20]) + size := binary.LittleEndian.Uint64(h[20:28]) + bits := binary.LittleEndian.Uint64(h[28:36]) + if string(h[:4]) != "OLAH" || v != 8 || block == 0 || block > math.MaxInt64 || size > math.MaxInt64 || bits == 0 || bits > 8*1024*1024 { + return nil, fmt.Errorf("invalid OLAH v8 header") + } + count := size / block + if size%block != 0 { + count++ + } + if count > bits { + return nil, fmt.Errorf("bitmap too small") + } + mask := make([]byte, (bits+7)/8) + if _, err := f.ReadAt(mask, 36); err != nil { + return nil, err + } + offset := int64(36 + len(mask)) + if st.Size()-offset < int64(size) { + return nil, fmt.Errorf("truncated cache payload") + } + for i := uint64(0); i < count; i++ { + if mask[i/8]&(1<<(i%8)) == 0 { + return nil, fmt.Errorf("missing cache block %d", i) + } + } + return &Payload{File: f, Offset: offset, Size: int64(size)}, nil +} + +func (p *Payload) ReadAt(b []byte, off int64) (int, error) { + if off < 0 { + return 0, fmt.Errorf("negative offset") + } + if len(b) == 0 { + return 0, nil + } + if off >= p.Size { + return 0, io.EOF + } + want := len(b) + if int64(len(b)) > p.Size-off { + b = b[:p.Size-off] + } + n, err := p.File.ReadAt(b, p.Offset+off) + if err == nil && n < want { + err = io.EOF + } + return n, err +} + +type Snapshot struct { + Files map[string]*Payload + Directories map[string]string +} + +// Pin validates the whole tree before exposing any file. It reads headers and +// bitmaps only, not the weights. Failure closes all previously opened files. +func Pin(sources map[string]string) (*Snapshot, error) { + s := &Snapshot{Files: make(map[string]*Payload)} + for name, source := range sources { + if name == "." || !validPath(name) { + s.Close() + return nil, fmt.Errorf("invalid model path %q", name) + } + p, err := Open(source) + if err != nil { + s.Close() + return nil, fmt.Errorf("%s: %w", name, err) + } + s.Files[name] = p + } + for name := range s.Files { + for parent := path.Dir(name); parent != "."; parent = path.Dir(parent) { + if _, ok := s.Files[parent]; ok { + s.Close() + return nil, fmt.Errorf("file/directory collision: %s", parent) + } + } + } + if len(s.Files) == 0 { + return nil, fmt.Errorf("empty model manifest") + } + return s, nil +} +func validPath(p string) bool { + return p != "" && !strings.ContainsAny(p, "\\\x00") && !strings.HasPrefix(p, "/") && path.Clean(p) == p && p != ".." && !strings.HasPrefix(p, "../") +} +func (s *Snapshot) Close() { + for _, p := range s.Files { + p.File.Close() + } +} diff --git a/deploy/cachemount-sidecar/source/pkg/cachemount/tree.go b/deploy/cachemount-sidecar/source/pkg/cachemount/tree.go new file mode 100644 index 0000000..3fb068e --- /dev/null +++ b/deploy/cachemount-sidecar/source/pkg/cachemount/tree.go @@ -0,0 +1,350 @@ +package cachemount + +import ( + "encoding/hex" + "encoding/json" + "fmt" + "io" + "io/fs" + "os" + "path" + "path/filepath" + "sort" + "strings" + + "dingospeed/pkg/hfprojection" + "dingospeed/pkg/repository" +) + +// Selection is an administrator-selected local model version. This service is +// local filesystem access, not a substitute for the web console's user ACLs. +type Selection struct { + CacheRoot string `json:"cache_root"` + Namespace string `json:"namespace"` + Repo string `json:"repo"` + Revision string `json:"revision"` +} +type TreeFile struct { + Path string `json:"path"` + Size int64 `json:"size"` + ContentID string `json:"content_id"` + Backing string `json:"backing"` +} +type Version struct { + Selection + Commit string `json:"commit"` + SourceFilePath string `json:"source_file_path"` + Status string `json:"status"` + Error string `json:"error,omitempty"` + Files []TreeFile `json:"files"` +} + +func readJSON(root, p string, out any) error { + if err := repository.SafePath(root, p); err != nil { + return err + } + f, err := os.Open(p) + if err != nil { + return err + } + defer f.Close() + b, err := io.ReadAll(io.LimitReader(f, 32*1024*1024+1)) + if err != nil { + return err + } + if len(b) > 32*1024*1024 { + return fmt.Errorf("metadata exceeds 32 MiB") + } + return json.Unmarshal(b, out) +} +func commitAt(root, p string) (string, error) { + var envelope struct { + Status int `json:"status_code"` + Content string `json:"content"` + } + if err := readJSON(root, p, &envelope); err != nil { + return "", err + } + if envelope.Status != 200 { + return "", fmt.Errorf("metadata status %d", envelope.Status) + } + b, err := hex.DecodeString(envelope.Content) + if err != nil { + return "", err + } + var meta struct { + SHA string `json:"sha"` + } + if err = json.Unmarshal(b, &meta); err != nil { + return "", err + } + return meta.SHA, repository.Segment(meta.SHA) +} + +// Discover uses registered identities, legacy upload manifests and HF/ModelScope cached +// metadata. It never registers repositories or downloads content. +func Discover(root string) ([]Selection, error) { + root, err := filepath.Abs(root) + if err != nil { + return nil, err + } + if _, err = os.Stat(root); err != nil { + return nil, err + } + descriptors, err := repository.List(root) + if err != nil { + return nil, err + } + keys := map[string]repository.RepoKey{} + for _, d := range descriptors { + if d.RepoType == "models" && d.Source == "hosted" { + keys[d.ID()] = d.RepoKey + } + } + local := filepath.Join(root, "api", "models", repository.Local) + err = filepath.WalkDir(local, func(p string, e fs.DirEntry, walkErr error) error { + if walkErr != nil { + return walkErr + } + if e.IsDir() { + return nil + } + if e.Name() != "dingo-local-manifest.json" { + return nil + } + repoRoot := filepath.Dir(filepath.Dir(filepath.Dir(p))) + for _, k := range keys { + if filepath.Clean(k.APIRoot(root)) == repoRoot { + return nil + } + } + rel, err := filepath.Rel(local, repoRoot) + if err != nil { + return err + } + k := repository.RepoKey{Namespace: repository.Local, RepoType: "models", Repo: filepath.ToSlash(rel)} + if err = k.Validate(); err != nil { + return err + } + keys[k.ID()] = k + return nil + }) + if err != nil && !os.IsNotExist(err) { + return nil, err + } + out := []Selection{} + for _, k := range keys { + entries, err := os.ReadDir(filepath.Join(k.APIRoot(root), "revision")) + if err != nil { + return nil, err + } + revisions := map[string]string{} + pointed := map[string]bool{} + for _, e := range entries { + if !e.IsDir() { + continue + } + p := k.Revision(root, e.Name()) + if _, err := os.Stat(filepath.Join(p, "dingo-superseded.json")); err == nil { + continue + } else if !os.IsNotExist(err) { + return nil, err + } + commit, err := commitAt(root, filepath.Join(p, "meta_get.json")) + if os.IsNotExist(err) { + continue + } + if err != nil { + return nil, err + } + revisions[e.Name()] = commit + if e.Name() != commit { + pointed[commit] = true + } + } + for rev, c := range revisions { + if rev == c && pointed[c] { + continue + } + out = append(out, Selection{root, k.Namespace, k.Repo, rev}) + } + } + for _, namespace := range []string{repository.HuggingFace, repository.ModelScope} { + catalog, err := (hfprojection.Reader{Root: root, Namespace: namespace}).Catalog("models") + if err != nil { + return nil, err + } + for _, repo := range catalog.Repos { + if repo.Error != "" { + return nil, fmt.Errorf("%s/%s: %s", namespace, repo.Repo, repo.Error) + } + for _, rev := range repo.Revisions { + out = append(out, Selection{root, namespace, repo.Repo, rev.Name}) + } + } + } + sort.Slice(out, func(i, j int) bool { + return out[i].Namespace+"/"+out[i].Repo+"/"+out[i].Revision < out[j].Namespace+"/"+out[j].Repo+"/"+out[j].Revision + }) + return out, nil +} + +func resolve(s Selection) (Version, error) { + v := Version{Selection: s, Status: "unavailable", Files: []TreeFile{}} + k := repository.RepoKey{Namespace: s.Namespace, RepoType: "models", Repo: s.Repo} + if err := k.Validate(); err != nil { + return v, err + } + if err := repository.Segment(s.Revision); err != nil { + return v, err + } + if !filepath.IsAbs(s.CacheRoot) { + return v, fmt.Errorf("cache_root must be absolute") + } + if s.Namespace == repository.HuggingFace || s.Namespace == repository.ModelScope { + m, paths, err := (hfprojection.Reader{Root: s.CacheRoot, Namespace: s.Namespace}).MountSources("models", s.Repo, s.Revision) + if err != nil { + return v, err + } + v.Commit = m.Commit + for _, f := range m.Files { + v.Files = append(v.Files, TreeFile{f.Path, f.Size, f.OID, paths[f.Path]}) + } + } else { + if s.Namespace != repository.Local { + d, err := repository.Read(s.CacheRoot, k) + if err != nil { + return v, err + } + if d.Source != "hosted" { + return v, fmt.Errorf("not a hosted repository") + } + } + commit, err := commitAt(s.CacheRoot, filepath.Join(k.Revision(s.CacheRoot, s.Revision), "meta_get.json")) + if err != nil { + return v, err + } + v.Commit = commit + var manifest []struct { + Path string `json:"path"` + SHA string `json:"sha256"` + Size int64 `json:"size"` + } + if err = readJSON(s.CacheRoot, filepath.Join(k.Revision(s.CacheRoot, commit), "dingo-local-manifest.json"), &manifest); err != nil { + return v, err + } + if manifest == nil { + return v, fmt.Errorf("null manifest") + } + for _, f := range manifest { + digest, err := hex.DecodeString(f.SHA) + if err != nil || len(digest) != 32 || f.Size < 0 { + return v, fmt.Errorf("invalid manifest entry %s", f.Path) + } + p := k.Blob(s.CacheRoot, f.SHA) + if err = repository.SafePath(s.CacheRoot, p); err != nil { + return v, err + } + v.Files = append(v.Files, TreeFile{f.Path, f.Size, f.SHA, p}) + } + } + seen := map[string]bool{} + for _, f := range v.Files { + if err := repository.Relative(f.Path); err != nil { + return v, err + } + if seen[f.Path] { + return v, fmt.Errorf("duplicate path %s", f.Path) + } + seen[f.Path] = true + } + sort.Slice(v.Files, func(i, j int) bool { return v.Files[i].Path < v.Files[j].Path }) + return v, nil +} + +// Historical HF caches can contain ordinary files. Accept those only with a +// known matching size; malformed OLAH must never fall back to ordinary bytes. +func openTreePayload(name string, size int64, allowPlain bool) (*Payload, error) { + p, err := Open(name) + if err == nil || !allowPlain || size < 0 { + return p, err + } + f, openErr := os.Open(name) + if openErr != nil { + return nil, openErr + } + var magic [4]byte + _, _ = f.ReadAt(magic[:], 0) + st, statErr := f.Stat() + if statErr != nil || !st.Mode().IsRegular() || st.Size() != size || string(magic[:]) == "OLAH" { + f.Close() + return nil, err + } + return &Payload{File: f, Size: size}, nil +} + +// PrepareTree pins every ready version atomically. Unavailable versions remain +// visible as error directories, never as misleading partial/empty models. +func PrepareTree(selections []Selection, mountpoint string) (*Snapshot, []Version, error) { + snapshot := &Snapshot{Files: map[string]*Payload{}, Directories: map[string]string{}} + versions := []Version{} + prefixes := []string{} + for _, s := range selections { + k := repository.RepoKey{Namespace: s.Namespace, RepoType: "models", Repo: s.Repo} + if err := k.Validate(); err != nil { + snapshot.Close() + return nil, nil, err + } + if err := repository.Segment(s.Revision); err != nil { + snapshot.Close() + return nil, nil, err + } + prefix := path.Join(s.Namespace, s.Repo, s.Revision) + for _, old := range prefixes { + if prefix == old || strings.HasPrefix(prefix, old+"/") || strings.HasPrefix(old, prefix+"/") { + snapshot.Close() + return nil, nil, fmt.Errorf("ambiguous version paths %s and %s", old, prefix) + } + } + prefixes = append(prefixes, prefix) + v, err := resolve(s) + v.SourceFilePath = filepath.Join(mountpoint, filepath.FromSlash(prefix)) + opened := map[string]*Payload{} + if err == nil { + for _, f := range v.Files { + var p *Payload + p, err = openTreePayload(f.Backing, f.Size, s.Namespace == repository.HuggingFace) + if err != nil { + break + } + opened[f.Path] = p + if f.Size >= 0 && p.Size != f.Size { + err = fmt.Errorf("size mismatch: %s", f.Path) + break + } + } + for name := range opened { + for parent := path.Dir(name); parent != "."; parent = path.Dir(parent) { + if _, ok := opened[parent]; ok { + err = fmt.Errorf("file/directory collision: %s", parent) + } + } + } + } + if err != nil { + for _, p := range opened { + p.File.Close() + } + v.Error = err.Error() + snapshot.Directories[prefix] = v.Error + } else { + v.Status = "ready" + snapshot.Directories[prefix] = "" + for name, p := range opened { + snapshot.Files[prefix+"/"+name] = p + } + } + versions = append(versions, v) + } + return snapshot, versions, nil +} diff --git a/deploy/cachemount-sidecar/source/pkg/hfprojection/index.go b/deploy/cachemount-sidecar/source/pkg/hfprojection/index.go new file mode 100644 index 0000000..e559270 --- /dev/null +++ b/deploy/cachemount-sidecar/source/pkg/hfprojection/index.go @@ -0,0 +1,124 @@ +package hfprojection + +import ( + "os" + "sync" +) + +// Index is the read-only projection used by repository catalogue queries. A +// refresh scans the existing HF cache once and atomically replaces one repo +// type; manifest reads then reuse that snapshot without walking the cache. +type Index struct { + mu sync.RWMutex + refreshMu sync.Mutex + scopes map[string]indexedScope +} + +type indexedScope struct { + catalog Catalog + manifests map[string]Manifest + errors map[string]error +} + +func NewIndex() *Index { return &Index{scopes: make(map[string]indexedScope)} } + +var DefaultIndex = NewIndex() + +func scopeKey(root, repoType string) string { return root + "\x00" + repoType } +func manifestKey(repo, revision string) string { return repo + "\x00" + revision } + +// Refresh rebuilds only from local cache files. It never calls HF and never +// writes the provider cache or its database. +func (i *Index) Refresh(root, repoType string) (Catalog, error) { + i.refreshMu.Lock() + defer i.refreshMu.Unlock() + return i.refresh(root, repoType) +} + +func (i *Index) refresh(root, repoType string) (Catalog, error) { + reader := Reader{Root: root} + catalog, err := reader.Catalog(repoType) + if err != nil { + return Catalog{}, err + } + scope := indexedScope{catalog: catalog, manifests: make(map[string]Manifest), errors: make(map[string]error)} + for _, repo := range catalog.Repos { + if repo.Error != "" { + continue + } + for _, revision := range repo.Revisions { + manifest, readErr := reader.Manifest(repoType, repo.Repo, revision.Name) + key := manifestKey(repo.Repo, revision.Name) + if readErr != nil { + scope.errors[key] = readErr + continue + } + scope.manifests[key] = manifest + commitKey := manifestKey(repo.Repo, revision.Commit) + if _, seen := scope.manifests[commitKey]; seen { + continue + } + commitManifest := manifest + commitManifest.Revision = revision.Commit + scope.manifests[commitKey] = commitManifest + } + } + i.mu.Lock() + i.scopes[scopeKey(root, repoType)] = scope + i.mu.Unlock() + return catalog, nil +} + +func (i *Index) Manifest(root, repoType, repo, revision string) (Manifest, error) { + key := scopeKey(root, repoType) + i.mu.RLock() + _, exists := i.scopes[key] + i.mu.RUnlock() + if !exists { + i.refreshMu.Lock() + i.mu.RLock() + _, exists = i.scopes[key] + i.mu.RUnlock() + if !exists { + _, err := i.refresh(root, repoType) + if err != nil { + i.refreshMu.Unlock() + return Manifest{}, err + } + } + i.refreshMu.Unlock() + } + i.mu.RLock() + scope := i.scopes[key] + manifest, ok := scope.manifests[manifestKey(repo, revision)] + readErr := scope.errors[manifestKey(repo, revision)] + i.mu.RUnlock() + if readErr != nil { + return Manifest{}, readErr + } + if !ok { + return Manifest{}, os.ErrNotExist + } + return manifest, nil +} + +// Invalidate makes the next catalogue or manifest read rebuild local facts. +// It is called after a cache task changes HF bytes. +func (i *Index) Invalidate(root, repoType string) { + if i == nil { + return + } + i.refreshMu.Lock() + defer i.refreshMu.Unlock() + i.mu.Lock() + defer i.mu.Unlock() + if repoType == "" { + for key := range i.scopes { + if len(key) > len(root) && key[:len(root)+1] == root+"\x00" { + delete(i.scopes, key) + } + } + return + } + delete(i.scopes, scopeKey(root, repoType)) +} diff --git a/deploy/cachemount-sidecar/source/pkg/hfprojection/mount_sources.go b/deploy/cachemount-sidecar/source/pkg/hfprojection/mount_sources.go new file mode 100644 index 0000000..472a0bb --- /dev/null +++ b/deploy/cachemount-sidecar/source/pkg/hfprojection/mount_sources.go @@ -0,0 +1,28 @@ +package hfprojection + +import "fmt" + +// MountSources resolves one observed manifest, without re-reading a mutable +// branch for each file. The caller must validate and pin every payload before +// publishing the mount. Paths are local administrative data, not an HTTP API. +func (r Reader) MountSources(typ, repo, revision string) (Manifest, map[string]string, error) { + m, err := r.Manifest(typ, repo, revision) + if err != nil { + return m, nil, err + } + if !m.ManifestComplete { + return m, nil, fmt.Errorf("complete cached repository manifest required") + } + paths := make(map[string]string, len(m.Files)) + for _, f := range m.Files { + if f.localPath == "" { + return m, nil, fmt.Errorf("missing content: %s", f.Path) + } + p, err := r.contentPath(m.RepoKey, f.localPath) + if err != nil { + return m, nil, fmt.Errorf("%s: %w", f.Path, err) + } + paths[f.Path] = p + } + return m, paths, nil +} diff --git a/deploy/cachemount-sidecar/source/pkg/hfprojection/projection.go b/deploy/cachemount-sidecar/source/pkg/hfprojection/projection.go new file mode 100644 index 0000000..eb373d5 --- /dev/null +++ b/deploy/cachemount-sidecar/source/pkg/hfprojection/projection.go @@ -0,0 +1,633 @@ +// Package hfprojection reads existing HF cache facts without registration, +// network access, database access, or opening any file for writing. +package hfprojection + +import ( + "encoding/binary" + "encoding/hex" + "encoding/json" + "errors" + "fmt" + "io" + "io/fs" + "math" + "os" + "path/filepath" + "sort" + "strings" + + "dingospeed/pkg/repository" +) + +type Revision struct { + Name string `json:"name"` + Commit string `json:"commit"` +} +type Repo struct { + repository.RepoKey + Revisions []Revision `json:"revisions"` + Access string `json:"access"` + Error string `json:"error,omitempty"` +} +type Catalog struct { + Namespace string `json:"namespace"` + RepoType string `json:"repoType"` + Repos []Repo `json:"repos"` +} +type File struct { + Path string `json:"path"` + OID string `json:"oid"` + Size int64 `json:"size"` + CachedBytes int64 `json:"cachedBytes"` + CacheStatus string `json:"cacheStatus"` + localPath string +} +type Manifest struct { + repository.RepoKey + Revision string `json:"revision"` + Commit string `json:"commit"` + MetadataAvailable bool `json:"metadataAvailable"` + ManifestComplete bool `json:"manifestComplete"` + Access string `json:"access"` + Files []File `json:"files"` +} + +// Reader has only a filesystem root. No downloader or database can be injected. +type Reader struct { + Root string + Namespace string +} + +func (r Reader) namespace() string { + if r.Namespace != "" { + return r.Namespace + } + return repository.HuggingFace +} + +func (r Reader) key(typ, repo string) (repository.RepoKey, error) { + k := repository.RepoKey{Namespace: r.namespace(), RepoType: typ, Repo: repo} + if err := k.Validate(); err != nil { + return k, err + } + p := strings.Split(repo, "/") + if len(p) > 2 || (r.namespace() == repository.HuggingFace && (strings.EqualFold(p[0], repository.Local) || strings.EqualFold(p[0], repository.ModelScope))) { + return k, fmt.Errorf("not an upstream HF repository") + } + return k, nil +} +func (r Reader) safe(p string) error { return repository.SafePath(r.Root, p) } +func (r Reader) dirs(p string) ([]os.DirEntry, error) { + if err := r.safe(p); err != nil { + return nil, err + } + v, err := os.ReadDir(p) + if errors.Is(err, os.ErrNotExist) { + return nil, nil + } + return v, err +} + +// Catalog discovers legacy layouts, including repositories without a marker. +// It never traverses hosted/modelscope roots or guesses identity from blobs. +func (r Reader) Catalog(typ string) (Catalog, error) { + out := Catalog{Namespace: r.namespace(), RepoType: typ, Repos: []Repo{}} + if _, err := r.key(typ, "placeholder"); err != nil { + return out, err + } + if err := r.safe(r.Root); err != nil { + return out, err + } + if _, err := os.Stat(r.Root); err != nil { + return out, err + } + names := map[string]bool{} + if r.namespace() == repository.ModelScope { + descriptors, err := repository.List(r.Root) + if err != nil { + return out, err + } + for _, d := range descriptors { + if d.Namespace == r.namespace() && d.RepoType == typ { + names[d.Repo] = true + } + } + } else { + for _, area := range []string{"api", "files"} { + base := filepath.Join(r.Root, area, typ) + owners, err := r.dirs(base) + if err != nil { + return out, err + } + for _, owner := range owners { + if !owner.IsDir() || strings.EqualFold(owner.Name(), repository.Local) || strings.EqualFold(owner.Name(), repository.ModelScope) { + continue + } + children, err := r.dirs(filepath.Join(base, owner.Name())) + if err != nil { + return out, err + } + nested := map[string]bool{} + for _, child := range children { + if !child.IsDir() { + continue + } + entries, err := r.dirs(filepath.Join(base, owner.Name(), child.Name())) + if err != nil { + return out, err + } + for _, e := range entries { + if e.IsDir() && cacheArea(e.Name(), area) { + names[owner.Name()+"/"+child.Name()] = true + nested[child.Name()] = true + } + } + } + // A real two-segment repository can itself be named revision, + // paths-info or resolve. Prefer its own layout over treating its + // name as the enclosing single-segment repository's cache area. + for _, child := range children { + if child.IsDir() && cacheArea(child.Name(), area) && !nested[child.Name()] { + names[owner.Name()] = true + } + } + } + } + } + for name := range names { + k, err := r.key(typ, name) + if err != nil { + return out, err + } + repo := Repo{RepoKey: k, Revisions: []Revision{}, Access: "unknown"} + repo.Revisions, err = r.revisions(k) + if err == nil { + repo.Access, err = r.access(k, repo.Revisions) + } + if err != nil { + repo.Error = "cached repository metadata is unreadable or invalid" + repo.Revisions = []Revision{} + } + out.Repos = append(out.Repos, repo) + } + sort.Slice(out.Repos, func(i, j int) bool { return out.Repos[i].Repo < out.Repos[j].Repo }) + return out, nil +} +func cacheArea(name, area string) bool { + return area == "api" && (name == "revision" || name == "paths-info") || area == "files" && name == "resolve" +} + +type metadata struct { + SHA string `json:"sha"` + Private *bool `json:"private"` + Gated json.RawMessage `json:"gated"` + Siblings []struct { + Path string `json:"rfilename"` + Size *int64 `json:"size"` + OID string `json:"blobId"` + LFS *struct { + OID string `json:"sha256"` + Size int64 `json:"size"` + } `json:"lfs"` + } `json:"siblings"` +} + +func metadataAccess(m metadata) string { + if m.Private != nil && *m.Private { + return "restricted" + } + gated := strings.TrimSpace(string(m.Gated)) + if gated != "" && gated != "false" && gated != "null" { + return "restricted" + } + if m.Private != nil && !*m.Private && gated == "false" { + return "public" + } + return "unknown" +} + +// No cached permission fact is inferred from successful download or a repo +// name. A missing permission field is unknown, not permission to publish it. +func (r Reader) access(k repository.RepoKey, revisions []Revision) (string, error) { + access := "public" + // Hidden fixed-commit metadata can still contain a restrictive permission + // fact. Deduplicating the displayed branch list must not drop that fact. + seen := map[string]bool{} + all := append([]Revision(nil), revisions...) + for _, rev := range all { + seen[rev.Name] = true + } + entries, err := r.dirs(filepath.Join(k.APIRoot(r.Root), "revision")) + if err != nil { + return "unknown", err + } + for _, entry := range entries { + if entry.IsDir() && !seen[entry.Name()] { + all = append(all, Revision{Name: entry.Name()}) + } + } + if len(all) == 0 { + return "unknown", nil + } + for _, rev := range all { + var m metadata + err := r.decode(filepath.Join(k.Revision(r.Root, rev.Name), "meta_get.json"), &m) + if errors.Is(err, os.ErrNotExist) { + access = "unknown" + continue + } + if err != nil { + return "unknown", err + } + switch metadataAccess(m) { + case "restricted": + return "restricted", nil + case "unknown": + access = "unknown" + } + } + return access, nil +} + +func (r Reader) decode(p string, out any) error { + if err := r.safe(p); err != nil { + return err + } + f, err := os.Open(p) + if err != nil { + return err + } + defer f.Close() + // Cache responses are bounded; malformed/truncated data must fail closed. + b, err := io.ReadAll(io.LimitReader(f, 32*1024*1024+1)) + if err != nil { + return err + } + if len(b) > 32*1024*1024 { + return fmt.Errorf("cached metadata exceeds projection limit") + } + var wrapper struct { + Status int `json:"status_code"` + Content string `json:"content"` + } + if err = json.Unmarshal(b, &wrapper); err != nil { + return fmt.Errorf("invalid cached metadata: %w", err) + } + if wrapper.Status != 200 { + return fmt.Errorf("cached metadata is not successful (HTTP %d)", wrapper.Status) + } + b, err = hex.DecodeString(wrapper.Content) + if err != nil { + return fmt.Errorf("invalid cached metadata encoding: %w", err) + } + if err = json.Unmarshal(b, out); err != nil { + return fmt.Errorf("invalid cached response body: %w", err) + } + return nil +} +func (r Reader) revisions(k repository.RepoKey) ([]Revision, error) { + byName := map[string]string{} + entries, err := r.dirs(filepath.Join(k.APIRoot(r.Root), "revision")) + if err != nil { + return nil, err + } + for _, e := range entries { + if !e.IsDir() { + continue + } + if err := repository.Segment(e.Name()); err != nil { + return nil, err + } + var meta metadata + err := r.decode(filepath.Join(k.Revision(r.Root, e.Name()), "meta_get.json"), &meta) + if errors.Is(err, os.ErrNotExist) { + continue + } + if err != nil { + return nil, err + } + if err := repository.Segment(meta.SHA); err != nil { + return nil, fmt.Errorf("invalid cached commit") + } + byName[e.Name()] = meta.SHA + } + if k.Namespace != repository.ModelScope { + for _, base := range []string{filepath.Join(k.APIRoot(r.Root), "paths-info"), filepath.Join(k.FilesRoot(r.Root), "resolve")} { + entries, err := r.dirs(base) + if err != nil { + return nil, err + } + for _, e := range entries { + if e.IsDir() { + if err := repository.Segment(e.Name()); err != nil { + return nil, err + } + if _, ok := byName[e.Name()]; !ok { + byName[e.Name()] = e.Name() + } + } + } + } + } + out := []Revision{} + pointed := map[string]bool{} + for name, commit := range byName { + if name != commit { + pointed[commit] = true + } + } + for name, commit := range byName { + if name == commit && pointed[commit] { + continue + } + out = append(out, Revision{name, commit}) + } + sort.Slice(out, func(i, j int) bool { return out[i].Name < out[j].Name }) + return out, nil +} + +func (r Reader) Manifest(typ, repo, revision string) (Manifest, error) { + k, err := r.key(typ, repo) + out := Manifest{RepoKey: k, Revision: revision, Files: []File{}, Access: "unknown"} + if err != nil { + return out, err + } + if err = repository.Segment(revision); err != nil { + return out, err + } + revisions, err := r.revisions(k) + if err != nil { + return out, err + } + for _, rev := range revisions { + if rev.Name == revision { + out.Commit = rev.Commit + break + } + } + // Fixed commits remain directly readable even when hidden from the branch + // list to avoid counting the same cached revision twice. + if out.Commit == "" { + for _, rev := range revisions { + if rev.Commit == revision { + out.Commit = revision + break + } + } + } + if out.Commit == "" { + return out, os.ErrNotExist + } + out.Access, err = r.access(k, revisions) + if err != nil { + return out, err + } + files := map[string]File{} + var meta metadata + err = r.decode(filepath.Join(k.Revision(r.Root, revision), "meta_get.json"), &meta) + if errors.Is(err, os.ErrNotExist) && revision != out.Commit { + err = r.decode(filepath.Join(k.Revision(r.Root, out.Commit), "meta_get.json"), &meta) + } + if err != nil && !errors.Is(err, os.ErrNotExist) { + return out, err + } + if err == nil { + if meta.SHA != out.Commit { + return out, fmt.Errorf("cached revision changed; refresh projection") + } + out.MetadataAvailable = true + out.ManifestComplete = meta.Siblings != nil + for _, s := range meta.Siblings { + if err := repository.Relative(s.Path); err != nil { + return out, err + } + v := File{Path: s.Path, OID: s.OID, Size: -1, CacheStatus: "unknown"} + if s.Size != nil { + v.Size = *s.Size + } + if s.LFS != nil { + v.OID = s.LFS.OID + v.Size = s.LFS.Size + } + files[s.Path] = v + } + } + infoRoot := filepath.Join(k.APIRoot(r.Root), "paths-info", out.Commit) + err = r.walk(infoRoot, func(p string, e fs.DirEntry) error { + if e.Name() != "paths-info_post.json" { + return nil + } + rel, _ := filepath.Rel(infoRoot, filepath.Dir(p)) + path := filepath.ToSlash(rel) + if err := repository.Relative(path); err != nil { + return err + } + var items []struct { + OID string `json:"oid"` + Size int64 `json:"size"` + LFS struct { + OID string `json:"oid"` + } `json:"lfs"` + } + if err := r.decode(p, &items); err != nil { + return err + } + if len(items) != 1 { + return fmt.Errorf("invalid cached file metadata") + } + v := files[path] + v.Path = path + v.Size = items[0].Size + v.OID = items[0].OID + if items[0].LFS.OID != "" { + v.OID = items[0].LFS.OID + } + files[path] = v + return nil + }) + if err != nil { + return out, err + } + resolveRoot := filepath.Join(k.FilesRoot(r.Root), "resolve", out.Commit) + err = r.walk(resolveRoot, func(p string, e fs.DirEntry) error { + rel, _ := filepath.Rel(resolveRoot, p) + path := filepath.ToSlash(rel) + if err := repository.Relative(path); err != nil { + return err + } + v, ok := files[path] + if !ok { + v = File{Path: path, Size: -1} + } + v.localPath = p + files[path] = v + return nil + }) + if err != nil { + return out, err + } + for _, v := range files { + if v.OID != "" { + if err := repository.Segment(v.OID); err != nil { + return out, fmt.Errorf("invalid cached object identity") + } + v.localPath = k.Blob(r.Root, v.OID) + } + // Catalogue manifests describe the cached upstream snapshot. Payload + // completeness belongs to transfer/file-serving paths and is deliberately + // not inspected while listing repositories or files. + v.CacheStatus = "unknown" + out.Files = append(out.Files, v) + } + sort.Slice(out.Files, func(i, j int) bool { return out.Files[i].Path < out.Files[j].Path }) + return out, nil +} + +func (r Reader) walk(root string, visit func(string, fs.DirEntry) error) error { + if err := r.safe(root); err != nil { + return err + } + err := filepath.WalkDir(root, func(p string, e fs.DirEntry, err error) error { + if err != nil { + return err + } + if e.IsDir() { + return nil + } + return visit(p, e) + }) + if errors.Is(err, os.ErrNotExist) { + if _, rootErr := os.Stat(root); errors.Is(rootErr, os.ErrNotExist) { + return nil + } + } + return err +} + +// Only resolve symlinks into this repository's blob directory. Never follow +// arbitrary links out of the cache or into a different repository. +func (r Reader) contentPath(k repository.RepoKey, p string) (string, error) { + if err := r.safe(filepath.Dir(p)); err != nil { + return "", err + } + info, err := os.Lstat(p) + if err != nil { + return "", err + } + if info.Mode()&os.ModeSymlink != 0 { + target, err := filepath.EvalSymlinks(p) + if err != nil { + return "", err + } + abs, _ := filepath.Abs(filepath.Join(k.FilesRoot(r.Root), "blobs")) + target, _ = filepath.Abs(target) + if filepath.Dir(target) != abs { + return "", fmt.Errorf("cached resolve link leaves repository blobs") + } + p = target + } + if err := r.safe(p); err != nil { + return "", err + } + return p, nil +} + +// inspect validates a bounded OLAH v8 header and actual bytes for every cached +// block. Corrupt OLAH is unknown, never treated as a complete plain file. +func inspect(f *os.File, expected int64) (status string, size, cached, offset int64, err error) { + info, err := f.Stat() + if err != nil { + return "unknown", 0, 0, 0, err + } + if !info.Mode().IsRegular() { + return "unknown", 0, 0, 0, fmt.Errorf("cache content is not a regular file") + } + var raw [36]byte + n, readErr := f.ReadAt(raw[:], 0) + if n < 4 || string(raw[:4]) != "OLAH" { + if expected < 0 || expected != info.Size() { + return "unknown", info.Size(), 0, 0, nil + } + return "complete", info.Size(), info.Size(), 0, nil + } + if readErr != nil { + return "unknown", expected, 0, 0, nil + } + version := binary.LittleEndian.Uint64(raw[4:12]) + block := binary.LittleEndian.Uint64(raw[12:20]) + length := binary.LittleEndian.Uint64(raw[20:28]) + bits := binary.LittleEndian.Uint64(raw[28:36]) + if version != 8 || block == 0 || block > math.MaxInt64 || length > math.MaxInt64 || bits > 8*1024*1024 || bits == 0 { + return "unknown", expected, 0, 0, nil + } + count := length / block + if length%block != 0 { + count++ + } + if count > bits { + return "unknown", expected, 0, 0, nil + } + mask := make([]byte, (bits+7)/8) + if _, err := f.ReadAt(mask, 36); err != nil { + return "unknown", expected, 0, 0, nil + } + size = int64(length) + offset = 36 + int64(len(mask)) + if expected >= 0 && expected != size { + return "unknown", size, 0, offset, nil + } + complete := true + for i := uint64(0); i < count; i++ { + start := i * block + end := start + min(block, length-start) + if mask[i/8]&(1<<(i%8)) != 0 && info.Size()-offset >= int64(end) { + cached += int64(end - start) + } else { + complete = false + } + } + if complete { + return "complete", size, cached, offset, nil + } + return "partial", size, cached, offset, nil +} + +var ErrIncomplete = errors.New("file is not completely cached") + +// OpenFile returns only an already complete local payload, pinned to an open +// read-only descriptor. Callers must close it; no cache manager is involved. +func (r Reader) OpenFile(typ, repo, revision, path string) (*os.File, *io.SectionReader, error) { + if err := repository.Relative(path); err != nil { + return nil, nil, err + } + m, err := r.Manifest(typ, repo, revision) + if err != nil { + return nil, nil, err + } + for _, v := range m.Files { + if v.Path == path { + if v.localPath == "" { + return nil, nil, os.ErrNotExist + } + p, err := r.contentPath(m.RepoKey, v.localPath) + if err != nil { + return nil, nil, err + } + f, err := os.Open(p) + if err != nil { + return nil, nil, err + } + state, size, _, offset, err := inspect(f, v.Size) + if err != nil { + f.Close() + return nil, nil, err + } + if state != "complete" { + f.Close() + return nil, nil, ErrIncomplete + } + return f, io.NewSectionReader(f, offset, size), nil + } + } + return nil, nil, os.ErrNotExist +} diff --git a/deploy/cachemount-sidecar/source/pkg/repository/repository.go b/deploy/cachemount-sidecar/source/pkg/repository/repository.go new file mode 100644 index 0000000..b3cfe0b --- /dev/null +++ b/deploy/cachemount-sidecar/source/pkg/repository/repository.go @@ -0,0 +1,383 @@ +// Package repository owns platform repository identity, layout and registration. +package repository + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "github.com/gofrs/flock" + "io/fs" + "net/url" + "os" + "path/filepath" + "strings" + "sync" + "time" + "unicode" + "unicode/utf8" +) + +const Marker = "repository.json" +const HuggingFace = "huggingface" +const ModelScope = "modelscope" +const Local = "dingo-local" +const Company = "datacanvas" + +type RepoKey struct { + Namespace string `json:"namespace"` + RepoType string `json:"repoType"` + Repo string `json:"repo"` +} + +func (k RepoKey) ID() string { return k.Namespace + "/" + k.Repo } +func (k RepoKey) Validate() error { + if k.RepoType != "models" && k.RepoType != "datasets" && k.RepoType != "spaces" { + return fmt.Errorf("invalid repoType") + } + if err := Segment(k.Namespace); err != nil { + return fmt.Errorf("namespace: %w", err) + } + if err := Relative(k.Repo); err != nil { + return fmt.Errorf("repo: %w", err) + } + return nil +} + +// ParseID is only for values already known to be full platform repository IDs. +func ParseID(repoType, id string) (RepoKey, error) { + p := strings.SplitN(id, "/", 2) + if len(p) != 2 { + return RepoKey{}, fmt.Errorf("repository ID must include namespace") + } + k := RepoKey{p[0], repoType, p[1]} + return k, k.Validate() +} + +func Segment(s string) error { + if s == "" || s == "." || s == ".." || !utf8.ValidString(s) || len(s) > 255 || strings.ContainsAny(s, `/\:*?"<>|`) || strings.TrimSpace(s) != s || strings.HasSuffix(s, ".") { + return fmt.Errorf("unsafe path segment %q", s) + } + for _, r := range s { + if unicode.IsControl(r) { + return fmt.Errorf("control character in path") + } + } + base := strings.ToLower(strings.SplitN(s, ".", 2)[0]) + if base == "con" || base == "prn" || base == "aux" || base == "nul" || (len(base) == 4 && (strings.HasPrefix(base, "com") || strings.HasPrefix(base, "lpt")) && base[3] >= '1' && base[3] <= '9') { + return fmt.Errorf("reserved platform path name") + } + return nil +} +func Relative(s string) error { + if len(s) > 1024 { + return fmt.Errorf("relative path exceeds 1024 bytes") + } + for _, p := range strings.Split(s, "/") { + if err := Segment(p); err != nil { + return err + } + } + return nil +} +func (k RepoKey) FilesRoot(root string) string { + if k.Namespace == ModelScope { + return filepath.Join(root, "modelscope", k.RepoType, filepath.FromSlash(k.Repo)) + } + return filepath.Join(root, "files", k.RepoType, k.storagePath()) +} +func (k RepoKey) APIRoot(root string) string { + if k.Namespace == ModelScope { + return filepath.Join(root, "api", k.RepoType, ModelScope, filepath.FromSlash(k.Repo)) + } + return filepath.Join(root, "api", k.RepoType, k.storagePath()) +} + +// Provider identity is virtual: HF keeps its original upstream directory and +// hosted namespaces live exclusively below the local upload root. +func (k RepoKey) storagePath() string { + if k.Namespace == HuggingFace { + return filepath.FromSlash(k.Repo) + } + // Existing pre-namespace uploads remain addressable until explicitly migrated. + if k.Namespace == Local { + return filepath.Join(Local, filepath.FromSlash(k.Repo)) + } + return filepath.Join(Local, k.Namespace, filepath.FromSlash(k.Repo)) +} +func (k RepoKey) Blob(root, sha string) string { return filepath.Join(k.FilesRoot(root), "blobs", sha) } +func (k RepoKey) Resolve(root, commit, path string) string { + return filepath.Join(k.FilesRoot(root), "resolve", commit, filepath.FromSlash(path)) +} +func EscapeURLPath(value string) string { + parts := strings.Split(value, "/") + for i := range parts { + parts[i] = url.PathEscape(parts[i]) + } + return strings.Join(parts, "/") +} +func (k RepoKey) PathsInfo(root, commit, path string) string { + return filepath.Join(k.APIRoot(root), "paths-info", commit, filepath.FromSlash(path)) +} +func (k RepoKey) Revision(root, revision string) string { + return filepath.Join(k.APIRoot(root), "revision", revision) +} + +type Descriptor struct { + RepoKey + Version int `json:"version"` + Source string `json:"source"` + Persistent bool `json:"persistent"` + Provider string `json:"provider,omitempty"` + UpstreamRepo string `json:"upstreamRepo,omitempty"` + Format string `json:"format"` +} + +func Hosted(k RepoKey) Descriptor { + return Descriptor{RepoKey: k, Version: 1, Source: "hosted", Persistent: true, Format: "dingcache"} +} +func Remote(k RepoKey) Descriptor { + return Descriptor{RepoKey: k, Version: 1, Source: "remote", Provider: k.Namespace, UpstreamRepo: k.Repo, Format: "dingcache"} +} +func (d Descriptor) Validate() error { + if err := d.RepoKey.Validate(); err != nil { + return err + } + if d.Version != 1 { + return fmt.Errorf("unsupported repository descriptor version") + } + switch d.Source { + case "hosted": + if !d.Persistent || d.Provider != "" || d.UpstreamRepo != "" || d.Format != "dingcache" || strings.EqualFold(d.Namespace, HuggingFace) || strings.EqualFold(d.Namespace, ModelScope) { + return fmt.Errorf("invalid hosted source policy") + } + case "remote": + if d.Persistent || (d.Provider != HuggingFace && d.Provider != ModelScope) || d.Namespace != d.Provider || d.UpstreamRepo != d.Repo { + return fmt.Errorf("invalid remote source policy") + } + if d.Format != "dingcache" { + return fmt.Errorf("invalid remote format") + } + default: + return fmt.Errorf("unknown repository source") + } + return nil +} +func Read(root string, k RepoKey) (Descriptor, error) { + if err := k.Validate(); err != nil { + return Descriptor{}, err + } + if err := SafePath(root, filepath.Join(k.APIRoot(root), Marker)); err != nil { + return Descriptor{}, err + } + b, err := os.ReadFile(filepath.Join(k.APIRoot(root), Marker)) + if err != nil { + return Descriptor{}, err + } + var d Descriptor + if err = json.Unmarshal(b, &d); err != nil { + return d, err + } + if d.RepoKey != k { + return d, fmt.Errorf("repository marker identity mismatch") + } + return d, d.Validate() +} + +func List(root string) ([]Descriptor, error) { + result := []Descriptor{} + err := filepath.WalkDir(filepath.Join(root, "api"), func(p string, e fs.DirEntry, err error) error { + if err != nil { + return err + } + if !e.IsDir() { + return nil + } + marker := filepath.Join(p, Marker) + info, statErr := os.Lstat(marker) + if errors.Is(statErr, os.ErrNotExist) { + return nil + } + if statErr != nil { + return statErr + } + if info.IsDir() { + return nil + } + b, err := os.ReadFile(marker) + if errors.Is(err, os.ErrNotExist) { + return nil + } + if err != nil { + return err + } + if err = SafePath(root, marker); err != nil { + return err + } + var d Descriptor + if err = json.Unmarshal(b, &d); err != nil { + return err + } + if err = d.Validate(); err != nil { + return err + } + if filepath.Clean(filepath.Join(d.APIRoot(root), Marker)) != filepath.Clean(marker) { + return fmt.Errorf("repository marker path mismatch: %s", marker) + } + result = append(result, d) + return filepath.SkipDir + }) + if errors.Is(err, os.ErrNotExist) { + return result, nil + } + return result, err +} + +var registrationMu sync.Mutex + +// Register serializes case and ancestor checks with marker creation. The lock +// file also rejects concurrent registration by another process sharing root. +func Register(root string, d Descriptor) error { + if err := d.Validate(); err != nil { + return err + } + if err := SafePath(root, filepath.Join(d.APIRoot(root), Marker)); err != nil { + return err + } + if err := SafePath(root, filepath.Join(d.FilesRoot(root), ".check")); err != nil { + return err + } + registrationMu.Lock() + defer registrationMu.Unlock() + if err := os.MkdirAll(filepath.Join(root, "api"), 0755); err != nil { + return err + } + lockPath := filepath.Join(root, "api", ".repository-registration.lock") + fileLock := flock.New(lockPath) + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + locked, err := fileLock.TryLockContext(ctx, 10*time.Millisecond) + if err != nil { + return err + } + if !locked { + return fmt.Errorf("repository registration busy") + } + defer fileLock.Unlock() + all, err := List(root) + if err != nil { + return err + } + for _, old := range all { + if old.RepoKey != d.RepoKey { + for _, pair := range [][2]string{{old.FilesRoot(root), d.FilesRoot(root)}, {old.APIRoot(root), d.APIRoot(root)}} { + a, b := strings.ToLower(filepath.Clean(pair[0])), strings.ToLower(filepath.Clean(pair[1])) + if a == b || strings.HasPrefix(a, b+string(filepath.Separator)) || strings.HasPrefix(b, a+string(filepath.Separator)) { + return fmt.Errorf("repository physical path conflict") + } + } + } + if strings.EqualFold(old.Namespace, d.Namespace) && old.Namespace != d.Namespace { + return fmt.Errorf("namespace case conflict") + } + if old.RepoType != d.RepoType || old.Namespace != d.Namespace { + continue + } + oldParts, newParts := strings.Split(old.Repo, "/"), strings.Split(d.Repo, "/") + for i := 0; i < len(oldParts) && i < len(newParts); i++ { + if !strings.EqualFold(oldParts[i], newParts[i]) { + break + } + if oldParts[i] != newParts[i] { + return fmt.Errorf("repository parent directory case conflict") + } + } + a, b := strings.ToLower(old.Repo), strings.ToLower(d.Repo) + if a == b { + if old.RepoKey == d.RepoKey && old == d { + return nil + } + return fmt.Errorf("repository identity or source conflict") + } + if strings.HasPrefix(a, b+"/") || strings.HasPrefix(b, a+"/") { + return fmt.Errorf("repository ancestor/descendant conflict") + } + } + dest := filepath.Join(d.APIRoot(root), Marker) + if err = os.MkdirAll(filepath.Dir(dest), 0755); err != nil { + return err + } + b, err := json.MarshalIndent(d, "", " ") + if err != nil { + return err + } + f, err := os.CreateTemp(filepath.Dir(dest), ".repository-*.tmp") + if err != nil { + return err + } + defer os.Remove(f.Name()) + if _, err = f.Write(b); err != nil { + f.Close() + return err + } + if err = f.Sync(); err != nil { + f.Close() + return err + } + if err = f.Close(); err != nil { + return err + } + return os.Rename(f.Name(), dest) +} + +// Locate finds the descriptor for a file by registered roots, never by guessing +// how many repository path segments precede blobs/resolve/provider-cache. +func Locate(root, file string) (Descriptor, string, error) { + all, err := List(root) + if err != nil { + return Descriptor{}, "", err + } + return LocateRegistered(root, file, all) +} + +// LocateRegistered reuses one validated registry snapshot for batch traversal. +func LocateRegistered(root, file string, all []Descriptor) (Descriptor, string, error) { + for _, d := range all { + rel, err := filepath.Rel(d.FilesRoot(root), file) + if err == nil && rel != "." && rel != ".." && !strings.HasPrefix(rel, ".."+string(filepath.Separator)) && !filepath.IsAbs(rel) { + return d, filepath.ToSlash(rel), nil + } + } + return Descriptor{}, "", os.ErrNotExist +} + +func SafePath(root, target string) error { + rootAbs, err := filepath.Abs(root) + if err != nil { + return err + } + targetAbs, err := filepath.Abs(target) + if err != nil { + return err + } + rel, err := filepath.Rel(rootAbs, targetAbs) + if err != nil { + return err + } + if rel == ".." || strings.HasPrefix(rel, ".."+string(filepath.Separator)) || filepath.IsAbs(rel) { + return fmt.Errorf("path escapes repository root") + } + for current := targetAbs; current != rootAbs; current = filepath.Dir(current) { + info, err := os.Lstat(current) + if os.IsNotExist(err) { + continue + } + if err != nil { + return err + } + if info.Mode()&os.ModeSymlink != 0 { + return fmt.Errorf("symlink in repository path: %s", current) + } + } + return nil +} diff --git a/deploy/cachemount-sidecar/source/pkg/repository/storage.go b/deploy/cachemount-sidecar/source/pkg/repository/storage.go new file mode 100644 index 0000000..0d3f602 --- /dev/null +++ b/deploy/cachemount-sidecar/source/pkg/repository/storage.go @@ -0,0 +1,47 @@ +package repository + +import ( + "fmt" + "strings" +) + +// FromStorage decodes unchanged remote identities and explicitly encoded hosted ones. +func FromStorage(repoType, org, repo string) (RepoKey, error) { + k := RepoKey{Namespace: "huggingface", RepoType: repoType, Repo: repo} + if strings.HasPrefix(org, "dingo-local/") { + k.Namespace = strings.TrimPrefix(org, "dingo-local/") + } else if strings.HasPrefix(org, "modelscope/") { + k.Namespace = "modelscope" + k.Repo = strings.TrimPrefix(org, "modelscope/") + "/" + repo + } else if org != "" { + k.Repo = org + "/" + repo + } + return k, k.Validate() +} + +// Storage leaves HF owner/repo exactly as old SQL and old peers expect. +func (k RepoKey) Storage() (org, repo string, err error) { + if err = k.Validate(); err != nil { + return + } + repo = k.Repo + switch k.Namespace { + case "huggingface": + parts := strings.SplitN(repo, "/", 2) + if len(parts) == 2 { + org, repo = parts[0], parts[1] + } + case "modelscope": + parts := strings.SplitN(repo, "/", 2) + if len(parts) != 2 { + return "", "", fmt.Errorf("ModelScope needs owner/repo") + } + org, repo = "modelscope/"+parts[0], parts[1] + default: + org = "dingo-local/" + k.Namespace + } + if k.Namespace != "huggingface" && (len(org) > 100 || len(repo) > 100 || len(org)+1+len(repo) > 100) { + err = fmt.Errorf("repository identity exceeds unchanged SQL VARCHAR(100) capacity") + } + return +} diff --git a/deploy/cachemount-sidecar/source/sidecar.py b/deploy/cachemount-sidecar/source/sidecar.py new file mode 100644 index 0000000..2a6cd34 --- /dev/null +++ b/deploy/cachemount-sidecar/source/sidecar.py @@ -0,0 +1,117 @@ +"""Expose the repository tree. No Spinfield or inference process starts here.""" +import json +import os +from pathlib import Path +import signal +import subprocess +import sys +import tempfile +import time + +STATE = Path("/run/cachemount") +READY = STATE / "ready.json" +stopping = False + + +def stop(signum, frame): + global stopping + stopping = True + + +def configuration(): + names = ("DINGO_CACHE_ROOT", "DINGO_MODEL_MOUNT") + values = [os.environ.get(n, "") for n in names] + if not all(values): + raise ValueError("required environment: " + ", ".join(names)) + raw, target = values + if not Path(raw).is_absolute() or not Path(target).is_absolute(): + raise ValueError("input and output must be absolute container paths") + root = Path(raw).resolve(strict=True) + mount = Path(target).resolve() + if not root.is_dir() or root == mount or root in mount.parents or mount in root.parents: + raise ValueError("input and output must be separate directory trees") + if not os.path.ismount(mount.parent): + raise ValueError("output must be a child of the dedicated shared volume mount") + return root, mount + + +def check(): + state = json.loads(READY.read_text()) + os.kill(state["pid"], 0) + if not os.path.ismount(state["mount"]): + raise RuntimeError("mount missing") + list(Path(state["mount"]).iterdir()) + for version in state["versions"]: + model = Path(version["source_file_path"]) + list(model.iterdir()) + if version["files"]: + with (model / version["files"][0]["path"]).open("rb") as f: + f.read(1) + + +def run(): + signal.signal(signal.SIGTERM, stop) + signal.signal(signal.SIGINT, stop) + STATE.mkdir(parents=True, exist_ok=True) + READY.unlink(missing_ok=True) + root, mount = configuration() + if os.path.ismount(mount): + raise RuntimeError("mount already exists; recreate this Pod after a mount failure") + mount.mkdir(exist_ok=True) + work = Path(tempfile.mkdtemp(prefix="session-", dir=STATE)) + report = work / "mapping.json" + helper = subprocess.Popen(["/usr/local/bin/cachemount", "-tree-root", str(root), "-report", str(report), str(mount)]) + try: + deadline = time.monotonic() + 120 + while not stopping: + if helper.poll() is not None: + raise RuntimeError("cachemount exited during startup") + if report.exists(): + try: + versions = json.loads(report.read_text()) + break + except json.JSONDecodeError: + pass + if time.monotonic() > deadline: + raise RuntimeError("mount startup timed out") + time.sleep(.1) + if stopping: + return + usable = [v for v in versions if v["status"] == "ready"] + if not usable: + raise RuntimeError("no usable repository versions; inspect cachemount logs") + state = {"pid": helper.pid, "mount": str(mount), "versions": usable} + pending = work / "ready.json" + pending.write_text(json.dumps(state)) + pending.replace(READY) + check() + print(json.dumps({"ready": True, "model_root": str(mount), "available_versions": len(usable), + "unavailable_versions": len(versions)-len(usable)}), flush=True) + while not stopping: + if helper.poll() is not None: + raise RuntimeError("cachemount exited; recreate the Pod before reloading the model") + time.sleep(.2) + finally: + READY.unlink(missing_ok=True) + if helper.poll() is None: + helper.terminate() + try: + helper.wait(timeout=20) + except subprocess.TimeoutExpired: + print("unmount busy; stopping helper; recreate Pod", file=sys.stderr) + helper.kill() + helper.wait() + raise RuntimeError("clean unmount failed") + + +if __name__ == "__main__": + try: + if sys.argv[1:] == ["check"]: + check() + elif not sys.argv[1:]: + run() + else: + raise ValueError("usage: sidecar.py [check]") + except Exception as exc: + print(str(exc), file=sys.stderr) + sys.exit(1) diff --git a/deploy/cachemount/README.md b/deploy/cachemount/README.md new file mode 100644 index 0000000..a01469b --- /dev/null +++ b/deploy/cachemount/README.md @@ -0,0 +1,19 @@ +# 消费者启动适配 + +`consumer_supervisor.py` 与 `healthcheck.py` 用于同容器启动 cachemount 和消费者。默认命令 `/modelfleet`，也可以传入其他消费者命令。 + +镜像需包含 Linux cachemount 二进制 `/usr/local/bin/cachemount`、Python 3、fuse3，以及在 passwd/group 中登记的消费者 UID/GID。挂载目录必须为空且归该用户所有；容器需访问 `/dev/fuse` 并允许 FUSE 挂载。容器运行时需按其安全策略配置 FUSE 所需权限。 + +环境变量：`DINGO_CACHE_ROOT` 默认 `/repos`，`DINGO_MODEL_MOUNT` 默认 `/mnt/dingo-models`，`DINGO_MOUNT_STATE` 默认 `/run/cachemount`。当前 Spinfield 健康检查按默认状态目录和 localhost:8082 配置；更改状态目录时应同时适配健康检查。 + +```text +python3 /opt/cachemount/consumer_supervisor.py /modelfleet +``` + +只读挂入原始 repos。挂载和应用必须使用同一个 UID、同一挂载命名空间。应用就绪应同时满足挂载和自身 HTTP 健康；helper 失败会停止应用，以便容器编排系统重新启动整个容器。 + +模型目录格式为 `/mnt/dingo-models/{namespace}/{完整仓库名}/{revision}`，可作为模型加载库的目录参数。完整仓库名可以包含 `/`，不能只保留最后一级名称。 + +如果实际模型加载进程运行在另一个容器、Pod 或节点，需在该消费者中启动 helper，或将节点上已挂载的正文目录挂入消费者；仅传递路径字符串不能让其他容器访问目录。挂载固定启动时的清单，更新版本需停止读者并重新挂载，期间源文件不能被原地修改。 + +构建及可重复验证步骤见 [只读挂载说明](../../docs/cachemount.md) 和 [统一模型目录](../../docs/model-tree-acceptance.md)。 diff --git a/deploy/cachemount/consumer_supervisor.py b/deploy/cachemount/consumer_supervisor.py new file mode 100644 index 0000000..3333ffc --- /dev/null +++ b/deploy/cachemount/consumer_supervisor.py @@ -0,0 +1,86 @@ +#!/usr/bin/env python3 +"""Run cachemount and a consumer in the same mount namespace and UID. +Exit on either child's failure. Never keep the application healthy after the +mount helper dies. Each boot uses a fresh mapping report; no stale readiness. +""" +import json +import os +from pathlib import Path +import signal +import subprocess +import sys +import tempfile +import time + +stopping = False + + +def stop(signum, frame): + global stopping + stopping = True + + +def terminate(p): + if p and p.poll() is None: + p.terminate() + try: + p.wait(timeout=15) + except subprocess.TimeoutExpired: + p.kill() + p.wait() + + +def main(): + signal.signal(signal.SIGTERM, stop) + signal.signal(signal.SIGINT, stop) + root = os.environ.get("DINGO_CACHE_ROOT", "/repos") + mount = os.environ.get("DINGO_MODEL_MOUNT", "/mnt/dingo-models") + state = Path(os.environ.get("DINGO_MOUNT_STATE", "/run/cachemount")) + state.mkdir(parents=True, exist_ok=True) + ready = state/"ready.json" + ready.unlink(missing_ok=True) + session = Path(tempfile.mkdtemp(prefix="session-", dir=state)) + report = session/"mapping.json" + helper = app = None + try: + helper = subprocess.Popen(["/usr/local/bin/cachemount", "-tree-root", root, "-report", str(report), mount]) + deadline = time.monotonic()+60 + while not stopping: + if helper.poll() is not None: + raise RuntimeError("cachemount exited before readiness") + if os.path.ismount(mount) and report.exists(): + try: + versions = json.loads(report.read_text()) + break + except json.JSONDecodeError: + pass + if time.monotonic() >= deadline: + raise RuntimeError("mount readiness timed out") + time.sleep(.1) + if stopping: + return 0 + if not any(v["status"] == "ready" for v in versions): + raise RuntimeError("no usable model version") + app = subprocess.Popen(sys.argv[1:] or ["/modelfleet"]) + temp = session/"ready.json" + temp.write_text(json.dumps({"uid": os.getuid(), "helper_pid": helper.pid, "consumer_pid": app.pid, + "mount": mount, "mapping": str(report)})) + temp.replace(ready) + print("mount ready; consumer started", flush=True) + while not stopping: + if helper.poll() is not None or app.poll() is not None or not os.path.ismount(mount): + raise RuntimeError("mount or consumer exited; stopping the pair") + time.sleep(.2) + return 0 + finally: + ready.unlink(missing_ok=True) + terminate(app) + terminate(helper) + + +if __name__ == "__main__": + try: + sys.exit(main()) + except Exception as e: + print(str(e), file=sys.stderr, flush=True) + sys.exit(1) diff --git a/deploy/cachemount/healthcheck.py b/deploy/cachemount/healthcheck.py new file mode 100644 index 0000000..3ae272a --- /dev/null +++ b/deploy/cachemount/healthcheck.py @@ -0,0 +1,11 @@ +import json +import os +from pathlib import Path +import urllib.request + +state = json.loads(Path("/run/cachemount/ready.json").read_text()) +assert os.path.ismount(state["mount"]) +os.kill(state["helper_pid"], 0) +os.kill(state["consumer_pid"], 0) +with urllib.request.urlopen("http://127.0.0.1:8082/api/v1/status", timeout=3) as response: + assert response.status == 200 diff --git a/docker/config/config.yaml b/docker/config/config.yaml index 5be053b..1d01949 100644 --- a/docker/config/config.yaml +++ b/docker/config/config.yaml @@ -75,7 +75,7 @@ diskClean: enabled: false cacheSizeLimit: 41781441855488 #38T 41781441855488 cacheCleanStrategy: "LRU" #LRU,FIFO,LARGE_FIRST - collectTimePeriod: 1 #定期检测磁盘使用量时间周期，单位小时（H） + collectTimePeriod: 24 #定期检测磁盘使用量时间周期，单位小时（H） dynamicProxy: enabled: false diff --git a/docs/cachemount.md b/docs/cachemount.md new file mode 100644 index 0000000..93b9c8e --- /dev/null +++ b/docs/cachemount.md @@ -0,0 +1,85 @@ +# DingCache 训练目录只读挂载 + +统一目录树模式及可重复验证步骤见 [统一模型目录：使用与用户验收](model-tree-acceptance.md)。下文的单目录入口继续保留。 + +消费者容器的启动、健康检查及模型目录接入见 [消费者启动适配](../deploy/cachemount/README.md)。 + +独立 Linux FUSE 工具，将现有 OLAH v8 缓存文件的正文展示成普通模型文件。不修改缓存格式，不启动 Speed/Scheduler，不回源下载，不生成权重副本。训练进程使用挂载目录作为本地模型路径。 + +## 构建及使用 + +需要 Linux、可访问的 `/dev/fuse` 和 `fusermount3`。在 dingospeed 仓库运行： + +```sh +CGO_ENABLED=0 go build -o cachemount ./cmd/cachemount +mkdir -p /path/to/empty-mount +./cachemount -cache-root /absolute/existing/cache -repo owner/model -revision main /path/to/empty-mount +``` + +HF 模式使用已有缓存元数据，固定启动时的清单和 commit。要求元数据包含完整 siblings 清单，并且所有文件已完整缓存为 OLAH v8。缺失、缺块、损坏、清单大小不符或目录冲突均拒绝整个挂载，不隐藏失败文件。尚不支持历史普通文件与 OLAH 混合目录。 + +其他来源可提供管理员准备的 JSON 清单；不要求调整原目录结构： + +```json +{ + "config.json": "/absolute/cache/config-blob", + "model-00001-of-00002.safetensors": "/absolute/cache/weight-blob-1", + "model-00002-of-00002.safetensors": "/absolute/cache/weight-blob-2" +} +``` + +```sh +./cachemount -manifest /absolute/files.json /path/to/empty-mount +``` + +清单键是训练侧相对路径，值是已有缓存文件绝对路径。清单由可信本地管理员提供；它不是面向任意用户的鉴权接口。工具默认只允许挂载用户访问，不启用 `allow_other`。跨用户、容器及分布式训练的挂载传播需要单独配置和验证。 + +训练完成并释放所有 mmap/文件引用后，另一个终端执行： + +```sh +fusermount3 -u /path/to/empty-mount +``` + +也可以向挂载进程发送 SIGINT/SIGTERM；忙碌时保留服务并报告错误，释放读者后可再次发送信号。checkpoint 和训练输出必须写到另一个可写目录。 + +## 数据及生命周期 + +- 挂载前只读取文件头和分块位图，检查完整性和物理长度；不全量读取权重，也不在启动时重算 SHA256。 +- 随机读映射为 `缓存头长度 + 逻辑偏移`。支持正常文件读取和只读 mmap，文件大小对外排除缓存头。 +- 所有源文件描述符保留到卸载。Linux 下源文件被 unlink、rename 或同路径替换后，挂载仍读取原 inode；底层空间在描述符关闭后才能释放。 +- **源文件必须在挂载期间保持内容不变。** 保留描述符不防止其他进程原地写入、truncate 或打洞。本工具没有接入 Speed 的写入/清理租约；不能宣称对任意原地修改提供快照隔离。读取遇到截断返回 EIO，但已缓存页可能仍可读取。 +- 已完成权重应按内容身份管理。挂载不会自动跟随分支更新；更新模型需要结束读者并重新挂载。 +- 无持久化权重副本；内核可能同时缓存底层和 FUSE 页，仍会占用内存。每个清单文件占用一个源描述符，大量分片需检查 `ulimit -n`。 +- 首版采用用户态读取转发，尚未调优大模型吞吐。FUSE 服务异常退出会导致未完成的读取失败，不能当作高可用存储层。 + +## 可重复验证 + +```sh +DINGO_TEST_FUSE=1 CGO_ENABLED=0 go test -v ./pkg/cachemount +CGO_ENABLED=0 go test ./pkg/hfprojection -run TestMountSources +CGO_ENABLED=0 go vet ./pkg/cachemount ./cmd/cachemount +python3 test/cachemount/validate.py --binary /absolute/cachemount --work /absolute/new-test-dir +python3 test/cachemount/validate.py --hf --binary /absolute/cachemount --work /absolute/another-new-test-dir +``` + +Python 测试只在新目录生成样本，验证摘要、mmap、跨进程读取、只读限制、缓存原内容不变，并记录挂载准备、普通文件与挂载文件的热缓存读取、导出落盘耗时。安装 PyTorch 后会实际执行 mmap 加载、5 步 CPU 训练、checkpoint 保存和恢复；未安装时报告 unavailable，不能算训练通过。 + +另安装 transformers、safetensors、numpy 后，测试会本地生成随机 tiny GPT2，保存为 safetensors，经挂载目录离线 `from_pretrained` 加载，比较原始目录与挂载目录输出，再训练 3 步、保存恢复并继续训练。无需下载上游模型。 + +Go 测试包含真实内核挂载、并发读取、缺块和损坏拒绝、目录冲突、截断错误、删除后读取，以及 512 GiB **逻辑大小的稀疏文件**头部准备测试。稀疏文件仅证明启动不扫描正文，不代表读取了 512 GiB 真实权重。 + +测试报告写入上述命令指定的新测试目录。小型模型和合成样本的结果不能代表实际大模型、GPU、多节点或长期运行情况；热缓存测试不代表磁盘性能，不能按样本线性外推大模型耗时。 + +### 已有大文件只读验证 + +`test/cachemount/validate_existing.py` 对现有完整 OLAH 文件执行两遍全量 SHA256：直接读取正文、经 FUSE 读取，均与仓库清单摘要比较；随后验证随机读取、跨块/末尾读取、mmap 及四线程读取。源文件只读打开，不复制正文、不清空系统页缓存，所有输出写入新的测试目录。 + +```sh +python3 -u test/cachemount/validate_existing.py \ + --binary /absolute/cachemount \ + --source /absolute/repos/files/models/owner/repo/blobs/HASH \ + --expected-sha256 HASH \ + --work /absolute/new-validation-dir +``` + +报告分别记录总读校验时间与 read 调用累计时间。后者排除了用户态摘要计算时间，但仍受缓存、调度及存储路径影响；顺序为先直读后挂载，不是受控冷缓存对照。Windows 盘经 WSL DrvFS/9p 访问的结果不能当作原生 Linux NVMe 吞吐。 diff --git a/docs/hf-readonly-projection.md b/docs/hf-readonly-projection.md new file mode 100644 index 0000000..8827337 --- /dev/null +++ b/docs/hf-readonly-projection.md @@ -0,0 +1,25 @@ +# HF 缓存只读投射 + +模型仓库展示使用以下独立 GET 接口，保留原有下载接口、HF owner/repo 数据库身份和缓存布局。 + +- `/api/repositories/{repoType}/huggingface/local-catalog`：发现本地 HF 仓库及已知 revision，不要求 `repository.json`。 +- `/api/repositories/{repoType}/huggingface/local-manifest?repo=Qwen/demo&revision=main`：返回已缓存的上游清单，不在列清单时读取文件正文或检查块完整性。 +- `/api/repositories/{repoType}/huggingface/local-file?repo=Qwen/demo&revision=main&path=README.md`：只读取已完整缓存的内容，支持 HTTP Range；缺失返回 404，未完整/未知返回 409。 + +三者只使用只读文件操作，无数据库、网络、注册、缓存创建或缓存修复行为。独立展示索引由 Spinfield 维护，不写回 HF 原目录。不要用原 `metadata/tree/file` 的回源读取代替这些展示接口。 + +仓库身份为 `namespace=huggingface`，repo 保留上游 owner/name。扫描排除 `dingo-local` 和 `modelscope` 物理根，不改变公司与个人上传归属。发现基于已有 `api///` 和 `files///` 下的缓存结构，未下载且无任何本地缓存事实的上游仓库不在该目录中。 + +`access` 为 public/restricted/unknown：只有元数据明确 `private=false` 且 `gated=false` 才确认为 public；任一已知版本受限即 restricted，属性缺失即 unknown。控制面将后两者限制为管理员可见，不能因为节点曾下载成功就推断可公开。这里只反映本地最后已知元数据，不声称实时查询了上游权限。分支指针已指向的 commit 不再重复列为 revision；没有分支指针时保留 commit 作为恢复浏览入口。 + +这些端点与现有 Speed 仓库接口共用内部服务信任边界；面向用户的鉴权由 Spinfield 执行。Speed 本身不是面向不可信用户的权限网关。 + +`metadataAvailable` 表示可读本地上游元数据；`manifestComplete` 表示其中存在 siblings 清单，不表示实时上游状态或本地文件完整。元数据缺失时仍展示 paths-info/resolve 中已知文件，清单完整性为 false。损坏的元数据明确报错，目录以单仓库 error 表达，不冒充空库。 + +清单中的 `cacheStatus` 固定为 unknown，`cachedBytes` 不表示已经核验的缓存量；`size=-1` 表示文件大小未知，`oid` 可为空。列出文件名不代表该文件可读。实际读取时才校验文件路径、链接目标、OLAH v8 头、块位图和物理长度；指向仓库外或其他仓库的链接被拒绝。普通文件要求已知大小与实际大小一致。HF 摘要不被强制视为 SHA256，读取检查也不重新计算内容摘要。 + +数据目录可能在下载过程中变化；读取错误须保留并显示异常，不可将错误当作删除或完整空清单。接口输出是一次读取观察，不提供跨文件事务快照。当前 cached metadata wrapper 上限 32 MiB，超限明确报错。 + +仅剩孤立 blobs、没有 revision/paths-info/resolve 证据的目录不会被猜测成可展示仓库。实际读取时，以 OLAH 魔数开头但头部无效的文件会被拒绝（包括无法与损坏容器区分的少见普通历史文件），不根据物理长度相等就当作完整内容返回。 + +验证：`go test ./pkg/hfprojection ./pkg/repository ./internal/router ./internal/handler ./internal/dao`。隔离浏览器验收物料位于父项目 `artifacts/hf-projection-20260917`，测试缓存以只读 mount 提供。 diff --git a/docs/model-tree-acceptance.md b/docs/model-tree-acceptance.md new file mode 100644 index 0000000..2bc5f02 --- /dev/null +++ b/docs/model-tree-acceptance.md @@ -0,0 +1,111 @@ +# 统一模型目录：使用与用户验收 + +本文说明统一模型目录的挂载方式及可重复验证步骤。验收报告由使用者在自己的测试环境中生成。 + +## 约定和实现 + +``` +source_file_path = /mnt/dingo-models/{namespace}/{完整仓库名}/{revision} +``` + +读取 `source_file_path/config.json`、权重分片或子目录文件，得到原始正文。namespace、完整仓库名、revision 保留原名；内部 SHA 和 blob 路径不出现在用户目录中。上传个人/公司 namespace 沿用 repository.json 的身份，旧上传兼容 dingo-local。HF 为 huggingface。 + +支持两种模式：`-tree-root` 自动发现一个本地缓存根的 HF/上传仓库版本；`-tree-config` 明确选择要提供的版本，支持多个缓存根。配置是选择仓库版本，不是手工填写每个文件的映射。 + +本次只提供 models，不包含 ModelScope、远程节点取回或同步。目录反映**挂载时固定的本地仓库清单**，不是不断跟随更新的 main。仓库页面应选择同一节点、同一版本，并在验收期间保持不变。使用结束后卸载、重新启动可取得新版本。已有授权规则不由本地 FUSE 替代：运行工具的管理员决定暴露范围，默认仅挂载用户访问，不启用 allow_other。 + +如果某个已发现版本缺文件、缺块或清单不可用：版本目录仍在，进入/读取返回 EIO；日志和映射报告写明原因。绝不显示成一个看似正常的部分目录。身份格式错误、版本路径有歧义或无法安全发现仓库则整个启动失败。合法空版本保留空目录。 + +```mermaid +flowchart LR + A[HF缓存元数据] --> C[版本解析与文件清单] + B[上传仓库原始清单] --> C + C --> D[固定原始路径与内容映射] + D --> E[/mnt/dingo-models/namespace/repo/revision] + U[用户控件选择文件夹] --> E + E --> F[只读 FUSE：正文偏移映射] + F --> G[已有 DingCache 正文] +``` + +上传读取与仓库文件树相同的 `dingo-local-manifest.json`，HF 复用现有 hfprojection 完整清单解析。挂载不写回仓库元数据、不迁移格式、不复制权重。保留源文件描述符到卸载，可保护 Linux unlink/rename 后的读取；**仍要求内容不被外部原地修改或 truncate，尚未接入全局清理租约**。 + +统一树模式也允许 HF 历史缓存中的普通文件，但要求清单大小已知且相符。具有 OLAH 魔数但头部损坏的文件不会退化成普通文件暴露。上传文件继续要求合法完整 OLAH。 + +## 先用看得懂的小样本验收 + +在有 FUSE 的 Linux/WSL 中，进入 dingospeed 源码目录。使用一个新的验收目录；报告也不覆盖旧文件。 + +```sh +CGO_ENABLED=0 go build -o ./cachemount ./cmd/cachemount +python3 test/cachemount/accept_tree.py demo --work "$HOME/dingo-tree-demo" +``` + +样本含两个完整模型目录和一个故意缺块的版本；有嵌套文件路径，并另外保存一份原始文件用于独立对照。权重是小型合成字节，不是可训练模型。 + +首次使用时，请管理员创建空目录 `/mnt/dingo-models` 并授权当前挂载用户使用。目录若非空，程序拒绝挂载。然后在终端 A 运行： + +```sh +./cachemount -tree-config "$HOME/dingo-tree-demo/selection.json" \ + -report "$HOME/dingo-tree-demo/mapping.json" /mnt/dingo-models +``` + +保持终端 A 运行。在终端 B 亲自看和读： + +```sh +find /mnt/dingo-models/huggingface/demo/HFModel/main -type f +find /mnt/dingo-models/datacanvas/team/UploadModel/v1 -type f +cat /mnt/dingo-models/datacanvas/team/UploadModel/v1/config.json +cmp "$HOME/dingo-tree-demo/originals/datacanvas/team/UploadModel/v1/model.safetensors" \ + /mnt/dingo-models/datacanvas/team/UploadModel/v1/model.safetensors +ls /mnt/dingo-models/datacanvas/team/Incomplete/v1 +``` + +你应看到：两个正常目录的三份文件均保留原名和子目录；cat 输出可读 JSON；cmp 无输出且退出码为 0；最后一个故意缺块的目录明确报 I/O 错误，不能正常进入空目录。 + +生成浏览器可打开的现场对照报告： + +```sh +python3 test/cachemount/accept_tree.py check \ + --report "$HOME/dingo-tree-demo/mapping.json" \ + --originals "$HOME/dingo-tree-demo/originals" \ + --output "$HOME/dingo-tree-demo/acceptance.html" --full +``` + +打开 acceptance.html。每个版本左右并排显示“仓库预期清单”和“实际目录”，逐文件列出大小和读取结果；配置文件展示内部 OLAH 头的十六进制，以及用户实际读到的 JSON。原始样本对照使用独立文件，不是把挂载结果自己和自己比较。 + +**认可条件：路径和文件数一致、缺失/额外均为空、原始样本内容一致、配置是正常 JSON、缺块目录被明确拒绝。任何一项不满足都不验收。** 这份报告只验证文件系统映射，不声称真实模型训练通过。 + +## 再用你自己的仓库验收 + +先停止读者，卸载上述样本： + +```sh +fusermount3 -u /mnt/dingo-models +``` + +自动发现真实仓库（将 `/absolute/repos` 替换为实际缓存根目录，并使用新报告文件）： + +```sh +./cachemount -tree-root /absolute/repos \ + -report "$HOME/dingo-real-mapping.json" /mnt/dingo-models +``` + +终端日志会打印每个实际 `source_file_path` 及 ready/unavailable。也可用 `-tree-config` 只选择你在仓库页面看到的版本： + +```json +[ + {"cache_root":"/absolute/repos","namespace":"huggingface","repo":"Qwen/demo","revision":"main"}, + {"cache_root":"/absolute/repos","namespace":"datacanvas","repo":"team/model","revision":"v1"} +] +``` + +在仓库页面选择相同节点及版本，逐项对照文件夹的原名、层级和大小，再在对方真实控件里选择打印的目录，使用其现有模型库加载。 + +```sh +python3 test/cachemount/accept_tree.py check --report "$HOME/dingo-real-mapping.json" \ + --output "$HOME/dingo-real-acceptance.html" +``` + +默认小文件全量比较，大于 16 MiB 的文件只比较头/中/尾，并明确标注“抽样，未全量”。要对大权重完整验收，增加 `--full`，将实际读取所有内容；有 64 位内容摘要时同时比对 SHA256，不要把默认抽样当成全量正确性证明。真实验收不要传 `--originals`，该参数只用于小型独立样本。 + +报告的预期树来自挂载时固定清单，不能代替与实际仓库页面的人工对照。最后一项验收必须是对方原有控件与模型加载库成功使用该目录。若加载进程在容器内，应先将此挂载目录暴露到容器中；服务器存在路径不等于其他机器也有这个路径。 diff --git a/docs/scheduler-discovery.md b/docs/scheduler-discovery.md new file mode 100644 index 0000000..c22e5e5 --- /dev/null +++ b/docs/scheduler-discovery.md @@ -0,0 +1,14 @@ +# Scheduler 注册端点与无令牌配置 + +2026-09-22：为先调通控制面链路，`/api/scheduler-registration` 的 GET/PUT +不再检查 DINGO_NODE_MANAGEMENT_TOKEN。调用方无须 Authorization；保留字段校验、 +已绑定身份不可变及活动缓存任务保护。传输设置原本就不要求该令牌。 + +注册请求新增 `managementUrl` / `downloadUrl`，同时保存到 `.scheduler-registration.json`。 +ModelFleet 下发时传自己的可访问地址；独立部署可配置 `upload.advertiseUrl`。 +未显式指定时，管理 URL 由注册 host 与 upload.port 生成；下载 URL 由注册 host/port 生成。 +需要 HTTPS、容器端口映射或路径前缀时填写完整 URL，不依靠默认推导。 + +Register gRPC 新增字段 5/6 上报上述端点，已有字段编号和业务协议保持不变。 +Scheduler `/api/v1/nodes/health?endpoints=true` 用于 ModelFleet 自动发现。 +本地 Standalone 和库存独立性保持不变；Scheduler 不可达不影响本地文件操作。 diff --git a/go.mod b/go.mod index a890084..f6ec8b2 100644 --- a/go.mod +++ b/go.mod @@ -13,16 +13,17 @@ require ( github.com/gofrs/flock v0.8.1 github.com/google/uuid v1.6.0 github.com/google/wire v0.6.0 + github.com/hanwen/go-fuse/v2 v2.8.0 github.com/klauspost/compress v1.18.0 github.com/labstack/echo/v4 v4.13.3 github.com/labstack/gommon v0.4.2 github.com/patrickmn/go-cache v2.1.0+incompatible github.com/prometheus/client_golang v1.22.0 + github.com/prometheus/client_model v0.6.1 github.com/robfig/cron/v3 v3.0.1 github.com/shirou/gopsutil v3.21.11+incompatible go.uber.org/zap v1.24.0 golang.org/x/sync v0.13.0 - golang.org/x/sys v0.35.0 google.golang.org/grpc v1.73.0 google.golang.org/protobuf v1.36.6 gopkg.in/natefinch/lumberjack.v2 v2.2.1 @@ -48,7 +49,6 @@ require ( github.com/mattn/go-isatty v0.0.20 // indirect github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect github.com/pkg/errors v0.9.1 // indirect - github.com/prometheus/client_model v0.6.1 // indirect github.com/prometheus/common v0.62.0 // indirect github.com/prometheus/procfs v0.15.1 // indirect github.com/twitchyliquid64/golang-asm v0.15.1 // indirect @@ -60,6 +60,7 @@ require ( golang.org/x/arch v0.15.0 // indirect golang.org/x/crypto v0.37.0 // indirect golang.org/x/net v0.39.0 // indirect + golang.org/x/sys v0.35.0 // indirect golang.org/x/text v0.24.0 // indirect google.golang.org/genproto/googleapis/rpc v0.0.0-20250324211829-b45e905df463 // indirect ) diff --git a/go.sum b/go.sum index 1a06626..15e1ae2 100644 --- a/go.sum +++ b/go.sum @@ -8,22 +8,14 @@ github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM= github.com/beorn7/perks v1.0.1/go.mod h1:G2ZrVWU2WbWT9wwq4/hrbKbnv/1ERSJQ0ibhJ6rlkpw= github.com/bytedance/gopkg v0.1.3 h1:TPBSwH8RsouGCBcMBktLt1AymVo2TVsBVCY4b6TnZ/M= github.com/bytedance/gopkg v0.1.3/go.mod h1:576VvJ+eJgyCzdjS+c4+77QF3p7ubbtiKARP3TxducM= -github.com/bytedance/sonic v1.13.2 h1:8/H1FempDZqC4VqjptGo14QQlJx8VdZJegxs6wwfqpQ= -github.com/bytedance/sonic v1.13.2/go.mod h1:o68xyaF9u2gvVBuGHPlUVCy+ZfmNNO5ETf1+KgkJhz4= github.com/bytedance/sonic v1.15.2 h1:90H+rcF/FwLXwfB1cudOLq/je83n683Utf4Cbp0xHCo= github.com/bytedance/sonic v1.15.2/go.mod h1:mT2NbXunuaEbnZ+mRIX/vYqKISmgEuHFDI4UzmKx2SA= -github.com/bytedance/sonic/loader v0.1.1/go.mod h1:ncP89zfokxS5LZrJxl5z0UJcsk4M4yY2JpfqGeCtNLU= -github.com/bytedance/sonic/loader v0.2.4 h1:ZWCw4stuXUsn1/+zQDqeE7JKP+QO47tz7QCNan80NzY= -github.com/bytedance/sonic/loader v0.2.4/go.mod h1:N8A3vUdtUebEY2/VQC0MyhYeKUFosQU6FxH2JmUe6VI= github.com/bytedance/sonic/loader v0.5.1 h1:Ygpfa9zwRCCKSlrp5bBP/b/Xzc3VxsAW+5NIYXrOOpI= github.com/bytedance/sonic/loader v0.5.1/go.mod h1:AR4NYCk5DdzZizZ5djGqQ92eEhCCcdf5x77udYiSJRo= 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/cloudwego/base64x v0.1.5 h1:XPciSp1xaq2VCSt6lF0phncD4koWyULpl5bUxbfCyP4= -github.com/cloudwego/base64x v0.1.5/go.mod h1:0zlkT4Wn5C6NdauXdJRhSKRlJvmclQ1hhJgA0rcu/8w= github.com/cloudwego/base64x v0.1.6 h1:t11wG9AECkCDk5fMSoxmufanudBtJ+/HemLstXDLI2M= github.com/cloudwego/base64x v0.1.6/go.mod h1:OFcloc187FXDaYHvrNIjxSe8ncn0OOM8gEHfghB2IPU= -github.com/cloudwego/iasm v0.2.0/go.mod h1:8rXZaNYT2n95jn+zTI1sDr+IgcD2GVs0nlbbQPiEFhY= github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E= github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= @@ -56,8 +48,6 @@ github.com/go-redsync/redsync/v4 v4.13.0 h1:49X6GJfnbLGaIpBBREM/zA4uIMDXKAh1NDkv github.com/go-redsync/redsync/v4 v4.13.0/go.mod h1:HMW4Q224GZQz6x1Xc7040Yfgacukdzu7ifTDAKiyErQ= github.com/gofrs/flock v0.8.1 h1:+gYjHKf32LDeiEEFhQaotPbLuUXjY5ZqxKgXy7n59aw= github.com/gofrs/flock v0.8.1/go.mod h1:F1TvTiK9OcQqauNUHlbJvyl9Qa1QvF/gOUDKA14jxHU= -github.com/gofrs/flock v0.12.1 h1:MTLVXXHf8ekldpJk3AKicLij9MdwOWkZ+a/jHHZby9E= -github.com/gofrs/flock v0.12.1/go.mod h1:9zxTsyu5xtJ9DK+1tFZyibEV7y3uwDxPPfbxeeHCoD0= github.com/golang/protobuf v1.5.4 h1:i7eJL8qZTpSEXOPTxNKhASYpMn+8e5Q6AdndVa1dWek= github.com/golang/protobuf v1.5.4/go.mod h1:lnTiLA8Wa4RWRcIUkrtSVa5nRhsEGBg48fD6rSs7xps= github.com/gomodule/redigo v1.8.9 h1:Sl3u+2BI/kk+VEatbj0scLdrFhjPmbxOc1myhDP41ws= @@ -70,6 +60,8 @@ github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= github.com/google/wire v0.6.0 h1:HBkoIh4BdSxoyo9PveV8giw7ZsaBOvzWKfcg/6MrVwI= github.com/google/wire v0.6.0/go.mod h1:F4QhpQ9EDIdJ1Mbop/NZBRB+5yrR6qg3BnctaoUk6NA= +github.com/hanwen/go-fuse/v2 v2.8.0 h1:wV8rG7rmCz8XHSOwBZhG5YcVqcYjkzivjmbaMafPlAs= +github.com/hanwen/go-fuse/v2 v2.8.0/go.mod h1:yE6D2PqWwm3CbYRxFXV9xUd8Md5d6NG0WBs5spCswmI= github.com/hashicorp/errwrap v1.0.0/go.mod h1:YH+1FKiLXxHSkmPseP+kNlulaMuP3n2brvKWEqk/Jc4= github.com/hashicorp/errwrap v1.1.0 h1:OxrOeh75EUXMY8TBjag2fzXGZ40LB6IKw45YeGUDY2I= github.com/hashicorp/errwrap v1.1.0/go.mod h1:YH+1FKiLXxHSkmPseP+kNlulaMuP3n2brvKWEqk/Jc4= @@ -77,10 +69,8 @@ github.com/hashicorp/go-multierror v1.1.1 h1:H5DkEtf6CXdFp0N0Em5UCwQpXMWke8IA0+l github.com/hashicorp/go-multierror v1.1.1/go.mod h1:iw975J/qwKPdAO1clOe2L8331t/9/fmwbPZ6JB6eMoM= github.com/klauspost/compress v1.18.0 h1:c/Cqfb0r+Yi+JtIEq73FWXVkRonBlf0CRNYc8Zttxdo= github.com/klauspost/compress v1.18.0/go.mod h1:2Pp+KzxcywXVXMr50+X0Q/Lsb43OQHYWRCY2AiWywWQ= -github.com/klauspost/cpuid/v2 v2.0.9/go.mod h1:FInQzS24/EEf25PyTYn52gqo7WaD8xa0213Md/qVLRg= github.com/klauspost/cpuid/v2 v2.2.10 h1:tBs3QSyvjDyFTq3uoc/9xFpCuOsJQFNPiAhYdw2skhE= github.com/klauspost/cpuid/v2 v2.2.10/go.mod h1:hqwkgyIinND0mEev00jJYCxPNVRVXFQeu1XKlok6oO0= -github.com/knz/go-libedit v1.10.1/go.mod h1:MZTVkCWyz0oBc7JOWP3wNAzd002ZbM/5hgShxwh4x8M= github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE= github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk= github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= @@ -97,6 +87,8 @@ github.com/mattn/go-colorable v0.1.14 h1:9A9LHSqF/7dyVVX6g0U9cwm9pG3kP9gSzcuIPHP github.com/mattn/go-colorable v0.1.14/go.mod h1:6LmQG8QLFO4G5z1gPvYEzlUgJ2wF+stgPZH1UqBm1s8= github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY= github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y= +github.com/moby/sys/mountinfo v0.7.2 h1:1shs6aH5s4o5H2zQLn796ADW1wMrIwHsyJ2v9KouLrg= +github.com/moby/sys/mountinfo v0.7.2/go.mod h1:1YOa8w8Ih7uW0wALDUgT1dTTSBrZ+HiBLGws92L2RU4= github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 h1:C3w9PqII01/Oq1c1nUAm88MOHcQC9l5mIlSMApZMrHA= github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822/go.mod h1:+n7T8mK8HuQTcFwEeznm/DIxMOiR9yIdICNftLE1DvQ= github.com/patrickmn/go-cache v2.1.0+incompatible h1:HRMgzkcYKYpi3C8ajMPV8OFXaaRUnok+kx1WdO15EQc= @@ -127,10 +119,8 @@ github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+ github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw= github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo= github.com/stretchr/objx v0.5.2/go.mod h1:FRsXN1f5AsAjCGJKqEizvkpNtU+EGNCLh3NxZ/8L+MA= -github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU= -github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4= github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo= github.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOfJA= github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY= @@ -244,4 +234,3 @@ gopkg.in/natefinch/lumberjack.v2 v2.2.1/go.mod h1:YD8tP3GAjkrDg1eZH7EGmyESg/lsYs gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= -nullprogram.com/x/optparse v1.0.0/go.mod h1:KdyPE+Igbe0jQUrVfMqDMeJQIJZEuyV7pjYmp6pbG50= diff --git a/internal/dao/cache_admin_dao.go b/internal/dao/cache_admin_dao.go index 2570f70..93e4c1b 100644 --- a/internal/dao/cache_admin_dao.go +++ b/internal/dao/cache_admin_dao.go @@ -16,6 +16,7 @@ package dao import ( "encoding/hex" + "encoding/json" "fmt" "io/fs" "net/http" @@ -29,6 +30,7 @@ import ( "dingospeed/pkg/common" "dingospeed/pkg/config" "dingospeed/pkg/consts" + "dingospeed/pkg/repository" "dingospeed/pkg/util" "github.com/bytedance/sonic" @@ -64,8 +66,10 @@ func NewCacheAdminDao(fileDao *FileDao) *CacheAdminDao { // CacheRepo 是仓库维度的汇总，供页面左侧的仓库树使用。 type CacheRepo struct { + Namespace string `json:"namespace"` + Repo string `json:"repo"` RepoType string `json:"repoType"` - OrgRepo string `json:"orgRepo"` + OrgRepo string `json:"-"` Source string `json:"source"` FileCount int `json:"fileCount"` TotalSize int64 `json:"totalSize"` @@ -80,8 +84,10 @@ type CacheRepo struct { // 同一份内容在同一个快照里被多个路径引用，那也是两行。行的身份是 (path, sha)， // 与一级删除的粒度严格一致。 type CacheFileRow struct { + Namespace string `json:"namespace"` + Repo string `json:"repo"` RepoType string `json:"repoType"` - OrgRepo string `json:"orgRepo"` + OrgRepo string `json:"-"` Path string `json:"path"` Sha string `json:"sha"` Size int64 `json:"size"` @@ -104,6 +110,8 @@ type CacheFileRow struct { // 一个 sha 一个文件，而不是每个仓库一份索引：并发删除各写各的文件， // 不需要读改写合并，天然原子。 type RecycleEntry struct { + Namespace string `json:"namespace"` + Repo string `json:"repo"` RepoType string `json:"repoType"` OrgRepo string `json:"orgRepo"` Sha string `json:"sha"` @@ -128,10 +136,12 @@ type RecycleRow struct { // DeleteItem 是删除请求里的一条。二级删除只用 RepoType/OrgRepo/Sha。 type DeleteItem struct { - RepoType string `json:"repoType"` - OrgRepo string `json:"orgRepo"` - Path string `json:"path"` - Sha string `json:"sha"` + Namespace string `json:"namespace"` + Repo string `json:"repo"` + RepoType string `json:"repoType"` + OrgRepo string `json:"-"` + Path string `json:"path"` + Sha string `json:"sha"` } // DeleteResult 逐条返回，不因为其中一条失败就整批回滚： @@ -235,11 +245,11 @@ func repoSource(orgRepo string) string { } func repoFilesRoot(repoType, orgRepo string) string { - return filepath.Join(config.SysConfig.Repos(), "files", repoType, filepath.FromSlash(orgRepo)) + return RepositoryKey(repoType, orgRepo).FilesRoot(config.SysConfig.Repos()) } func repoApiRoot(repoType, orgRepo string) string { - return filepath.Join(config.SysConfig.Repos(), "api", repoType, filepath.FromSlash(orgRepo)) + return RepositoryKey(repoType, orgRepo).APIRoot(config.SysConfig.Repos()) } func recycleEntryPath(repoType, orgRepo, sha string) string { @@ -390,7 +400,17 @@ func collectRemoteRefs(idx *repoIndex) { if len(parts) < 2 { return nil } - addRef(readPathsInfoOid(path), parts[0], strings.Join(parts[1:], "/")) + oid := readPathsInfoOid(path) + if stat := idx.Blobs[oid]; !stat.Exists { + if content, err := readCacheContent(path); err == nil { + var infos []common.PathsInfo + if sonic.Unmarshal(content, &infos) == nil && len(infos) > 0 { + stat.Size = infos[0].Size + idx.Blobs[oid] = stat + } + } + } + addRef(oid, parts[0], strings.Join(parts[1:], "/")) return nil }) @@ -483,20 +503,16 @@ type repoKey struct { // listRepoKeys 枚举所有仓库。 // -// orgRepo 的层级不固定（有 org 时是 org/repo，没有时就是 repo），所以不能按固定深度 -// glob，只能向下找到“含 blobs 或 resolve 子目录”的那一层作为仓库根，找到就不再深入。 +// 身份完全来自经过校验的描述符，不从 blobs/resolve 目录名推测仓库边界。 func listRepoKeys() []repoKey { - keys := make(map[repoKey]struct{}) - repos := config.SysConfig.Repos() - for _, repoType := range []string{"models", "datasets", "spaces"} { - scanRepoRoots(filepath.Join(repos, "files", repoType), repoType, []string{"blobs", "resolve"}, keys) - // paths-info 和 recycle 都是仓库根下的数据子树。它们本身足以证明仓库 - // 存在；命中后必须停止向下遍历，否则仓库发现会退化为扫描全部文件。 - scanRepoRoots(filepath.Join(repos, "api", repoType), repoType, []string{"revision", "paths-info", recycleDirName}, keys) - } - result := make([]repoKey, 0, len(keys)) - for k := range keys { - result = append(result, k) + all, err := repository.List(config.SysConfig.Repos()) + if err != nil { + zap.S().Errorf("repository enumeration failed: %v", err) + return nil + } + result := make([]repoKey, 0, len(all)) + for _, d := range all { + result = append(result, repoKey{RepoType: d.RepoType, OrgRepo: d.ID()}) } sort.Slice(result, func(i, j int) bool { if result[i].RepoType != result[j].RepoType { @@ -507,25 +523,6 @@ func listRepoKeys() []repoKey { return result } -func scanRepoRoots(root, repoType string, markers []string, out map[repoKey]struct{}) { - _ = filepath.WalkDir(root, func(path string, entry fs.DirEntry, err error) error { - if err != nil || !entry.IsDir() || path == root { - return nil - } - for _, marker := range markers { - if entry.Name() == marker { - // 命中标记目录，它的父目录就是仓库根。 - rel, relErr := filepath.Rel(root, filepath.Dir(path)) - if relErr == nil && rel != "." && !strings.HasPrefix(rel, "..") { - out[repoKey{RepoType: repoType, OrgRepo: filepath.ToSlash(rel)}] = struct{}{} - } - return fs.SkipDir - } - } - return nil - }) -} - // --------------------------------------------------------------------------- // 一级列表 // --------------------------------------------------------------------------- @@ -537,6 +534,7 @@ func (d *CacheAdminDao) ListRepos() []*CacheRepo { rows := indexRows(idx) repo := &CacheRepo{ RepoType: key.RepoType, + Namespace: RepositoryKey(key.RepoType, key.OrgRepo).Namespace, Repo: RepositoryKey(key.RepoType, key.OrgRepo).Repo, OrgRepo: key.OrgRepo, Source: idx.Source, FileCount: len(rows), @@ -597,7 +595,8 @@ func indexRows(idx *repoIndex) []*CacheFileRow { row, ok := byKey[key] if !ok { row = &CacheFileRow{ - RepoType: idx.RepoType, + RepoType: idx.RepoType, + Namespace: RepositoryKey(idx.RepoType, idx.OrgRepo).Namespace, Repo: RepositoryKey(idx.RepoType, idx.OrgRepo).Repo, OrgRepo: idx.OrgRepo, Path: ref.Path, Sha: sha, @@ -684,7 +683,8 @@ func indexOrphans(idx *repoIndex) []*RecycleRow { continue } entry = RecycleEntry{ - RepoType: idx.RepoType, + RepoType: idx.RepoType, + Namespace: RepositoryKey(idx.RepoType, idx.OrgRepo).Namespace, Repo: RepositoryKey(idx.RepoType, idx.OrgRepo).Repo, OrgRepo: idx.OrgRepo, Sha: sha, Size: stat.Size, @@ -694,6 +694,8 @@ func indexOrphans(idx *repoIndex) []*RecycleRow { UnlinkedAt: stat.ModTime.Unix(), } } + entry.Namespace = RepositoryKey(idx.RepoType, idx.OrgRepo).Namespace + entry.Repo = RepositoryKey(idx.RepoType, idx.OrgRepo).Repo if entry.Size == 0 { entry.Size = stat.Size } @@ -785,7 +787,18 @@ func (d *CacheAdminDao) softDeleteRepo(key repoKey, items []DeleteItem) []*Delet uploadRepoLocks.Lock(lockKey) defer uploadRepoLocks.Unlock(lockKey) + if IsLocalOrgRepo(key.OrgRepo) { + if err := NewUploadDao(d.fileDao, nil).recoverInventoryLocked(RepositoryKey(key.RepoType, key.OrgRepo)); err != nil { + return failAll(items, err.Error()) + } + } + if err := validateRepositoryReferences(key.RepoType, key.OrgRepo); err != nil { + return failAll(items, err.Error()) + } idx := buildRepoIndex(key.RepoType, key.OrgRepo) + if idx.Source == CacheSourceUpload { + return d.journalSoftDelete(idx, items) + } results := make([]*DeleteResult, 0, len(items)) for _, item := range items { results = append(results, d.softDeleteOne(idx, item)) @@ -846,7 +859,8 @@ func (d *CacheAdminDao) softDeleteOne(idx *repoIndex, item DeleteItem) *DeleteRe stat := idx.Blobs[item.Sha] paths := []string{item.Path} if err := writeRecycleEntry(RecycleEntry{ - RepoType: idx.RepoType, + RepoType: idx.RepoType, + Namespace: RepositoryKey(idx.RepoType, idx.OrgRepo).Namespace, Repo: RepositoryKey(idx.RepoType, idx.OrgRepo).Repo, OrgRepo: idx.OrgRepo, Sha: item.Sha, Size: stat.Size, @@ -1060,6 +1074,11 @@ func (d *CacheAdminDao) purgeOne(item DeleteItem) *DeleteResult { // 重新算一遍引用而不是信任前端传来的状态：从页面加载到点确认之间， // 完全可能有一次发布或下载把这份内容重新引用起来。 + if err := validateRepositoryReferences(item.RepoType, item.OrgRepo); err != nil { + result.Status = "failed" + result.Reason = err.Error() + return result + } idx := buildRepoIndex(item.RepoType, item.OrgRepo) if len(idx.BySha[item.Sha]) > 0 { removeRecycleEntry(item.RepoType, item.OrgRepo, item.Sha) @@ -1089,6 +1108,15 @@ func (d *CacheAdminDao) purgeOne(item DeleteItem) *DeleteResult { // 的动作，必须与正在写入的分块上传互斥。Linux 上少了这条，unlink 之后分块 writer // 会继续往孤儿 inode 上写，位图更新全部丢失且无人知晓。 func reclaimBlobFile(repoType, orgRepo, sha string) error { + if err := validateRepoKey(repoKey{RepoType: repoType, OrgRepo: orgRepo}); err != nil { + return err + } + if err := repository.Segment(sha); err != nil { + return err + } + if err := repository.SafePath(config.SysConfig.Repos(), BlobPath(repoType, orgRepo, sha)); err != nil { + return err + } blobPath := localBlobPath(repoType, orgRepo, sha) blobKey := uploadBlobLockKey(repoType, orgRepo, sha) uploadBlobLocks.Lock(blobKey) @@ -1105,16 +1133,16 @@ func reclaimBlobFile(repoType, orgRepo, sha string) error { } func validateRepoKey(key repoKey) error { - if key.RepoType != "models" && key.RepoType != "datasets" && key.RepoType != "spaces" { - return fmt.Errorf("invalid repoType: %s", key.RepoType) + k, err := repository.ParseID(key.RepoType, key.OrgRepo) + if err != nil { + return err } - if key.OrgRepo == "" { - return fmt.Errorf("orgRepo is required") + d, err := repository.Read(config.SysConfig.Repos(), k) + if err != nil { + return fmt.Errorf("repository policy unavailable: %w", err) } - // orgRepo 直接参与拼路径，必须挡住 .. 与绝对路径。 - clean := filepath.ToSlash(filepath.Clean(filepath.FromSlash(key.OrgRepo))) - if clean != key.OrgRepo || strings.HasPrefix(clean, "..") || strings.HasPrefix(clean, "/") || filepath.IsAbs(clean) { - return fmt.Errorf("invalid orgRepo: %s", key.OrgRepo) + if d.Format != "dingcache" { + return fmt.Errorf("provider cache management requires provider-aware operation") } return nil } @@ -1126,3 +1154,48 @@ func failAll(items []DeleteItem, reason string) []*DeleteResult { } return results } + +// Validate all durable reference records before deletion. Listing can be best +// effort; interpreting an unreadable record as an unreferenced blob cannot. +func validateRepositoryReferences(repoType, id string) error { + if err := validateRepoKey(repoKey{RepoType: repoType, OrgRepo: id}); err != nil { + return err + } + root := repoApiRoot(repoType, id) + return filepath.WalkDir(root, func(p string, e fs.DirEntry, err error) error { + if err != nil { + return err + } + if e.IsDir() { + return nil + } + if !strings.HasSuffix(e.Name(), ".json") { + return nil + } + if err = repository.SafePath(config.SysConfig.Repos(), p); err != nil { + return err + } + b, err := os.ReadFile(p) + if err != nil { + return err + } + var value interface{} + return sonic.Unmarshal(b, &value) + }) +} + +// Persisted tombstones retain their historical full-ID field. The HTTP row +// exposes only independent namespace/repo fields. +func (r RecycleRow) MarshalJSON() ([]byte, error) { + type plain RecycleRow + b, err := json.Marshal(plain(r)) + if err != nil { + return nil, err + } + var fields map[string]interface{} + if err = json.Unmarshal(b, &fields); err != nil { + return nil, err + } + delete(fields, "orgRepo") + return json.Marshal(fields) +} diff --git a/internal/dao/cache_admin_dao_test.go b/internal/dao/cache_admin_dao_test.go index edf8e41..90141cd 100644 --- a/internal/dao/cache_admin_dao_test.go +++ b/internal/dao/cache_admin_dao_test.go @@ -8,6 +8,7 @@ import ( "dingospeed/internal/downloader" "dingospeed/pkg/config" + "dingospeed/pkg/repository" "dingospeed/pkg/util" "github.com/bytedance/sonic" @@ -61,6 +62,9 @@ func writeBlobContent(t *testing.T, blobPath string, content []byte) { func writePathsInfo(t *testing.T, fileDao *FileDao, repoType, orgRepo, commit, path, oid string, size int64) { t.Helper() + if err := repository.Register(config.SysConfig.Repos(), repository.Remote(RepositoryKey(repoType, orgRepo))); err != nil { + t.Fatal(err) + } body, err := sonic.Marshal([]map[string]interface{}{{ "type": "file", "oid": oid, @@ -83,6 +87,9 @@ func writePathsInfo(t *testing.T, fileDao *FileDao, repoType, orgRepo, commit, p // seedRemoteFile 造一份“从上游拉下来并缓存住了”的文件。 func seedRemoteFile(t *testing.T, fileDao *FileDao, orgRepo, commit, path string, content []byte) string { t.Helper() + if err := repository.Register(config.SysConfig.Repos(), repository.Remote(RepositoryKey("models", orgRepo))); err != nil { + t.Fatal(err) + } etag := sha256Hex(content) blobPath := BlobPath("models", orgRepo, etag) writeBlobContent(t, blobPath, content) @@ -118,6 +125,9 @@ func findOrphan(rows []*RecycleRow, sha string) *RecycleRow { func TestListRepoKeysStopsAtPathsInfo(t *testing.T) { _, _, _ = newTestCacheAdminDao(t) orgRepo := "dingo-local/paths-heavy" + if err := RegisterHosted("models", "dingo-local", "paths-heavy"); err != nil { + t.Fatal(err) + } // revision 是一个合法仓库标记，但这里故意把同名目录放在 paths-info // 深处。仓库发现如果进入数据子树，会把这个深层路径误判成另一个仓库。 @@ -215,7 +225,7 @@ func TestListFilesKeepsOneRowPerPathForSharedContent(t *testing.T) { func TestListFilesShowsRemoteCache(t *testing.T) { admin, u, _ := newTestCacheAdminDao(t) - const orgRepo = "Qwen/Qwen2.5-0.5B" + const orgRepo = "huggingface/Qwen/Qwen2.5-0.5B" content := []byte("remote model shard") etag := seedRemoteFile(t, u.fileDao, orgRepo, "abc123", "model.safetensors", content) @@ -239,7 +249,7 @@ func TestListFilesShowsRemoteCache(t *testing.T) { // 部分缓存必须显示成部分：页面上把半个文件显示成“完整”会让人误以为删了也能重下。 func TestListFilesReportsPartialCache(t *testing.T) { admin, u, _ := newTestCacheAdminDao(t) - const orgRepo = "Qwen/partial" + const orgRepo = "huggingface/Qwen/partial" full := make([]byte, testBlockSize*4) for i := range full { @@ -290,7 +300,7 @@ func clearLastBlock(t *testing.T, blobPath string) { // 就会以 O_RDWR 打开并可能重写头部，页面刷一下就把缓存改了。 func TestListFilesDoesNotMutateCache(t *testing.T) { admin, u, _ := newTestCacheAdminDao(t) - const orgRepo = "Qwen/readonly" + const orgRepo = "huggingface/Qwen/readonly" content := []byte("must not be touched") etag := seedRemoteFile(t, u.fileDao, orgRepo, "c1", "f.bin", content) @@ -325,7 +335,7 @@ func TestListReposAggregatesBothSources(t *testing.T) { payload := []byte("published") mustStage(t, u, deferredParam("a.bin", payload), payload) mustPublish(t, u, publishParam("main", manifestItem("a.bin", payload))) - seedRemoteFile(t, u.fileDao, "Qwen/remote-demo", "c1", "f.bin", []byte("remote")) + seedRemoteFile(t, u.fileDao, "huggingface/Qwen/remote-demo", "c1", "f.bin", []byte("remote")) repos := admin.ListRepos() var upload, remote *CacheRepo @@ -592,7 +602,7 @@ func TestReuploadAfterSoftDeleteReusesContentAndVoidsTombstone(t *testing.T) { func TestSoftDeleteRemoteRemovesResolveAndPathsInfo(t *testing.T) { admin, u, _ := newTestCacheAdminDao(t) - const orgRepo = "Qwen/Qwen2.5-0.5B" + const orgRepo = "huggingface/Qwen/Qwen2.5-0.5B" content := []byte("cached from upstream") etag := seedRemoteFile(t, u.fileDao, orgRepo, "c1", "model.bin", content) @@ -629,7 +639,7 @@ func TestSoftDeleteRemoteRemovesResolveAndPathsInfo(t *testing.T) { // 删掉了，把它显示成“待彻底删除”会怂恿用户删掉正常的缓存。 func TestRemoteBlobWithoutTombstoneIsNotListedAsOrphan(t *testing.T) { admin, u, _ := newTestCacheAdminDao(t) - const orgRepo = "Qwen/lru-victim" + const orgRepo = "huggingface/Qwen/lru-victim" content := []byte("resolve was evicted by lru") etag := seedRemoteFile(t, u.fileDao, orgRepo, "c1", "f.bin", content) @@ -824,7 +834,7 @@ func TestLegacyUnreferencedContentStillUsesBlobMTime(t *testing.T) { // 远端缓存只回收有墓碑的：无差别回收等于改动磁盘清理对公开模型的既有行为。 func TestCleanupRecycledBlobsOnlyTouchesTombstonedRemoteContent(t *testing.T) { admin, u, _ := newTestCacheAdminDao(t) - const orgRepo = "Qwen/Qwen2.5-0.5B" + const orgRepo = "huggingface/Qwen/Qwen2.5-0.5B" deleted := []byte("deleted by the operator") untouched := []byte("still referenced") diff --git a/internal/dao/cache_reconcile_test.go b/internal/dao/cache_reconcile_test.go new file mode 100644 index 0000000..cf308cc --- /dev/null +++ b/internal/dao/cache_reconcile_test.go @@ -0,0 +1,65 @@ +package dao + +import ( + "context" + "errors" + "testing" + + "dingospeed/internal/data" + "dingospeed/internal/downloader" + "dingospeed/pkg/config" + "dingospeed/pkg/consts" + pb "dingospeed/pkg/proto/manager" + "dingospeed/pkg/repository" + "google.golang.org/grpc" + "google.golang.org/protobuf/types/known/emptypb" +) + +type reconcileClient struct { + storageClient + fail bool +} + +func (c *reconcileClient) SyncFileProcess(ctx context.Context, r *pb.SyncFileProcessReq, opts ...grpc.CallOption) (*emptypb.Empty, error) { + if ctx.Err() != nil { + return nil, ctx.Err() + } + c.sync = r + if c.fail { + return nil, errors.New("scheduler unavailable") + } + return &emptypb.Empty{}, nil +} + +func TestCacheReconcileIdentityAndRetry(t *testing.T) { + old := config.SysConfig + defer func() { config.SysConfig = old }() + config.SysConfig = &config.Config{} + config.SysConfig.Scheduler.Mode = consts.SchedulerModeCluster + config.SysConfig.Scheduler.OriginMode = consts.SchedulerModeCluster + config.SysConfig.Scheduler.Discovery.InstanceId = "node-2" + data.NewBaseData() + c := &reconcileClient{} + d := NewDownloaderDao(&SchedulerDao{Client: c}) + ctx, cancel := context.WithCancel(context.Background()) + cancel() + p := &downloader.TaskParam{Context: ctx, DataType: "models", FileName: "weights.bin", Etag: "hash", FileSize: 7, + RepoKey: repository.RepoKey{Namespace: "huggingface", RepoType: "models", Repo: "Qwen/demo"}} + for i := 0; i < 2; i++ { + d.reconcileCachedFile(p) + e := c.sync.FileProcessEntries[0] + if e.Org != "Qwen" || e.Repo != "demo" || e.InstanceId != "node-2" || e.ProcessId != 0 || e.StartPos != 7 || e.EndPos != 7 || e.Status != consts.StatusDownloaded { + t.Fatalf("incorrect reconciliation: %v", e) + } + } + c.fail = true + d.reconcileCachedFile(p) + select { + case op := <-data.GetLocalOperationChan(): + if op.Type != consts.OperationProcess || op.IdentityVersion != 2 { + t.Fatalf("wrong retry: %+v", op) + } + default: + t.Fatal("failed reconciliation was not queued for retry") + } +} diff --git a/internal/dao/downloader_dao.go b/internal/dao/downloader_dao.go index c237f7f..b95d68a 100644 --- a/internal/dao/downloader_dao.go +++ b/internal/dao/downloader_dao.go @@ -16,18 +16,21 @@ package dao import ( "context" + "errors" "fmt" "net/http" + "net/url" "sync" "time" + "dingospeed/internal/data" "dingospeed/internal/downloader" "dingospeed/pkg/common" "dingospeed/pkg/config" "dingospeed/pkg/consts" myerr "dingospeed/pkg/error" "dingospeed/pkg/proto/manager" - "dingospeed/pkg/util" + "dingospeed/pkg/repository" "go.uber.org/zap" ) @@ -44,6 +47,14 @@ func NewDownloaderDao(schedulerDao *SchedulerDao) *DownloaderDao { // 整个文件 func (d *DownloaderDao) FileDownload(startPos, endPos int64, isInnerRequest bool, taskParam *downloader.TaskParam) error { + key, identityErr := repository.ParseID(taskParam.DataType, taskParam.OrgRepo) + if identityErr != nil { + return identityErr + } + taskParam.RepoKey = key + if !taskParam.LocalOnly && key.Namespace == repository.HuggingFace && config.SysConfig.IsCluster() && taskParam.FileSize > config.SysConfig.GetMinimumFileSize() { + taskParam.OnCacheComplete = func() { d.reconcileCachedFile(taskParam) } + } dingCacheManager := downloader.GetInstance() dingFile, err := dingCacheManager.GetDingFile(taskParam.BlobsFile, taskParam.FileSize) if err != nil { @@ -53,13 +64,15 @@ func (d *DownloaderDao) FileDownload(startPos, endPos int64, isInnerRequest bool taskParam.DingFile = dingFile tasks, err := d.constructTask(startPos, endPos, isInnerRequest, taskParam) if err != nil { + dingCacheManager.ReleasedDingFile(taskParam.BlobsFile) return err } + // Both OutResult and DoTask read TaskSize. Initialize it before either starts. + for _, task := range tasks { + task.SetTaskSize(len(tasks)) + } go func() { - defer close(taskParam.ResponseChan) - defer func() { - dingCacheManager.ReleasedDingFile(taskParam.BlobsFile) - }() + var wg sync.WaitGroup wg.Add(1) go func() { @@ -72,7 +85,11 @@ func (d *DownloaderDao) FileDownload(startPos, endPos int64, isInnerRequest bool } task := tasks[i] if i == 0 { - task.GetResponseChan() <- []byte{} // 先建立长连接 + select { + case task.GetResponseChan() <- []byte{}: + case <-taskParam.Context.Done(): + return + } } task.OutResult() } @@ -87,6 +104,25 @@ func (d *DownloaderDao) FileDownload(startPos, endPos int64, isInnerRequest bool }() } wg.Wait() // 等待协程池所有远程下载任务执行完毕 + complete, _ := analysisFilePosition(dingFile, 0, taskParam.FileSize) + dingCacheManager.ReleasedDingFile(taskParam.BlobsFile) + close(taskParam.ResponseChan) + if taskParam.CacheResult != nil { + var result error + if !complete { + result = fmt.Errorf("file cache incomplete") + if taskParam.Context.Err() != nil { + result = taskParam.Context.Err() + } + } + for _, child := range tasks { + if remote, ok := child.(*downloader.RemoteFileTask); ok && remote.ResultError != nil && !errors.Is(remote.ResultError, context.Canceled) { + result = remote.ResultError + break + } + } + taskParam.CacheResult <- result + } }() return nil } @@ -108,7 +144,7 @@ func (d *DownloaderDao) constructTask(startPos, endPos int64, isInnerRequest boo return nil, myerr.NewAppendCode(http.StatusNotFound, "Entry not found") } // isInnerRequest为true，即内部请求，是已经被调度过后，设置为内部域名的请求，这种请求将不会再次参与调度，直接做下载即可。 - if !isInnerRequest && config.SysConfig.IsCluster() && !fileComplete { + if !taskParam.LocalOnly && !isInnerRequest && config.SysConfig.IsCluster() && !fileComplete { if response, err := d.getRequestDomainScheduler(taskParam.DataType, taskParam.OrgRepo, taskParam.FileName, taskParam.Etag, curPos, endPos, taskParam.FileSize); err != nil { zap.S().Errorf("getRequestDomainScheduler err.%v", err) goto localTask @@ -122,15 +158,18 @@ func (d *DownloaderDao) constructTask(startPos, endPos int64, isInnerRequest boo } speedDomain := fmt.Sprintf("http://%s:%d", response.Host, response.Port) // 此刻向该节点发起远程下载请求 if endPos <= response.MaxOffset { + taskParam.Peer = true taskParam.Domain = speedDomain speedTasks := getContiguousRanges(curPos, endPos, taskParam) tasks = append(tasks, speedTasks...) } else { // 需要重新拆分任务 + taskParam.Peer = true taskParam.Domain = speedDomain beforeTasks := getContiguousRanges(curPos, response.MaxOffset, taskParam) tasks = append(tasks, beforeTasks...) - taskParam.Domain = config.SysConfig.GetHFURLBase() + taskParam.Peer = false + taskParam.Domain = upstreamDomain(taskParam) afterTasks := getContiguousRanges(response.MaxOffset, endPos, taskParam) tasks = append(tasks, afterTasks...) } @@ -144,20 +183,32 @@ func (d *DownloaderDao) constructTask(startPos, endPos int64, isInnerRequest boo } localTask: - taskParam.Domain = config.SysConfig.GetHFURLBase() + taskParam.Peer = false + taskParam.Domain = upstreamDomain(taskParam) tasks = getContiguousRanges(startPos, endPos, taskParam) return tasks, nil } +func upstreamDomain(p *downloader.TaskParam) string { + if p.Source != nil { + return p.Source.Domain + } + return config.SysConfig.GetHFURLBase() +} + func (d *DownloaderDao) getRequestDomainScheduler(dataType, orgRepo, fileName, etag string, startPos, endPos, fileSize int64) (*manager.SchedulerFileResponse, error) { - org, repo := util.SplitOrgRepo(orgRepo) + key, err := repository.ParseID(dataType, orgRepo) + if err != nil { + return nil, err + } + org, repo := key.Namespace, key.Repo response, err := d.schedulerDao.SchedulerFile(&manager.SchedulerFileRequest{ DataType: dataType, Org: org, Repo: repo, Name: fileName, Etag: etag, - InstanceId: config.SysConfig.Scheduler.Discovery.InstanceId, + InstanceId: config.SysConfig.Registration().NodeID, StartPos: startPos, EndPos: endPos, FileSize: fileSize, @@ -190,13 +241,18 @@ func doTask(ctx context.Context, tasks []common.DownloadTask) { return } task := tasks[i] - task.SetTaskSize(taskLen) if err := pool.Submit(ctx, task); err != nil { zap.S().Errorf("submit task err.%v", err) return } if config.SysConfig.GetRemoteFileRangeWaitTime() != 0 { - time.Sleep(config.SysConfig.GetRemoteFileRangeWaitTime()) + timer := time.NewTimer(config.SysConfig.GetRemoteFileRangeWaitTime()) + select { + case <-timer.C: + case <-ctx.Done(): + timer.Stop() + return + } } } } @@ -317,25 +373,68 @@ func createCacheTask(taskNo int, start, end int64, taskParam *downloader.TaskPar cache.DingFile = taskParam.DingFile cache.TaskSize = taskParam.TaskSize cache.FileName = taskParam.FileName - cache.OrgRepo = taskParam.OrgRepo + cache.RepoKey = taskParam.RepoKey + cache.OrgRepo = taskParam.RepoKey.ID() cache.ResponseChan = taskParam.ResponseChan + cache.OnComplete = taskParam.OnCacheComplete return cache } +// Sync by file identity, without scheduling another download or resetting an +// existing process. StartPos equals the verified full prefix, so the existing +// monotonic SQL guard accepts both missing and partially reported progress. +func (d *DownloaderDao) reconcileCachedFile(p *downloader.TaskParam) { + entry := &manager.FileProcessEntry{ + DataType: p.DataType, Org: p.RepoKey.Namespace, Repo: p.RepoKey.Repo, + Name: p.FileName, Etag: p.Etag, FileSize: p.FileSize, + InstanceId: config.SysConfig.Registration().NodeID, + StartPos: p.FileSize, EndPos: p.FileSize, Status: consts.StatusDownloaded, + } + if err := d.schedulerDao.SyncFileProcess(&manager.SyncFileProcessReq{FileProcessEntries: []*manager.FileProcessEntry{entry}}); err != nil { + zap.S().Errorf("reconcile cached file %s/%s: %v", p.OrgRepo, p.FileName, err) + data.WriteLocalOperationChan(consts.OperationProcess, &data.FileProcessParam{ + Datatype: p.DataType, Org: p.RepoKey.Namespace, Repo: p.RepoKey.Repo, + Name: p.FileName, Etag: p.Etag, FileSize: p.FileSize, + StartPos: p.FileSize, EndPos: p.FileSize, Status: consts.StatusDownloaded, + }) + } +} + func createRemoteTask(taskNo int, start, end int64, taskParam *downloader.TaskParam) *downloader.RemoteFileTask { remote := downloader.NewRemoteFileTask(taskNo, start, end) + remote.Source = taskParam.Source + remote.LocalOnly = taskParam.LocalOnly remote.Context = taskParam.Context remote.DingFile = taskParam.DingFile remote.Authorization = taskParam.Authorization remote.Domain = taskParam.Domain remote.Uri = taskParam.Uri + remote.UpstreamURI = taskParam.Uri + remote.Peer = taskParam.Peer + if taskParam.Peer { + remote.Uri = peerFileURI(taskParam) + } remote.Queue = make(chan []byte, getQueueSize(remote.RangeStartPos, remote.RangeEndPos)) remote.ResponseChan = taskParam.ResponseChan remote.TaskSize = taskParam.TaskSize remote.FileName = taskParam.FileName - remote.OrgRepo = taskParam.OrgRepo + remote.RepoKey = taskParam.RepoKey + remote.OrgRepo = taskParam.RepoKey.ID() remote.DataType = taskParam.DataType remote.Etag = taskParam.Etag remote.Cancel = taskParam.Cancel return remote } + +func peerFileURI(p *downloader.TaskParam) string { + if p.RepoKey.Namespace == repository.HuggingFace { + prefix := "" + if p.RepoKey.RepoType != "models" { + prefix = "/" + p.RepoKey.RepoType + } + return prefix + "/" + repository.EscapeURLPath(p.RepoKey.Repo) + "/resolve/" + url.PathEscape(p.Revision) + "/" + repository.EscapeURLPath(p.FileName) + } + + q := url.Values{"repo": {p.RepoKey.Repo}, "revision": {p.Revision}, "path": {p.FileName}} + return "/api/repositories/" + url.PathEscape(p.RepoKey.RepoType) + "/" + url.PathEscape(p.RepoKey.Namespace) + "/file?" + q.Encode() +} diff --git a/internal/dao/file_dao.go b/internal/dao/file_dao.go index 65e5552..b252db7 100644 --- a/internal/dao/file_dao.go +++ b/internal/dao/file_dao.go @@ -32,7 +32,9 @@ import ( "dingospeed/pkg/common" "dingospeed/pkg/config" "dingospeed/pkg/consts" + "dingospeed/pkg/dependency" myerr "dingospeed/pkg/error" + "dingospeed/pkg/repository" "dingospeed/pkg/util" "github.com/bytedance/sonic" @@ -72,8 +74,8 @@ func (f *FileDao) CheckCommitHf(repoType, orgRepo, commit, authorization string) } func (f *FileDao) GetFileCommitSha(repoType, orgRepo, commit, authorization string, source string) (string, error) { - metaShaKey := GetMetaShaRepoKey(orgRepo, commit, authorization) - if v, ok := f.baseData.Cache.Get(metaShaKey); ok { + metaShaKey := GetMetaShaRepoKey(repoType+"/"+orgRepo, commit, authorization) + if v, ok := f.baseData.Cache.Get(metaShaKey); ok && (IsLocalOrgRepo(orgRepo) || !config.SysConfig.Online()) { return v.(string), nil } var ( @@ -83,10 +85,8 @@ func (f *FileDao) GetFileCommitSha(repoType, orgRepo, commit, authorization stri if IsLocalOrgRepo(orgRepo) { commitSha, err = f.GetCommitHfOffline(repoType, orgRepo, commit) if err != nil { - var accessErr *util.FileAccessError - if errors.As(err, &accessErr) { - zap.S().Errorw("local metadata storage unavailable", "kind", accessErr.Kind, "path", accessErr.Path, "error", accessErr.Err) - return "", myerr.Wrap("local metadata storage unavailable", err) + if !errors.Is(err, os.ErrNotExist) { + return "", localMetadataError(err) } return "", myerr.NewAppendCode(http.StatusNotFound, fmt.Sprintf("%s is not found", orgRepo)) } @@ -97,20 +97,18 @@ func (f *FileDao) GetFileCommitSha(repoType, orgRepo, commit, authorization stri } commitSha, err = f.GetCommitHfOffline(repoType, orgRepo, commit) if err != nil { - var accessErr *util.FileAccessError - if errors.As(err, &accessErr) { - zap.S().Errorw("metadata cache storage unavailable", "kind", accessErr.Kind, "path", accessErr.Path, "error", accessErr.Err) - return "", myerr.Wrap("metadata cache storage unavailable", err) - } if source == "file" { // 若只是发起文件下载（先在线后离线），将不会校验meta文件是否存在，没有就创建，主要是看文件本身是否存在。 goto remoteRequestMeta } + if !errors.Is(err, os.ErrNotExist) { + return "", localMetadataError(err) + } zap.S().Warnf("getFileCommitSha GetCommitHfOffline err.%v", err) return "", myerr.NewAppendCode(http.StatusNotFound, fmt.Sprintf("%s is not found", orgRepo)) } f.baseData.Cache.Set(metaShaKey, commitSha, config.SysConfig.GetDefaultExpiration()) - f.baseData.Cache.Set(GetMetaShaRepoKey(orgRepo, commitSha, authorization), commitSha, config.SysConfig.GetDefaultExpiration()) + f.baseData.Cache.Set(GetMetaShaRepoKey(repoType+"/"+orgRepo, commitSha, authorization), commitSha, config.SysConfig.GetDefaultExpiration()) return commitSha, nil remoteRequestMeta: @@ -131,7 +129,7 @@ remoteRequestMeta: commitSha = sha } f.baseData.Cache.Set(metaShaKey, commitSha, config.SysConfig.GetDefaultExpiration()) - f.baseData.Cache.Set(GetMetaShaRepoKey(orgRepo, commitSha, authorization), commitSha, config.SysConfig.GetDefaultExpiration()) + f.baseData.Cache.Set(GetMetaShaRepoKey(repoType+"/"+orgRepo, commitSha, authorization), commitSha, config.SysConfig.GetDefaultExpiration()) return commitSha, nil } @@ -151,25 +149,35 @@ func (f *FileDao) getCommitHfRemote(repoType, orgRepo, commit, authorization str zap.S().Errorf("unmarshal content:%s, error:%v", string(resp.Body), err) return http.StatusInternalServerError, "", err } + if err := repository.Segment(sha.Sha); err != nil { + return 502, "", err + } return resp.StatusCode, sha.Sha, nil } func (f *FileDao) RemoteRequestMeta(method, repoType, orgRepo, revision, authorization string) (*common.Response, error) { + return f.RemoteRequestMetaContext(context.Background(), method, repoType, orgRepo, revision, authorization) +} +func (f *FileDao) RemoteRequestMetaContext(ctx context.Context, method, repoType, orgRepo, revision, authorization string) (*common.Response, error) { + upstream, err := UpstreamRepo(repoType, orgRepo) + if err != nil { + return nil, err + } var reqUri string if revision == "" { - reqUri = fmt.Sprintf("/api/%s/%s", repoType, orgRepo) + reqUri = fmt.Sprintf("/api/%s/%s", repoType, repository.EscapeURLPath(upstream)) } else { - reqUri = fmt.Sprintf("/api/%s/%s/revision/%s", repoType, orgRepo, revision) + reqUri = fmt.Sprintf("/api/%s/%s/revision/%s", repoType, repository.EscapeURLPath(upstream), repository.EscapeURLPath(revision)) } headers := map[string]string{} if authorization != "" { headers["authorization"] = authorization } - return util.RetryRequest(func() (*common.Response, error) { + return util.RetryRequestContext(ctx, func() (*common.Response, error) { if method == consts.RequestTypeHead { - return util.Head(reqUri, headers) + return util.HeadContext(ctx, reqUri, headers) } else if method == consts.RequestTypeGet { - return util.Get(reqUri, headers) + return util.GetContext(ctx, reqUri, headers) } else { return nil, fmt.Errorf("request method err") } @@ -177,35 +185,51 @@ func (f *FileDao) RemoteRequestMeta(method, repoType, orgRepo, revision, authori } func (f *FileDao) GetCommitHfOffline(repoType, orgRepo, commit string) (string, error) { - apiPath := fmt.Sprintf("%s/api/%s/%s/revision/%s/meta_get.json", config.SysConfig.Repos(), repoType, orgRepo, commit) + apiPath := filepath.Join(RepositoryKey(repoType, orgRepo).Revision(config.SysConfig.Repos(), commit), "meta_get.json") exists, err := util.PathExists(apiPath) if err != nil { + util.ObserveFileAccessFailure("stat", err) return "", err } if exists { cacheContent, err := f.ReadCacheRequest(apiPath) if err != nil { - if accessErr, ok := util.ClassifyFileAccessError(apiPath, err); ok { - return "", accessErr - } + util.ObserveFileAccessFailure("read", err) return "", err } var sha CommitHfSha + dependency.Default.Observe(dependency.MetadataRead, true, time.Now()) if err = sonic.Unmarshal(cacheContent.OriginContent, &sha); err != nil { zap.S().Errorf("unmarshal error.%v", err) return "", myerr.Wrap("Unmarshal err", err) } return sha.Sha, nil } - return "", myerr.New(fmt.Sprintf("apiPath file not exist, %s", apiPath)) + return "", fmt.Errorf("apiPath file not exist, %s: %w", apiPath, os.ErrNotExist) +} + +// Keep the cause for internal diagnosis, without exposing paths in HTTP errors. +func localMetadataError(err error) error { + if _, ok := util.ClassifyFileAccessError("", err); ok { + return myerr.WrapCode(http.StatusServiceUnavailable, "local metadata storage unavailable", err) + } + return myerr.WrapCode(http.StatusInternalServerError, "local metadata invalid", err) } func (f *FileDao) FileGetGenerator(c echo.Context, repoType, orgRepo, commit, fileName, method string) error { + upstream := orgRepo + if !IsLocalOrgRepo(orgRepo) { + var err error + upstream, err = UpstreamRepo(repoType, orgRepo) + if err != nil { + return err + } + } var hfUri string if repoType == "models" { - hfUri = fmt.Sprintf("/%s/resolve/%s/%s", orgRepo, commit, fileName) + hfUri = fmt.Sprintf("/%s/resolve/%s/%s", repository.EscapeURLPath(upstream), repository.EscapeURLPath(commit), repository.EscapeURLPath(fileName)) } else { - hfUri = fmt.Sprintf("/%s/%s/resolve/%s/%s", repoType, orgRepo, commit, fileName) + hfUri = fmt.Sprintf("/%s/%s/resolve/%s/%s", repoType, repository.EscapeURLPath(upstream), repository.EscapeURLPath(commit), repository.EscapeURLPath(fileName)) } authorization := c.Request().Header.Get("Authorization") // _file_realtime_stream @@ -226,23 +250,79 @@ func (f *FileDao) FileGetGenerator(c echo.Context, repoType, orgRepo, commit, fi zap.S().Warnf("repo:%s, commit:%s, fileName:%s is directory", orgRepo, commit, fileName) return util.ErrorEntryNotFound(c) } - respHeaders, etag, startPos, endPos := constructRespHeader(c, pathInfo, commit, fileName) + return f.ServePreparedFile(c, RepositoryKey(repoType, orgRepo), commit, fileName, method, pathInfo, hfUri, nil) +} + +// ServePreparedFile serves normalized provider metadata through the existing cache pipeline. +func (f *FileDao) ServePreparedFile(c echo.Context, key repository.RepoKey, commit, fileName, method string, pathInfo *common.PathsInfo, uri string, source *downloader.RemoteSource) error { + var err error + repoType, orgRepo := key.RepoType, key.ID() + authorization := c.Request().Header.Get("Authorization") + var respHeaders map[string]string + var etag string + var startPos, endPos int64 + if source == nil { + respHeaders, etag, startPos, endPos = constructRespHeader(c, pathInfo, commit, fileName) + } else { + var err error + startPos, endPos, err = util.FileRange(c.Request().Header.Get("Range"), pathInfo.Size) + if err != nil { + c.Response().Header().Set("Content-Range", fmt.Sprintf("bytes */%d", pathInfo.Size)) + return echo.NewHTTPError(416, err.Error()) + } + etag = pathInfo.Oid + if pathInfo.Lfs.Oid != "" { + etag = pathInfo.Lfs.Oid + } + respHeaders = map[string]string{"Content-Type": "application/octet-stream", "Content-Length": fmt.Sprint(endPos - startPos), "Accept-Ranges": "bytes", "ETag": fmt.Sprintf("%q", etag), "X-Repo-Commit": commit} + if c.Request().Header.Get("Range") != "" { + respHeaders["Content-Range"] = fmt.Sprintf("bytes %d-%d/%d", startPos, endPos-1, pathInfo.Size) + } + } + if err := repository.Segment(etag); err != nil { + return echo.NewHTTPError(502, "invalid upstream content identity") + } + if !IsLocalOrgRepo(orgRepo) { + if err := repository.Register(config.SysConfig.Repos(), repository.Remote(RepositoryKey(repoType, orgRepo))); err != nil { + return err + } + } blobsFile := BlobPath(repoType, orgRepo, etag) filesPath := ResolvePath(repoType, orgRepo, commit, fileName) + if err = repository.SafePath(config.SysConfig.Repos(), blobsFile); err != nil { + return err + } + if err = repository.SafePath(config.SysConfig.Repos(), filepath.Dir(filesPath)); err != nil { + return err + } + if IsLocalOrgRepo(orgRepo) { + if complete, _, e := inspectCompleteBlob(blobsFile, pathInfo.Size); e != nil || !complete { + return echo.NewHTTPError(503, "hosted content unavailable") + } + } if err = f.ConstructBlobsAndFileFile(blobsFile, filesPath); err != nil { return util.ErrorProxyError(c) } if method == consts.RequestTypeHead { + if source != nil { + for k, v := range respHeaders { + c.Response().Header().Set(k, v) + } + return c.NoContent(http.StatusOK) + } return util.ResponseHeaders(c, http.StatusOK, respHeaders) } else if method == consts.RequestTypeGet { taskParam := &downloader.TaskParam{ TaskNo: 0, + Revision: commit, BlobsFile: blobsFile, FileName: fileName, FileSize: pathInfo.Size, OrgRepo: orgRepo, Authorization: authorization, - Uri: hfUri, + Uri: uri, + Source: source, + LocalOnly: key.Namespace == repository.ModelScope, DataType: repoType, Etag: etag, } @@ -314,7 +394,7 @@ func GetAnalysisFilePosition(dingFile *downloader.DingCache, startPos, endPos in // 以及 Windows 上与 filepath 系函数（MakeDirs、Abs/Rel 等）的分隔符不一致。 // filepath.Join 会做 Clean，两种情况都被吸收掉。 func BlobPath(repoType, orgRepo, etag string) string { - return filepath.Join(config.SysConfig.Repos(), "files", repoType, orgRepo, "blobs", etag) + return RepositoryKey(repoType, orgRepo).Blob(config.SysConfig.Repos(), etag) } // CopyLocalBlob streams the payload portion of one complete DingCache blob and @@ -359,10 +439,16 @@ func (f *FileDao) CopyLocalBlob(repoType, orgRepo, sha string, size int64, w io. // ResolvePath 是某个快照下仓库内路径对应的软链位置，与 BlobPath 同理只此一处。 func ResolvePath(repoType, orgRepo, commit, fileName string) string { - return filepath.Join(config.SysConfig.Repos(), "files", repoType, orgRepo, "resolve", commit, filepath.FromSlash(fileName)) + return RepositoryKey(repoType, orgRepo).Resolve(config.SysConfig.Repos(), commit, fileName) } +var fileReferenceLocks = newKeyedMutex() + func (f *FileDao) ConstructBlobsAndFileFile(blobsFile, filesPath string) (err error) { + // Reference creation may race across requests (and across FileDao instances). + // Only serialize this short filesystem operation, never the download itself. + fileReferenceLocks.Lock(filesPath) + defer fileReferenceLocks.Unlock(filesPath) if err = util.MakeDirs(blobsFile); err != nil { zap.S().Errorf("create %s dir err.%v", blobsFile, err) return err @@ -377,6 +463,13 @@ func (f *FileDao) ConstructBlobsAndFileFile(blobsFile, filesPath string) (err er return err } else { if !b { + // Windows may use the existing hard-link fallback instead of a symlink. + // It is already a valid reference, not a legacy file to delete/recreate. + srcInfo, srcErr := os.Stat(blobsFile) + dstInfo, dstErr := os.Stat(filesPath) + if srcErr == nil && dstErr == nil && os.SameFile(srcInfo, dstInfo) { + return nil + } zap.S().Infof("old data transfer, from %s to %s", filesPath, blobsFile) if blobFileExist := util.FileExists(blobsFile); blobFileExist { if err = util.DeleteFile(filesPath); err != nil { @@ -445,11 +538,14 @@ func (f *FileDao) InvalidateLocalManifest(repoType, orgRepo, commit string) { } func (f *FileDao) GetPathsInfo(hfUri, repoType, orgRepo, commit, authorization string, pathFileName string) (*common.PathsInfo, error) { + return f.GetPathsInfoContext(context.Background(), hfUri, repoType, orgRepo, commit, authorization, pathFileName) +} +func (f *FileDao) GetPathsInfoContext(ctx context.Context, hfUri, repoType, orgRepo, commit, authorization string, pathFileName string) (*common.PathsInfo, error) { var pathInfo *common.PathsInfo if pathFileName == "" { return nil, fmt.Errorf("pathFileName is null, %s/%s", orgRepo, commit) } - apiPathInfoPath := fmt.Sprintf("%s/api/%s/%s/paths-info/%s/%s/paths-info_post.json", config.SysConfig.Repos(), repoType, orgRepo, commit, pathFileName) + apiPathInfoPath := filepath.Join(RepositoryKey(repoType, orgRepo).PathsInfo(config.SysConfig.Repos(), commit, pathFileName), "paths-info_post.json") if IsLocalOrgRepo(orgRepo) { // 本地仓库的文件信息从该快照的清单派生，不依赖逐文件落盘的 paths-info 缓存。 manifest, err := f.ReadLocalManifest(repoType, orgRepo, commit) @@ -499,8 +595,12 @@ func (f *FileDao) GetPathsInfo(hfUri, repoType, orgRepo, commit, authorization s goto requestRemoteFileInfo requestRemoteFileInfo: - pathsInfoUri := fmt.Sprintf("/api/%s/%s/paths-info/%s", repoType, orgRepo, commit) - if response, err := f.requestFilePathInfo(pathsInfoUri, authorization, []string{pathFileName}); err != nil { + upstream, upstreamErr := UpstreamRepo(repoType, orgRepo) + if upstreamErr != nil { + return nil, upstreamErr + } + pathsInfoUri := fmt.Sprintf("/api/%s/%s/paths-info/%s", repoType, repository.EscapeURLPath(upstream), repository.EscapeURLPath(commit)) + if response, err := f.requestFilePathInfoContext(ctx, pathsInfoUri, authorization, []string{pathFileName}); err != nil { return nil, err } else { if !granted { @@ -517,7 +617,7 @@ requestRemoteFileInfo: return nil, myerr.NewAppendCode(http.StatusNotFound, "remoteRespPathsInfos is null") } if pathInfo.Size > consts.MAX_HTTP_DOWNLOAD_SIZE { - if resolveResp, err := f.requestFileResolve(hfUri, authorization); err != nil { + if resolveResp, err := f.requestFileResolveContext(ctx, hfUri, authorization); err != nil { return nil, err } else { pathInfo.XXetHash = resolveResp.GetKey(consts.HUGGINGFACE_HEADER_X_XET_HASH) @@ -538,16 +638,19 @@ requestRemoteFileInfo: } func (f *FileDao) requestFileResolve(fileResolveUri, authorization string) (*common.Response, error) { + return f.requestFileResolveContext(context.Background(), fileResolveUri, authorization) +} +func (f *FileDao) requestFileResolveContext(ctx context.Context, fileResolveUri, authorization string) (*common.Response, error) { headers := map[string]string{} if authorization != "" { headers["authorization"] = authorization } - response, err := util.RetryRequest(func() (*common.Response, error) { - return util.Head(fileResolveUri, headers) + response, err := util.RetryRequestContext(ctx, func() (*common.Response, error) { + return util.HeadContext(ctx, fileResolveUri, headers) }) if err != nil { zap.S().Errorf("req %s err.%v", fileResolveUri, err) - return nil, myerr.NewAppendCode(http.StatusInternalServerError, fmt.Sprintf("%v", err)) + return nil, fmt.Errorf("remote metadata request: %w", err) } // 非成功或重定向 if response.StatusCode != http.StatusOK && !(response.StatusCode >= http.StatusMultipleChoices && response.StatusCode <= http.StatusPermanentRedirect) { @@ -564,6 +667,9 @@ func (f *FileDao) requestFileResolve(fileResolveUri, authorization string) (*com } func (f *FileDao) requestFilePathInfo(pathsInfoUri, authorization string, filePaths []string) (*common.Response, error) { + return f.requestFilePathInfoContext(context.Background(), pathsInfoUri, authorization, filePaths) +} +func (f *FileDao) requestFilePathInfoContext(ctx context.Context, pathsInfoUri, authorization string, filePaths []string) (*common.Response, error) { reqData := map[string]interface{}{ "paths": filePaths, } @@ -575,11 +681,11 @@ func (f *FileDao) requestFilePathInfo(pathsInfoUri, authorization string, filePa if authorization != "" { headers["authorization"] = authorization } - if response, err := util.RetryRequest(func() (*common.Response, error) { - return util.Post(pathsInfoUri, "application/json", jsonData, headers) + if response, err := util.RetryRequestContext(ctx, func() (*common.Response, error) { + return util.PostContext(ctx, pathsInfoUri, "application/json", jsonData, headers) }); err != nil { zap.S().Errorf("req %s err.%v", pathsInfoUri, err) - return nil, myerr.NewAppendCode(http.StatusInternalServerError, fmt.Sprintf("%v", err)) + return nil, fmt.Errorf("remote metadata request: %w", err) } else if response.StatusCode != http.StatusOK { var errorResp common.ErrorResp if len(response.Body) > 0 { @@ -606,6 +712,14 @@ func (f *FileDao) FileChunkGet(c echo.Context, taskParam *downloader.TaskParam, if value := c.Request().Header.Get(consts.RequestSourceInner); value == "1" { isInnerRequest = true } + if taskParam.Source != nil || c.Request().Header.Get(common.LocalTransferHeader) != "" { + taskParam.CacheResult = make(chan error, 1) + } + ensureCache := taskParam.Source != nil && c.Request().Header.Get("X-Dingo-Ensure-Cache") == "1" + if ensureCache { + delete(respHeaders, "Content-Length") + respHeaders["Trailer"] = "X-Dingo-Cache-Complete, X-Dingo-Cache-Error" + } taskParam.Context = ctx taskParam.ResponseChan = responseChan taskParam.Cancel = cancel @@ -613,10 +727,29 @@ func (f *FileDao) FileChunkGet(c echo.Context, taskParam *downloader.TaskParam, if err := f.downloaderDao.FileDownload(startPos, endPos, isInnerRequest, taskParam); err != nil { return util.MultipleErrorProxyError(err, c) } - if err := util.ResponseStream(c, fileName, respHeaders, responseChan); err != nil { + streamErr := util.ResponseStream(c, fileName, respHeaders, responseChan) + if streamErr != nil { + cancel() + if taskParam.CacheResult != nil { + <-taskParam.CacheResult + } + err := streamErr zap.S().Errorf("FileChunkGet stream err.%v", err) return util.ErrorProxyTimeout(c) } + if taskParam.CacheResult != nil { + result := <-taskParam.CacheResult + if ensureCache { + c.Response().Header().Set("X-Dingo-Cache-Complete", fmt.Sprint(result == nil)) + if result != nil { + c.Response().Header().Set("X-Dingo-Cache-Error", util.CacheFailureReason(result)) + zap.S().Errorw("Could not complete file cache", "file", fileName, "error", result) + } + } + if result != nil && ensureCache { + return result + } + } return nil } diff --git a/internal/dao/file_dao_storage_error_test.go b/internal/dao/file_dao_storage_error_test.go index ecc6899..5cc7417 100644 --- a/internal/dao/file_dao_storage_error_test.go +++ b/internal/dao/file_dao_storage_error_test.go @@ -2,8 +2,14 @@ package dao import ( "errors" + "fmt" "net/http" + "net/http/httptest" + "os" + "path/filepath" "strings" + "sync/atomic" + "syscall" "testing" "dingospeed/internal/data" @@ -18,7 +24,7 @@ func TestGetCommitHfOfflinePreservesMissingBehavior(t *testing.T) { t.Cleanup(func() { config.SysConfig = oldConfig }) fileDao := NewFileDao(nil, data.NewBaseData(), nil) - _, err := fileDao.GetFileCommitSha("models", "org/repo", "main", "", "meta") + _, err := fileDao.GetFileCommitSha("models", "huggingface/org/repo", "main", "", "meta") var appErr myerr.Error if !errors.As(err, &appErr) || appErr.StatusCode() != http.StatusNotFound { t.Fatalf("missing metadata: got %v, want HTTP 404 application error", err) @@ -33,18 +39,120 @@ func TestGetFileCommitShaPreservesStorageFailure(t *testing.T) { t.Cleanup(func() { config.SysConfig = oldConfig }) fileDao := NewFileDao(nil, data.NewBaseData(), nil) - _, err := fileDao.GetFileCommitSha("models", "org/repo", "main", "", "meta") + _, err := fileDao.GetFileCommitSha("models", "huggingface/org/repo", "main", "", "meta") if err == nil { t.Fatal("expected inaccessible storage error") } - var accessErr *util.FileAccessError - if !errors.As(err, &accessErr) { - t.Fatalf("got %T %v, want wrapped FileAccessError", err, err) + var appErr myerr.Error + var access *util.FileAccessError + if !errors.As(err, &appErr) || appErr.StatusCode() != 503 || !errors.As(err, &access) { + t.Fatalf("storage failure lost: %v", err) + } +} + +func TestCommitObservationPreservesFallbackAndReadErrors(t *testing.T) { + var calls atomic.Int64 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { calls.Add(1); fmt.Fprint(w, `{"sha":"remote"}`) })) + defer server.Close() + old := config.SysConfig + defer func() { config.SysConfig = old }() + for _, tc := range []struct { + name, source, org string + remote bool + }{ + {"offline-meta", "meta", "huggingface/org/repo", false}, + {"offline-file-fallback", "file", "huggingface/org/repo", true}, + {"local-file-no-fallback", "file", "dingo-local/repo", false}, + } { + t.Run(tc.name, func(t *testing.T) { + config.SysConfig = &config.Config{Server: config.ServerConfig{Repos: "invalid\x00repos", HfScheme: "http", HfNetLoc: strings.TrimPrefix(server.URL, "http://")}, Retry: config.Retry{Attempts: 1}, Upload: config.Upload{Namespace: "dingo-local"}} + base := data.NewBaseData() + file := NewFileDao(nil, base, NewLockDao(base)) + before, observations := calls.Load(), util.FileAccessObservations() + sha, err := file.GetFileCommitSha("models", tc.org, "main", "", tc.source) + if tc.remote { + if err != nil || sha != "remote" || calls.Load()-before != 1 { + t.Fatalf("fallback changed: sha=%s err=%v calls=%d", sha, err, calls.Load()-before) + } + } else { + var app myerr.Error + if !errors.As(err, &app) || app.StatusCode() != 503 || calls.Load() != before { + t.Fatalf("local response changed: %v", err) + } + } + assertObservationDelta(t, observations, "stat") + }) + } + t.Run("stat-success-read-failure", func(t *testing.T) { + config.SysConfig = &config.Config{Server: config.ServerConfig{Repos: t.TempDir()}} + base := data.NewBaseData() + file := NewFileDao(nil, base, NewLockDao(base)) + path := filepath.Join(config.SysConfig.Repos(), "api/models/org/repo/revision/main/meta_get.json") + if err := os.MkdirAll(path, 0700); err != nil { + t.Fatal(err) + } + _, original := file.ReadCacheRequest(path) + before := util.FileAccessObservations() + _, got := file.GetCommitHfOffline("models", "huggingface/org/repo", "main") + var app myerr.Error + if !errors.As(got, &app) || got.Error() != original.Error() { + t.Fatalf("read error contract changed: %v / %v", got, original) + } + assertObservationDelta(t, before, "read") + }) +} + +func assertObservationDelta(t *testing.T, before []util.FileAccessObservation, operation string) { + t.Helper() + var delta uint64 + for i, value := range util.FileAccessObservations() { + if value.Operation == operation { + delta += value.Count - before[i].Count + } else if value.Count != before[i].Count { + t.Fatalf("unexpected operation: %+v", value) + } + } + if delta != 1 { + t.Fatalf("want one %s observation, got %d", operation, delta) } - if accessErr.Kind != util.FileAccessUnavailable { - t.Fatalf("got kind %q, want %q", accessErr.Kind, util.FileAccessUnavailable) +} + +func TestLocalMetadataErrorPreservesCauseAndHidesPath(t *testing.T) { + for _, cause := range []error{os.ErrPermission, syscall.Errno(5), syscall.Errno(107)} { + original := &os.PathError{Op: "read", Path: "private/cache/meta_get.json", Err: cause} + got := localMetadataError(myerr.Wrap("read failed", original)) + var app myerr.Error + if !errors.As(got, &app) || app.StatusCode() != 503 || !errors.Is(got, cause) { + t.Fatalf("access failure contract: %v", got) + } + if strings.Contains(got.Error(), original.Path) { + t.Fatal("private path exposed") + } } - if strings.Contains(err.Error(), "invalid\x00repos") { - t.Fatalf("public application error exposes storage path: %q", err.Error()) +} + +func TestHostedMetadataUsesLocalErrorsEvenWhenOnline(t *testing.T) { + old := config.SysConfig + defer func() { config.SysConfig = old }() + for _, online := range []bool{false, true} { + config.SysConfig = &config.Config{Server: config.ServerConfig{Repos: t.TempDir(), Online: online}, Upload: config.Upload{Namespace: "dingo-local"}} + base := data.NewBaseData() + lock := NewLockDao(base) + file := NewFileDao(nil, base, lock) + base.Cache.Set(GetMetaShaRepoKey("models/dingo-local/repo", "main", ""), "abc123", 0) + meta := NewMetaDao(file, lock, base) + _, err := meta.GetMetadata("models", "dingo-local/repo", "main", "get", "") + var app myerr.Error + if !errors.As(err, &app) || app.StatusCode() != 404 { + t.Fatalf("online=%v missing: %v", online, err) + } + p := filepath.Join(config.SysConfig.Repos(), "api/models/dingo-local/repo/revision/abc123/meta_get.json") + if err := os.MkdirAll(p, 0700); err != nil { + t.Fatal(err) + } + _, err = meta.GetMetadata("models", "dingo-local/repo", "main", "get", "") + if !errors.As(err, &app) || app.StatusCode() != 503 { + t.Fatalf("online=%v read failure: %v", online, err) + } } } diff --git a/internal/dao/file_reference_concurrent_test.go b/internal/dao/file_reference_concurrent_test.go new file mode 100644 index 0000000..f4c1394 --- /dev/null +++ b/internal/dao/file_reference_concurrent_test.go @@ -0,0 +1,48 @@ +package dao + +import ( + "bytes" + "os" + "path/filepath" + "sync" + "testing" +) + +func TestPreparedHardlinkRemainsAReference(t *testing.T) { + root := t.TempDir() + blob := filepath.Join(root, "blob") + ref := filepath.Join(root, "ref") + content := []byte("existing container bytes") + if err := os.WriteFile(blob, content, 0600); err != nil { + t.Fatal(err) + } + if err := os.Link(blob, ref); err != nil { + t.Skipf("hard links unsupported: %v", err) + } + original, err := os.Stat(ref) + if err != nil { + t.Fatal(err) + } + var wg sync.WaitGroup + for i := 0; i < 16; i++ { + wg.Add(1) + go func() { + defer wg.Done() + if err := (&FileDao{}).ConstructBlobsAndFileFile(blob, ref); err != nil { + t.Error(err) + } + }() + } + wg.Wait() + after, err := os.Stat(ref) + if err != nil { + t.Fatal(err) + } + if !os.SameFile(original, after) { + t.Fatal("existing hard link was recreated") + } + b, err := os.ReadFile(blob) + if err != nil || !bytes.Equal(b, content) { + t.Fatal("reference preparation modified blob", err) + } +} diff --git a/internal/dao/heartbeat_health.go b/internal/dao/heartbeat_health.go new file mode 100644 index 0000000..a55261d --- /dev/null +++ b/internal/dao/heartbeat_health.go @@ -0,0 +1,57 @@ +package dao + +import ( + "time" + + "dingospeed/pkg/config" + "dingospeed/pkg/dependency" + "dingospeed/pkg/proto/manager" + "dingospeed/pkg/util" + "github.com/google/uuid" +) + +var healthProcessID = uuid.NewString() +var healthStartedAt = time.Now().Unix() + +// Only read existing in-memory observations. Never probe storage in a heartbeat. +func heartbeatHealth(now time.Time) *manager.NodeHealthSnapshot { + period := config.SysConfig.Scheduler.Discovery.HeartbeatPeriod + if period < 1 { + period = 1 + } + if period > 86400 { + period = 86400 + } + snapshot := &manager.NodeHealthSnapshot{ + Version: 1, ProcessId: healthProcessID, StartedAt: healthStartedAt, + CollectedAt: max(now.Unix(), healthStartedAt), HeartbeatPeriodSeconds: uint32(period), + } + for _, capability := range []struct { + id dependency.ID + name string + }{ + {dependency.MetadataRead, "metadata_read"}, {dependency.MetadataWrite, "metadata_write"}, + } { + observation := dependency.Default.Snapshot(capability.id, now) + // Business observations can arrive after the caller captured now. Do not + // produce a report whose observation timestamps exceed its collection time. + snapshot.CollectedAt = max(snapshot.CollectedAt, healthUnix(observation.LastObservation), healthUnix(observation.LastFailure)) + snapshot.Capabilities = append(snapshot.Capabilities, &manager.CapabilityObservation{ + Capability: capability.name, State: int32(observation.State), Unresolved: observation.Unresolved, + LastObservation: healthUnix(observation.LastObservation), LastFailure: healthUnix(observation.LastFailure), + }) + } + for _, observation := range util.FileAccessObservations() { + snapshot.Errors = append(snapshot.Errors, &manager.StorageErrorCount{ + Operation: observation.Operation, Kind: string(observation.Kind), Count: observation.Count, + }) + } + return snapshot +} + +func healthUnix(t time.Time) int64 { + if t.IsZero() { + return 0 + } + return t.Unix() +} diff --git a/internal/dao/heartbeat_health_test.go b/internal/dao/heartbeat_health_test.go new file mode 100644 index 0000000..9bb2c9c --- /dev/null +++ b/internal/dao/heartbeat_health_test.go @@ -0,0 +1,82 @@ +package dao + +import ( + "context" + "errors" + "os" + "testing" + "time" + + "dingospeed/pkg/config" + "dingospeed/pkg/dependency" + pb "dingospeed/pkg/proto/manager" + "dingospeed/pkg/util" + "google.golang.org/grpc" + "google.golang.org/protobuf/proto" + "google.golang.org/protobuf/types/known/emptypb" +) + +type healthClient struct { + pb.ManagerClient + received *pb.HeartbeatRequest + err error +} + +func (c *healthClient) Heartbeat(_ context.Context, r *pb.HeartbeatRequest, _ ...grpc.CallOption) (*emptypb.Empty, error) { + encoded, err := proto.Marshal(r) + if err != nil { + return nil, err + } + c.received = &pb.HeartbeatRequest{} + if err = proto.Unmarshal(encoded, c.received); err != nil { + return nil, err + } + return &emptypb.Empty{}, c.err +} + +func TestHeartbeatCarriesMemoryHealthWithoutStorage(t *testing.T) { + oldConfig, oldMonitor := config.SysConfig, dependency.Default + t.Cleanup(func() { config.SysConfig = oldConfig; dependency.Default = oldMonitor }) + config.SysConfig = &config.Config{} + config.SysConfig.Id = 7 + config.SysConfig.Scheduler.Discovery.HeartbeatPeriod = 5 + config.SysConfig.Scheduler.Discovery.InstanceId = "hd-test" + config.SysConfig.Server.Repos = "\x00-inaccessible" + dependency.Default = &dependency.Monitor{} + client := &healthClient{err: errors.New("communication unavailable")} + dao := NewSchedulerDao() + dao.Client = client + util.ObserveFileAccessFailure("read", &os.PathError{Op: "read", Path: "unavailable-cache", Err: os.ErrPermission}) + if err := dao.Heartbeat(); err != client.err { + t.Fatal("heartbeat error semantics changed") + } + client.err = nil + if err := dao.Heartbeat(); err != nil { + t.Fatal(err) + } + h := client.received.Health + if client.received.Id != 7 || h == nil || h.Version != 1 || !h.Capabilities[0].Unresolved || h.Capabilities[1].State != 0 { + t.Fatalf("missing health: %v", client.received) + } + found := false + for _, e := range h.Errors { + if e.Operation == "read" && e.Kind == "permission_denied" && e.Count > 0 { + found = true + } + } + if !found { + t.Fatal("classified counter not transported") + } + if next := heartbeatHealth(time.Now()); next.ProcessId != h.ProcessId || next.StartedAt != h.StartedAt { + t.Fatal("process identity changes every heartbeat") + } + if len(h.Errors) > 16 || len(h.Capabilities) != 2 { + t.Fatal("unbounded report") + } + before := time.Now() + dependency.Default.Observe(dependency.MetadataRead, false, before.Add(time.Second)) + next := heartbeatHealth(before) + if next.CollectedAt < next.Capabilities[0].LastObservation { + t.Fatal("concurrent observation produced invalid report timestamp") + } +} diff --git a/internal/dao/inventory_journal.go b/internal/dao/inventory_journal.go new file mode 100644 index 0000000..1199c50 --- /dev/null +++ b/internal/dao/inventory_journal.go @@ -0,0 +1,260 @@ +package dao + +import ( + "dingospeed/pkg/config" + "dingospeed/pkg/inventory" + "dingospeed/pkg/repository" + "encoding/json" + "fmt" + "os" + "path/filepath" + "time" +) + +type metadataIntent struct { + Revision string + Commit string + Files []LocalManifestFile +} +type snapshotEdit struct { + OldCommit string + Tags []string + Files []LocalManifestFile +} +type deletionIntent struct { + Edits []snapshotEdit + Recycle []RecycleEntry +} + +func inventoryKey(k repository.RepoKey) inventory.Key { + return inventory.Key{Namespace: k.Namespace, RepoType: k.RepoType, Repo: k.Repo} +} + +// Inventory reads bypass serving caches, including during explicit disk audits. +func ReadInventorySnapshot(k repository.RepoKey, revision string) (string, []LocalManifestFile, error) { + b, err := readCacheContent(filepath.Join(k.Revision(config.SysConfig.Repos(), revision), "meta_get.json")) + if err != nil { + return "", nil, err + } + var meta struct { + Sha string `json:"sha"` + } + if err = json.Unmarshal(b, &meta); err != nil { + return "", nil, err + } + if len(meta.Sha) != 64 { + return "", nil, fmt.Errorf("invalid revision pointer") + } + f, err := readManifestFile(LocalManifestPath(k.RepoType, k.ID(), meta.Sha)) + return meta.Sha, f, err +} + +func (u *UploadDao) writeEffectiveMetadata(repoType, id, revision, commit string, files []LocalManifestFile) error { + k := RepositoryKey(repoType, id) + if err := u.recoverInventoryLocked(k); err != nil { + return err + } + if err := RegisterHosted(k.RepoType, k.Namespace, k.Repo); err != nil { + return err + } + p := metadataIntent{revision, commit, files} + return inventory.Run(config.SysConfig.Repos(), inventoryKey(k), "metadata", p, func() error { return u.writeEffectiveMetadataRaw(repoType, id, revision, commit, files) }) +} + +// Recovery replays only journaled repositories, under the same locks as normal writes. +func (u *UploadDao) RecoverInventory(k repository.RepoKey) error { + if err := k.Validate(); err != nil { + return err + } + lock := uploadRepoLockKey(k.RepoType, k.ID()) + repositoryLifecycle.Lock(lock) + defer repositoryLifecycle.Unlock(lock) + uploadRepoLocks.Lock(lock) + defer uploadRepoLocks.Unlock(lock) + return u.recoverInventoryLocked(k) +} +func (u *UploadDao) recoverInventoryLocked(k repository.RepoKey) error { + root := config.SysConfig.Repos() + s, err := inventory.Read(root) + if err != nil { + return err + } + e := s.Entries[inventoryKey(k).ID()] + if e == nil || e.Operation == nil { + return nil + } + inventory.Barrier.RLock() + defer inventory.Barrier.RUnlock() + switch e.Operation.Kind { + case "metadata": + var p metadataIntent + if err = json.Unmarshal(e.Operation.Data, &p); err != nil { + return err + } + if err = verifyManifestContent(k.RepoType, k.ID(), p.Files); err != nil { + return err + } + err = u.writeEffectiveMetadataRaw(k.RepoType, k.ID(), p.Revision, p.Commit, p.Files) + case "delete-files": + var p deletionIntent + if err = json.Unmarshal(e.Operation.Data, &p); err != nil { + return err + } + err = NewCacheAdminDao(u.fileDao).applyInventoryDeletion(k, p) + case "delete-repository": + err = u.removeEmptyRepository(k) + default: + return fmt.Errorf("unknown inventory operation %q", e.Operation.Kind) + } + if err != nil { + return err + } + return inventory.Finish(root, inventoryKey(k), e.Operation.Kind) +} + +// Capture desired manifests before any deletion; replay never depends on the old +// manifests still existing after a crash halfway through a multi-version delete. +func (d *CacheAdminDao) journalSoftDelete(idx *repoIndex, items []DeleteItem) []*DeleteResult { + k := RepositoryKey(idx.RepoType, idx.OrgRepo) + u := NewUploadDao(d.fileDao, nil) + if err := u.recoverInventoryLocked(k); err != nil { + return failAll(items, err.Error()) + } + idx = buildRepoIndex(idx.RepoType, idx.OrgRepo) + wanted := map[string]bool{} + results := make([]*DeleteResult, 0, len(items)) + p := deletionIntent{} + for _, item := range items { + r := &DeleteResult{DeleteItem: item, Status: "skipped", Reason: "no reference found"} + if item.Path == "" { + r.Status = "failed" + r.Reason = "path is required" + } else { + for _, ref := range idx.BySha[item.Sha] { + if ref.Path == item.Path { + r.Status = "deleted" + r.Reason = "" + wanted[item.Path+"\x00"+item.Sha] = true + } + } + } + results = append(results, r) + } + commits := map[string]bool{} + for _, refs := range idx.BySha { + for _, ref := range refs { + commits[ref.Commit] = true + } + } + for commit := range commits { + files, err := readManifestFile(LocalManifestPath(idx.RepoType, idx.OrgRepo, commit)) + if err != nil { + return failAll(items, err.Error()) + } + kept := make([]LocalManifestFile, 0, len(files)) + for _, f := range files { + if !wanted[f.Path+"\x00"+f.Sha256] { + kept = append(kept, f) + } + } + if len(kept) != len(files) { + p.Edits = append(p.Edits, snapshotEdit{commit, idx.TagsOf[commit], kept}) + } + } + for sha, refs := range idx.BySha { + all := true + paths := []string{} + tags := []string{} + for _, ref := range refs { + if !wanted[ref.Path+"\x00"+sha] { + all = false + } + paths = appendUnique(paths, ref.Path) + for _, tag := range idx.TagsOf[ref.Commit] { + tags = appendUnique(tags, tag) + } + } + if all { + p.Recycle = append(p.Recycle, RecycleEntry{RepoType: k.RepoType, Namespace: k.Namespace, Repo: k.Repo, OrgRepo: k.ID(), Sha: sha, Size: idx.Blobs[sha].Size, Source: CacheSourceUpload, Paths: paths, Revisions: tags, UnlinkedAt: time.Now().Unix()}) + } + } + if len(p.Edits) == 0 { + return results + } + err := inventory.Run(config.SysConfig.Repos(), inventoryKey(k), "delete-files", p, func() error { return d.applyInventoryDeletion(k, p) }) + if err != nil { + return failAll(items, err.Error()) + } + return results +} +func (d *CacheAdminDao) applyInventoryDeletion(k repository.RepoKey, p deletionIntent) error { + idx := &repoIndex{RepoType: k.RepoType, OrgRepo: k.ID(), Source: CacheSourceUpload, TagsOf: map[string][]string{}, BySha: map[string][]repoRef{}} + for _, edit := range p.Edits { + idx.TagsOf[edit.OldCommit] = edit.Tags + } + // Write every target before deleting any old snapshot: a target can also be + // another edit's old commit. Never remove a manifest a final tag still needs. + finals := map[string]bool{} + u := NewUploadDao(d.fileDao, nil) + for _, edit := range p.Edits { + commit, err := manifestCommit(edit.Files) + if err != nil { + return err + } + if len(edit.Files) == 0 && len(edit.Tags) == 0 { + continue + } + finals[commit] = true + if err = u.writeEffectiveMetadataRaw(k.RepoType, k.ID(), commit, commit, edit.Files); err != nil { + return err + } + for _, tag := range edit.Tags { + if err = u.writeMeta(k.RepoType, k.ID(), tag, commit, edit.Files); err != nil { + return err + } + if err = inventory.SyncParents(config.SysConfig.Repos(), filepath.Join(k.Revision(config.SysConfig.Repos(), tag), "meta_get.json")); err != nil { + return err + } + } + } + for _, edit := range p.Edits { + if !finals[edit.OldCommit] { + if err := d.dropSnapshot(idx, edit.OldCommit); err != nil { + return err + } + } + } + for _, entry := range p.Recycle { + if err := writeRecycleEntry(entry); err != nil { + return err + } + if err := inventory.SyncParents(config.SysConfig.Repos(), recycleEntryPath(k.RepoType, k.ID(), entry.Sha)); err != nil { + return err + } + } + return nil +} + +func (u *UploadDao) removeEmptyRepository(k repository.RepoKey) error { + if err := k.Validate(); err != nil { + return err + } + root := config.SysConfig.Repos() + for _, p := range []string{k.FilesRoot(root), k.APIRoot(root)} { + if err := repository.SafePath(root, p); err != nil { + return err + } + } + // The intent is recorded only after live references and active writers have + // been excluded. Replays cannot race a new same-name repository. + if err := os.RemoveAll(k.FilesRoot(root)); err != nil { + return err + } + if err := os.RemoveAll(k.APIRoot(root)); err != nil { + return err + } + if err := inventory.SyncParents(root, k.FilesRoot(root)); err != nil { + return err + } + return inventory.SyncParents(root, k.APIRoot(root)) +} diff --git a/internal/dao/meta_dao.go b/internal/dao/meta_dao.go index ef4eec2..6119f14 100644 --- a/internal/dao/meta_dao.go +++ b/internal/dao/meta_dao.go @@ -15,16 +15,22 @@ package dao import ( + "context" "errors" "fmt" "net/http" + "os" "path/filepath" "strings" + "syscall" + "time" "dingospeed/internal/data" "dingospeed/pkg/common" "dingospeed/pkg/config" + "dingospeed/pkg/dependency" myerr "dingospeed/pkg/error" + "dingospeed/pkg/repository" "dingospeed/pkg/util" "github.com/bytedance/sonic" @@ -73,26 +79,24 @@ func (m *MetaDao) WhoamiV2Generator(c echo.Context) error { } func (m *MetaDao) ReposGenerator(c echo.Context) error { - reposPath := config.SysConfig.Repos() - - datasets, _ := filepath.Glob(filepath.Join(reposPath, "api/datasets/*/*")) - datasetsRepos := util.ProcessPaths(datasets) - - models, _ := filepath.Glob(filepath.Join(reposPath, "api/models/*/*")) - modelsRepos := util.ProcessPaths(models) - - spaces, _ := filepath.Glob(filepath.Join(reposPath, "api/spaces/*/*")) - spacesRepos := util.ProcessPaths(spaces) - - return c.Render(http.StatusOK, "repos.html", map[string]interface{}{ - "datasets_repos": datasetsRepos, - "models_repos": modelsRepos, - "spaces_repos": spacesRepos, - }) + all, err := repository.List(config.SysConfig.Repos()) + if err != nil { + return echo.NewHTTPError(500, "repository registry unavailable") + } + groups := map[string]interface{}{"datasets_repos": []string{}, "models_repos": []string{}, "spaces_repos": []string{}} + for _, d := range all { + name := d.RepoType + "_repos" + groups[name] = append(groups[name].([]string), d.ID()) + } + return c.Render(http.StatusOK, "repos.html", groups) } func (m *MetaDao) RepoRefs(repoType string, orgRepo string, authorization string) (*common.Response, error) { - refsUri := fmt.Sprintf("/api/%s/%s/refs", repoType, orgRepo) + upstream, err := UpstreamRepo(repoType, orgRepo) + if err != nil { + return nil, err + } + refsUri := fmt.Sprintf("/api/%s/%s/refs", repoType, repository.EscapeURLPath(upstream)) headers := map[string]string{} if authorization != "" { headers["authorization"] = authorization @@ -107,6 +111,53 @@ func (m *MetaDao) ForwardRefs(originalReq echo.Context) (*http.Response, error) return util.ForwardRequest(originalReq) } +// RefreshPreheatMetadata resolves a moving HF branch once for an explicit task. +// Ordinary catalog browsing never calls this method. +func (m *MetaDao) RefreshPreheatMetadata(k repository.RepoKey, revision, authorization string) (*CommitHfSha, error) { + return m.RefreshPreheatMetadataContext(context.Background(), k, revision, authorization) +} +func (m *MetaDao) RefreshPreheatMetadataContext(ctx context.Context, k repository.RepoKey, revision, authorization string) (*CommitHfSha, error) { + if err := k.Validate(); err != nil { + return nil, err + } + if k.Namespace != repository.HuggingFace || !config.SysConfig.Online() { + return nil, fmt.Errorf("HF prewarm requires an online Hugging Face source") + } + if err := repository.Segment(revision); err != nil { + return nil, err + } + lock := m.lockDao.getMetaDataReqLock(GetMetaDataReqKey(k.RepoType, k.ID(), revision)) + lock.Lock() + defer lock.Unlock() + resp, err := m.fileDao.RemoteRequestMetaContext(ctx, "get", k.RepoType, k.ID(), revision, authorization) + if err != nil { + return nil, err + } + if resp.StatusCode != http.StatusOK { + return nil, myerr.NewAppendCode(resp.StatusCode, "HF prewarm metadata request failed") + } + var meta CommitHfSha + if err := sonic.Unmarshal(resp.Body, &meta); err != nil { + return nil, err + } + if err := repository.Segment(meta.Sha); err != nil { + return nil, err + } + for _, file := range meta.Siblings { + if err := repository.Relative(file.Rfilename); err != nil { + return nil, err + } + } + headers := resp.ExtractHeaders(resp.Headers) + if err := m.writeApiMetaFile(k.RepoType, k.ID(), meta.Sha, "get", resp.StatusCode, headers, resp.Body); err != nil { + return nil, err + } + if err := m.writeApiMetaFile(k.RepoType, k.ID(), revision, "get", resp.StatusCode, headers, resp.Body); err != nil { + return nil, err + } + return &meta, nil +} + func (m *MetaDao) GetMetadata(repoType, orgRepo, revision, method, authorization string) (*common.CacheContent, error) { var ( cacheContent *common.CacheContent @@ -120,9 +171,9 @@ func (m *MetaDao) GetMetadata(repoType, orgRepo, revision, method, authorization if err != nil { return nil, err } - apiDir := fmt.Sprintf("%s/api/%s/%s/revision/%s", config.SysConfig.Repos(), repoType, orgRepo, commitSha) - apiMetaPath := fmt.Sprintf("%s/%s", apiDir, fmt.Sprintf("meta_%s.json", method)) - if config.SysConfig.Online() { + apiDir := RepositoryKey(repoType, orgRepo).Revision(config.SysConfig.Repos(), commitSha) + apiMetaPath := filepath.Join(apiDir, "meta_"+method+".json") + if config.SysConfig.Online() && !IsLocalOrgRepo(orgRepo) { if util.FileExists(apiMetaPath) { if cacheContent, err = m.fileDao.ReadCacheRequest(apiMetaPath); err != nil { zap.S().Errorf("ReadCacheRequest err.%v", err) @@ -136,11 +187,21 @@ func (m *MetaDao) GetMetadata(repoType, orgRepo, revision, method, authorization } } } else { - if util.FileExists(apiMetaPath) { + exists, accessErr := util.PathExists(apiMetaPath) + if accessErr != nil { + util.ObserveFileAccessFailure("stat", accessErr) + return nil, localMetadataError(accessErr) + } + if exists { if cacheContent, err = m.fileDao.ReadCacheRequest(apiMetaPath); err != nil { + util.ObserveFileAccessFailure("read", err) zap.S().Errorf("ReadCacheRequest err.%v", err) - return nil, err + if errors.Is(err, os.ErrNotExist) { + return nil, myerr.NewAppendCode(http.StatusNotFound, fmt.Sprintf("%s not exist", orgRepo)) + } + return nil, localMetadataError(err) } + dependency.Default.Observe(dependency.MetadataRead, true, time.Now()) } else { return nil, myerr.NewAppendCode(http.StatusNotFound, fmt.Sprintf("%s not exist", orgRepo)) } @@ -149,6 +210,19 @@ func (m *MetaDao) GetMetadata(repoType, orgRepo, revision, method, authorization if err != nil { return nil, err } + if !IsLocalOrgRepo(orgRepo) { + k, err := repository.ParseID(repoType, orgRepo) + if err != nil { + return nil, err + } + if err = repository.Register(config.SysConfig.Repos(), repository.Remote(k)); err != nil { + // Recording a remote cache descriptor is not required to deliver an + // online response. Validation/conflict errors still fail closed. + if !config.SysConfig.Online() || !metadataCacheIOError(err) { + return nil, err + } + } + } return cacheContent, nil } @@ -188,60 +262,64 @@ func (m *MetaDao) requestAndSaveMeta(repoType, orgRepo, revision, commitSha, met return nil, myerr.NewAppendCode(resp.StatusCode, "request err") } extractHeaders := resp.ExtractHeaders(resp.Headers) + content := &common.CacheContent{StatusCode: resp.StatusCode, Headers: extractHeaders, OriginContent: resp.Body} + // This helper is used only for online metadata delivery. Keep failures in + // writeApiMetaFile visible to durable writers; tolerate filesystem cache + // failures here only after a complete successful upstream response exists. + cacheResult := func(err error) (*common.CacheContent, error) { + if metadataCacheIOError(err) { + return content, nil + } + return nil, err + } mainVersion := "main" if revision == mainVersion { err = m.writeApiMetaFile(repoType, orgRepo, revision, method, resp.StatusCode, extractHeaders, resp.Body) if err != nil { - m.logMetadataCacheWriteFailure(orgRepo, revision, method, err) + return cacheResult(err) } } else { - apiDir := fmt.Sprintf("%s/api/%s/%s/revision/%s", config.SysConfig.Repos(), repoType, orgRepo, mainVersion) - apiMetaPath := fmt.Sprintf("%s/%s", apiDir, fmt.Sprintf("meta_%s.json", method)) + apiDir := RepositoryKey(repoType, orgRepo).Revision(config.SysConfig.Repos(), mainVersion) + apiMetaPath := filepath.Join(apiDir, "meta_"+method+".json") if !util.FileExists(apiMetaPath) { err = m.writeApiMetaFile(repoType, orgRepo, mainVersion, method, resp.StatusCode, extractHeaders, resp.Body) // create main dir if err != nil { - m.logMetadataCacheWriteFailure(orgRepo, mainVersion, method, err) + return cacheResult(err) } } } err = m.writeApiMetaFile(repoType, orgRepo, commitSha, method, resp.StatusCode, extractHeaders, resp.Body) if err != nil { - m.logMetadataCacheWriteFailure(orgRepo, commitSha, method, err) + return cacheResult(err) } - return &common.CacheContent{ - StatusCode: resp.StatusCode, - Headers: extractHeaders, - OriginContent: resp.Body, - }, nil + return content, nil } -func (m *MetaDao) logMetadataCacheWriteFailure(orgRepo, revision, method string, err error) { - fields := []interface{}{"repo", orgRepo, "revision", revision, "method", method, "error", err} - var accessErr *util.FileAccessError - if errors.As(err, &accessErr) { - fields = append(fields, "kind", accessErr.Kind, "path", accessErr.Path) - } - zap.S().Warnw("serving remote metadata without cache", fields...) +// Only for cache persistence operations: ENOENT during a write is a failed +// cache write too. flock may return a bare errno rather than an os.PathError. +// Validation, namespace conflicts and lock contention are not I/O errors. +func metadataCacheIOError(err error) bool { + var pathErr *os.PathError + var linkErr *os.LinkError + var errno syscall.Errno + return errors.As(err, &pathErr) || errors.As(err, &linkErr) || errors.As(err, &errno) } func (m *MetaDao) writeApiMetaFile(repoType, orgRepo, commitSha, method string, statusCode int, extractHeaders map[string]string, body []byte) error { - apiDir := fmt.Sprintf("%s/api/%s/%s/revision/%s", config.SysConfig.Repos(), repoType, orgRepo, commitSha) - apiMetaPath := fmt.Sprintf("%s/%s", apiDir, fmt.Sprintf("meta_%s.json", method)) + apiDir := RepositoryKey(repoType, orgRepo).Revision(config.SysConfig.Repos(), commitSha) + apiMetaPath := filepath.Join(apiDir, "meta_"+method+".json") err := util.MakeDirs(apiMetaPath) if err != nil { + util.ObserveFileAccessFailure("mkdir", err) zap.S().Errorf("create %s dir err.%v", apiMetaPath, err) - if accessErr, ok := util.ClassifyFileAccessError(apiMetaPath, err); ok { - return accessErr - } return err } if err = m.fileDao.WriteCacheRequest(apiMetaPath, statusCode, extractHeaders, body); err != nil { + util.ObserveFileAccessFailure("write", err) zap.S().Errorf("writeCacheRequest err.%v", err) - if accessErr, ok := util.ClassifyFileAccessError(apiMetaPath, err); ok { - return accessErr - } return err } + dependency.Default.Observe(dependency.MetadataWrite, true, time.Now()) return nil } diff --git a/internal/dao/meta_dao_degraded_test.go b/internal/dao/meta_dao_degraded_test.go index b92439f..8126f60 100644 --- a/internal/dao/meta_dao_degraded_test.go +++ b/internal/dao/meta_dao_degraded_test.go @@ -1,19 +1,22 @@ package dao import ( + "errors" "net/http" "net/http/httptest" "os" "path/filepath" "strings" + "syscall" "testing" "dingospeed/internal/data" "dingospeed/pkg/config" "dingospeed/pkg/consts" + "dingospeed/pkg/util" ) -func TestRequestAndSaveMetaServesRemoteResponseWhenCacheUnavailable(t *testing.T) { +func TestRequestAndSaveMetaReturnsUpstreamWhenCacheUnavailable(t *testing.T) { const body = `{"sha":"remote-commit","id":"org/repo"}` server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "application/json") @@ -36,13 +39,34 @@ func TestRequestAndSaveMetaServesRemoteResponseWhenCacheUnavailable(t *testing.T baseData := data.NewBaseData() fileDao := NewFileDao(nil, baseData, NewLockDao(baseData)) metaDao := NewMetaDao(fileDao, nil, baseData) - got, err := metaDao.requestAndSaveMeta("models", "org/repo", "main", "remote-commit", consts.RequestTypeGet, "") - if err != nil { - t.Fatalf("remote metadata should survive cache failure: %v", err) + before := util.FileAccessObservations() + got, err := metaDao.requestAndSaveMeta("models", "huggingface/org/repo", "main", "remote-commit", consts.RequestTypeGet, "") + if err != nil || got == nil || got.StatusCode != 200 || string(got.OriginContent) != body { + t.Fatalf("successful upstream response lost: content=%v err=%v", got, err) } - if string(got.OriginContent) != body { - t.Fatalf("got body %q, want %q", got.OriginContent, body) + assertObservationDelta(t, before, "mkdir") +} + +func TestMetadataCacheWriteObservationPreservesOriginalError(t *testing.T) { + old := config.SysConfig + defer func() { config.SysConfig = old }() + config.SysConfig = &config.Config{Server: config.ServerConfig{Repos: t.TempDir()}} + base := data.NewBaseData() + file := NewFileDao(nil, base, NewLockDao(base)) + meta := NewMetaDao(file, nil, base) + path := filepath.Join(config.SysConfig.Repos(), "api/models/org/repo/revision/main/meta_get.json") + if err := os.MkdirAll(path, 0700); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(path, "keep"), []byte("keep"), 0600); err != nil { + t.Fatal(err) + } + before := util.FileAccessObservations() + err := meta.writeApiMetaFile("models", "huggingface/org/repo", "main", "get", 200, nil, []byte(`{"sha":"abc"}`)) + if err == nil { + t.Fatal("expected atomic replacement failure") } + assertObservationDelta(t, before, "write") } func TestRequestAndSaveMetaStillCachesWhenStorageAvailable(t *testing.T) { @@ -68,7 +92,7 @@ func TestRequestAndSaveMetaStillCachesWhenStorageAvailable(t *testing.T) { baseData := data.NewBaseData() fileDao := NewFileDao(nil, baseData, NewLockDao(baseData)) metaDao := NewMetaDao(fileDao, nil, baseData) - if _, err := metaDao.requestAndSaveMeta("models", "org/repo", "main", "remote-commit", consts.RequestTypeGet, ""); err != nil { + if _, err := metaDao.requestAndSaveMeta("models", "huggingface/org/repo", "main", "remote-commit", consts.RequestTypeGet, ""); err != nil { t.Fatal(err) } for _, revision := range []string{"main", "remote-commit"} { @@ -78,3 +102,77 @@ func TestRequestAndSaveMetaStillCachesWhenStorageAvailable(t *testing.T) { } } } + +func TestOnlineMetadataCacheFailureAtEachWrite(t *testing.T) { + const body = `{"sha":"remote-commit"}` + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("ETag", `"fixture"`) + if r.Method != "HEAD" { + _, _ = w.Write([]byte(body)) + } + })) + defer server.Close() + old := config.SysConfig + defer func() { config.SysConfig = old }() + for _, revision := range []string{"main", "v1"} { + for _, target := range []string{"main", "remote-commit"} { + for _, method := range []string{"get", "head"} { + t.Run(revision+"/"+target+"/"+method, func(t *testing.T) { + config.SysConfig = &config.Config{Server: config.ServerConfig{Online: true, Repos: t.TempDir(), HfScheme: "http", HfNetLoc: strings.TrimPrefix(server.URL, "http://")}, Retry: config.Retry{Attempts: 1}} + p := filepath.Join(config.SysConfig.Repos(), "api/models/org/repo/revision", target, "meta_"+method+".json") + if err := os.MkdirAll(p, 0700); err != nil { + t.Fatal(err) + } + // For a non-main revision, a pre-existing main cache is deliberately + // skipped; remove this case from failure-count assertions. + base := data.NewBaseData() + file := NewFileDao(nil, base, NewLockDao(base)) + meta := NewMetaDao(file, nil, base) + got, err := meta.requestAndSaveMeta("models", "huggingface/org/repo", revision, "remote-commit", method, "") + if err != nil || got == nil || got.StatusCode != 200 { + t.Fatalf("response lost: %v", err) + } + if method == "get" && string(got.OriginContent) != body { + t.Fatal("upstream body altered") + } + if method == "head" && len(got.OriginContent) != 0 { + t.Fatal("HEAD body altered") + } + if info, err := os.Stat(p); err != nil || !info.IsDir() { + t.Fatal("existing object replaced") + } + }) + } + } + } +} + +func TestIncompleteUpstreamMetadataStillFails(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Length", "100") + _, _ = w.Write([]byte("partial")) + })) + defer server.Close() + old := config.SysConfig + defer func() { config.SysConfig = old }() + config.SysConfig = &config.Config{Server: config.ServerConfig{Online: true, Repos: "invalid\x00repos", HfScheme: "http", HfNetLoc: strings.TrimPrefix(server.URL, "http://")}, Retry: config.Retry{Attempts: 1}} + base := data.NewBaseData() + file := NewFileDao(nil, base, NewLockDao(base)) + meta := NewMetaDao(file, nil, base) + if got, err := meta.requestAndSaveMeta("models", "huggingface/org/repo", "main", "commit", "get", ""); err == nil || got != nil { + t.Fatalf("incomplete upstream masked: %v %v", got, err) + } +} + +func TestMetadataCacheToleranceDoesNotMaskValidation(t *testing.T) { + for _, err := range []error{errors.New("namespace case conflict"), errors.New("repository registration busy"), errors.New("symlink in repository path")} { + if metadataCacheIOError(err) { + t.Fatalf("non-storage error tolerated: %v", err) + } + } + for _, err := range []error{&os.PathError{Op: "write", Err: os.ErrNotExist}, &os.LinkError{Op: "rename", Err: os.ErrPermission}, syscall.Errno(5)} { + if !metadataCacheIOError(err) { + t.Fatalf("cache persistence error not recognized: %v", err) + } + } +} diff --git a/internal/dao/namespace_cleanup_test.go b/internal/dao/namespace_cleanup_test.go new file mode 100644 index 0000000..44bcc84 --- /dev/null +++ b/internal/dao/namespace_cleanup_test.go @@ -0,0 +1,51 @@ +package dao + +import ( + "os" + "testing" + "time" +) + +func TestHostedCleanupIsolatesSameContentAcrossNamespacesAndRejectsBadManifest(t *testing.T) { + u, _ := newTestUploadDao(t) + content := []byte("shared content in independent namespaces") + p := uploadParam("nested/file.bin", content) + p.Namespace, p.Repo = "alice", "team/resolve/model-a" + alice := mustUpload(t, u, p, content) + p.Namespace, p.Deferred = "bob", true + bob := mustUpload(t, u, p, content) + old := time.Now().Add(-48 * time.Hour) + for _, ns := range []string{"alice", "bob"} { + if err := os.Chtimes(localBlobPath("models", ns+"/"+p.Repo, bob.Sha256), old, old); err != nil { + t.Fatal(err) + } + } + removed, err := u.CleanupUnreferencedBlobs(time.Hour) + if err != nil || removed != 1 { + t.Fatalf("cleanup removed=%d err=%v", removed, err) + } + if _, err := os.Stat(localBlobPath("models", "alice/"+p.Repo, alice.Sha256)); err != nil { + t.Fatal("alice's referenced content lost:", err) + } + if _, err := os.Stat(localBlobPath("models", "bob/"+p.Repo, bob.Sha256)); !os.IsNotExist(err) { + t.Fatal("bob's unreferenced content was not reclaimed") + } + + orphan := []byte("unreferenced but unsafe to reclaim without valid reference data") + p = uploadParam("new.bin", orphan) + p.Namespace, p.Repo, p.Deferred = "alice", "team/resolve/model-a", true + staged := mustUpload(t, u, p, orphan) + blob := localBlobPath("models", "alice/"+p.Repo, staged.Sha256) + if err := os.Chtimes(blob, old, old); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(LocalManifestPath("models", "alice/"+p.Repo, alice.Commit), []byte("broken manifest"), 0644); err != nil { + t.Fatal(err) + } + if _, err := u.CleanupUnreferencedBlobs(time.Hour); err == nil { + t.Fatal("bad manifest treated as no references") + } + if _, err := os.Stat(blob); err != nil { + t.Fatal("content reclaimed despite unreadable reference data:", err) + } +} diff --git a/internal/dao/namespace_peer_test.go b/internal/dao/namespace_peer_test.go new file mode 100644 index 0000000..6e6370f --- /dev/null +++ b/internal/dao/namespace_peer_test.go @@ -0,0 +1,63 @@ +package dao + +import ( + "dingospeed/internal/downloader" + "dingospeed/pkg/config" + "dingospeed/pkg/repository" + "dingospeed/pkg/util" + "io" + "net/http" + "net/http/httptest" + "net/url" + "testing" +) + +func TestPeerFileRequestPreservesNamespaceAndSeparatesCredentials(t *testing.T) { + old := config.SysConfig + defer func() { config.SysConfig = old }() + config.SysConfig = &config.Config{Server: config.ServerConfig{}, Download: config.Download{ReqTimeout: 2}} + key := repository.RepoKey{Namespace: "用户#50%", RepoType: "models", Repo: "team/resolve/model-a"} + p := &downloader.TaskParam{RepoKey: key, Revision: "commit-one", FileName: "weights/a b.bin"} + upstreamCalls := 0 + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + upstreamCalls++ + if r.Header.Get("X-Dingo-Service-Token") != "" { + t.Error("internal token reached upstream") + } + w.Write([]byte("ok")) + })) + defer upstream.Close() + peer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/api/repositories/models/用户#50%/file" { + t.Errorf("peer path %s", r.URL.Path) + } + if r.URL.Query().Get("repo") != key.Repo || r.URL.Query().Get("path") != p.FileName || r.URL.Query().Get("revision") != p.Revision { + t.Errorf("peer locator %s", r.URL.RawQuery) + } + if r.Header.Get("X-Dingo-Service-Token") != "" { + t.Error("unexpected service credential") + } + http.Redirect(w, r, upstream.URL, 302) + })) + defer peer.Close() + if _, err := url.Parse(peerFileURI(p)); err != nil { + t.Fatal(err) + } + if err := util.GetPeerStream(peer.URL, peerFileURI(p), map[string]string{"Authorization": "Bearer upstream-test"}, func(r *http.Response) error { + if r.StatusCode != 302 { + t.Errorf("peer redirect followed: %d", r.StatusCode) + } + return nil + }); err != nil { + t.Fatal(err) + } + if upstreamCalls != 0 { + t.Fatal("peer redirect must not receive service token") + } + if err := util.GetStream(upstream.URL, "/", map[string]string{"X-Dingo-Service-Token": "must-be-stripped"}, func(r *http.Response) error { _, err := io.Copy(io.Discard, r.Body); return err }); err != nil { + t.Fatal(err) + } + if upstreamCalls != 1 { + t.Fatalf("upstream calls %d", upstreamCalls) + } +} diff --git a/internal/dao/official_content.go b/internal/dao/official_content.go new file mode 100644 index 0000000..899427d --- /dev/null +++ b/internal/dao/official_content.go @@ -0,0 +1,19 @@ +package dao + +import "fmt" + +// VerifyPublishedSHA256 is the deliberately expensive official-management check. +// Completion bits alone cannot prove the content still matches its identity. +// The existing lightweight inventory/publication checks keep their behavior. +func VerifyPublishedSHA256(repoType, id string, files []LocalManifestFile) error { + for _, f := range files { + digest, err := hashDingCachePayload(localBlobPath(repoType, id, f.Sha256), f.Size) + if err != nil { + return fmt.Errorf("selected content %s is unreadable: %w", f.Path, err) + } + if digest != f.Sha256 { + return fmt.Errorf("selected content %s SHA256 changed", f.Path) + } + } + return nil +} diff --git a/internal/dao/official_content_test.go b/internal/dao/official_content_test.go new file mode 100644 index 0000000..0581316 --- /dev/null +++ b/internal/dao/official_content_test.go @@ -0,0 +1,28 @@ +package dao + +import ( + "os" + "testing" +) + +func TestOfficialContentRejectsSameSizeCorruption(t *testing.T) { + u, _ := newTestUploadDao(t) + body := []byte("official payload") + mustUpload(t, u, uploadParam("file.bin", body), body) + files := []LocalManifestFile{manifestItem("file.bin", body)} + if err := VerifyPublishedSHA256("models", "dingo-local/demo", files); err != nil { + t.Fatal(err) + } + path := localBlobPath("models", "dingo-local/demo", files[0].Sha256) + raw, err := os.ReadFile(path) + if err != nil { + t.Fatal(err) + } + raw[len(raw)-1] ^= 1 + if err = os.WriteFile(path, raw, 0600); err != nil { + t.Fatal(err) + } + if err = VerifyPublishedSHA256("models", "dingo-local/demo", files); err == nil { + t.Fatal("same-size changed payload accepted") + } +} diff --git a/internal/dao/preheat_cache_complete.go b/internal/dao/preheat_cache_complete.go new file mode 100644 index 0000000..07c835a --- /dev/null +++ b/internal/dao/preheat_cache_complete.go @@ -0,0 +1,97 @@ +package dao + +import ( + "os" + "path/filepath" + + "dingospeed/internal/downloader" + "dingospeed/pkg/config" + "dingospeed/pkg/repository" +) + +// PreheatCacheComplete checks current local references and blob headers rather +// than trusting a historical successful job after cache cleanup. +func PreheatCacheComplete(key repository.RepoKey, snapshot *CommitHfSha) bool { + if key.Validate() != nil || snapshot == nil || repository.Segment(snapshot.Sha) != nil { + return false + } + var total int64 + for _, file := range snapshot.Siblings { + if repository.Relative(file.Rfilename) != nil { + return false + } + path := ResolvePath(key.RepoType, key.ID(), snapshot.Sha, file.Rfilename) + path, ok := preheatContentPath(key, path) + if !ok { + return false + } + size, complete := completePreheatFile(path) + if !complete || size > snapshot.UsedStorage-total { + return false + } + total += size + } + return total == snapshot.UsedStorage +} + +// Resolve references may be regular files (copy/hard-link fallback) or links +// directly into this repository's blobs. Parent directories and blob targets +// must still pass the usual no-symlink checks. +func preheatContentPath(key repository.RepoKey, path string) (string, bool) { + root := config.SysConfig.Repos() + if repository.SafePath(root, filepath.Dir(path)) != nil { + return "", false + } + info, err := os.Lstat(path) + if err != nil { + return "", false + } + if info.Mode()&os.ModeSymlink != 0 { + target, err := os.Readlink(path) + if err != nil { + return "", false + } + if !filepath.IsAbs(target) { + target = filepath.Join(filepath.Dir(path), target) + } + target, err = filepath.Abs(target) + if err != nil { + return "", false + } + blobs, err := filepath.Abs(filepath.Join(key.FilesRoot(root), "blobs")) + if err != nil || filepath.Dir(target) != blobs { + return "", false + } + path = target + } + if repository.SafePath(root, path) != nil { + return "", false + } + return path, true +} + +func completePreheatFile(path string) (int64, bool) { + f, err := os.Open(path) + if err != nil { + return 0, false + } + defer f.Close() + info, err := f.Stat() + if err != nil || !info.Mode().IsRegular() { + return 0, false + } + if info.Size() == 0 { + return 0, true + } // Empty remote files have no cache header. + header := &downloader.DingCacheHeader{} + if header.Read(f) != nil || header.FileSize > uint64(info.Size()) || header.GetHeaderSize() > info.Size()-int64(header.FileSize) { + return 0, false + } + for i := uint64(0); i < header.BlockNumber; i++ { + present, err := header.BlockMask.Test(i) + if err != nil || !present { + return 0, false + } + } + return int64(header.FileSize), true +} diff --git a/internal/dao/preheat_cache_complete_test.go b/internal/dao/preheat_cache_complete_test.go new file mode 100644 index 0000000..6f332b3 --- /dev/null +++ b/internal/dao/preheat_cache_complete_test.go @@ -0,0 +1,109 @@ +package dao + +import ( + "encoding/binary" + "encoding/json" + "os" + "path/filepath" + "testing" + + "dingospeed/pkg/config" + "dingospeed/pkg/repository" +) + +func TestPreheatCacheCompleteReferences(t *testing.T) { + for _, provider := range []string{repository.HuggingFace, repository.ModelScope} { + for _, mode := range []string{"relative-link", "absolute-link", "regular", "missing", "external", "other-repository", "dangling", "chained", "parent-link", "partial", "truncated", "size-mismatch"} { + t.Run(provider+"/"+mode, func(t *testing.T) { + old := config.SysConfig + t.Cleanup(func() { config.SysConfig = old }) + root := t.TempDir() + config.SysConfig = &config.Config{Server: config.ServerConfig{Repos: root}} + key := repository.RepoKey{Namespace: provider, RepoType: "models", Repo: "owner/model"} + var snapshot CommitHfSha + if err := json.Unmarshal([]byte(`{"sha":"commit1","siblings":[{"rfilename":"model.bin"}]}`), &snapshot); err != nil { + t.Fatal(err) + } + snapshot.UsedStorage = 8 + payload := make([]byte, 45) + copy(payload, "OLAH") + binary.LittleEndian.PutUint64(payload[4:], 8) + binary.LittleEndian.PutUint64(payload[12:], 4) + binary.LittleEndian.PutUint64(payload[20:], 8) + binary.LittleEndian.PutUint64(payload[28:], 8) + payload[36] = 3 + copy(payload[37:], "abcdefgh") + if mode == "partial" { + payload[36] = 1 + } + if mode == "truncated" { + payload = payload[:40] + } + if mode == "size-mismatch" { + snapshot.UsedStorage = 9 + } + write := func(path string) { + t.Helper() + if err := os.MkdirAll(filepath.Dir(path), 0755); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(path, payload, 0600); err != nil { + t.Fatal(err) + } + } + link := func(target, path string) { + t.Helper() + if err := os.Symlink(target, path); err != nil { + t.Skipf("symlinks unavailable: %v", err) + } + } + blob := key.Blob(root, "content") + write(blob) + path := key.Resolve(root, snapshot.Sha, "model.bin") + if err := os.MkdirAll(filepath.Dir(path), 0755); err != nil { + t.Fatal(err) + } + switch mode { + case "regular": + write(path) + case "missing": + case "parent-link": + if err := os.Remove(filepath.Dir(path)); err != nil { + t.Fatal(err) + } + parent := t.TempDir() + write(filepath.Join(parent, "model.bin")) + link(parent, filepath.Dir(path)) + default: + target := blob + switch mode { + case "external": + target = filepath.Join(t.TempDir(), "content") + write(target) + case "other-repository": + other := key + other.Repo = "owner/other" + target = other.Blob(root, "content") + write(target) + case "dangling": + target = key.Blob(root, "absent") + case "chained": + target = key.Blob(root, "alias") + link(blob, target) + case "relative-link": + var err error + target, err = filepath.Rel(filepath.Dir(path), blob) + if err != nil { + t.Fatal(err) + } + } + link(target, path) + } + want := mode == "relative-link" || mode == "absolute-link" || mode == "regular" + if got := PreheatCacheComplete(key, &snapshot); got != want { + t.Fatalf("complete=%v want=%v", got, want) + } + }) + } + } +} diff --git a/internal/dao/preheat_metadata_test.go b/internal/dao/preheat_metadata_test.go new file mode 100644 index 0000000..ff33145 --- /dev/null +++ b/internal/dao/preheat_metadata_test.go @@ -0,0 +1,60 @@ +package dao + +import ( + "fmt" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "dingospeed/internal/data" + "dingospeed/pkg/config" + "dingospeed/pkg/repository" +) + +func TestRepeatedPreheatRefreshesMainAndPreservesSnapshots(t *testing.T) { + old := config.SysConfig + defer func() { config.SysConfig = old }() + commit, calls, status := "commit-a", 0, 200 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + calls++ + if r.URL.Path != "/api/models/owner/model/revision/main" { + t.Errorf("unexpected request %s", r.URL) + } + w.WriteHeader(status) + fmt.Fprintf(w, `{"sha":%q,"siblings":[{"rfilename":"config.json"}]}`, commit) + })) + defer server.Close() + config.SysConfig = &config.Config{Server: config.ServerConfig{Online: true, Repos: t.TempDir(), HfScheme: "http", HfNetLoc: strings.TrimPrefix(server.URL, "http://")}, Retry: config.Retry{Attempts: 1}} + base := data.NewBaseData() + locks := NewLockDao(base) + files := NewFileDao(nil, base, locks) + meta := NewMetaDao(files, locks, base) + key := repository.RepoKey{Namespace: repository.HuggingFace, RepoType: "models", Repo: "owner/model"} + for _, next := range []string{"commit-a", "commit-b", "commit-a", "commit-a"} { + commit = next + got, err := meta.RefreshPreheatMetadata(key, "main", "") + if err != nil || got.Sha != next { + t.Fatalf("refresh=%+v err=%v", got, err) + } + local, err := files.GetCommitHfOffline("models", key.ID(), "main") + if err != nil || local != next { + t.Fatalf("main=%s err=%v", local, err) + } + } + if calls != 4 { + t.Fatalf("upstream calls=%d", calls) + } + for _, rev := range []string{"commit-a", "commit-b"} { + if got, err := files.GetCommitHfOffline("models", key.ID(), rev); err != nil || got != rev { + t.Fatalf("snapshot=%s err=%v", got, err) + } + } + status = 403 + if _, err := meta.RefreshPreheatMetadata(key, "main", ""); err == nil { + t.Fatal("upstream denial reported success") + } + if got, err := files.GetCommitHfOffline("models", key.ID(), "main"); err != nil || got != "commit-a" { + t.Fatalf("failed refresh changed main: %s %v", got, err) + } +} diff --git a/internal/dao/provider_metadata.go b/internal/dao/provider_metadata.go new file mode 100644 index 0000000..68dd970 --- /dev/null +++ b/internal/dao/provider_metadata.go @@ -0,0 +1,39 @@ +package dao + +import ( + "bytes" + "encoding/json" + "path/filepath" + + "dingospeed/pkg/common" + "dingospeed/pkg/config" + "dingospeed/pkg/repository" + "dingospeed/pkg/util" +) + +// CacheProviderFileInfo uses the HF paths-info envelope. Immutable file metadata +// need not be replaced on every read. The short lock also spans path inspection, +// avoiding concurrent Stat/rename conflicts on Windows across provider handlers. +func (f *FileDao) CacheProviderFileInfo(key repository.RepoKey, commit, requestedRevision string, info *common.PathsInfo) error { + target := filepath.Join(key.PathsInfo(config.SysConfig.Repos(), commit, info.Path), "paths-info_post.json") + body, err := json.Marshal([]*common.PathsInfo{info}) + if err != nil { + return err + } + return f.CacheProviderResponse(target, commit, requestedRevision, body) +} + +func (f *FileDao) CacheProviderResponse(target, commit, requestedRevision string, body []byte) error { + fileReferenceLocks.Lock(target) + defer fileReferenceLocks.Unlock(target) + if err := repository.SafePath(config.SysConfig.Repos(), target); err != nil { + return err + } + if cached, err := f.ReadCacheRequest(target); err == nil && bytes.Equal(cached.OriginContent, body) { + return nil + } + if err := util.MakeDirs(target); err != nil { + return err + } + return f.WriteCacheRequest(target, 200, map[string]string{"x-repo-commit": commit, "x-requested-revision": requestedRevision}, body) +} diff --git a/internal/dao/repository_dao.go b/internal/dao/repository_dao.go new file mode 100644 index 0000000..caf2395 --- /dev/null +++ b/internal/dao/repository_dao.go @@ -0,0 +1,36 @@ +package dao + +import ( + "dingospeed/pkg/config" + "dingospeed/pkg/repository" + "fmt" +) + +func RegisterHosted(repoType, namespace, repo string) error { + k := repository.RepoKey{Namespace: namespace, RepoType: repoType, Repo: repo} + if repoType == "spaces" { + return fmt.Errorf("spaces uploads are not supported") + } + if _, _, err := k.Storage(); err != nil { + return err + } + return repository.Register(config.SysConfig.Repos(), repository.Hosted(k)) +} + +func RepositoryKey(repoType, id string) repository.RepoKey { + k, _ := repository.ParseID(repoType, id) + return k +} + +func UpstreamRepo(repoType, id string) (string, error) { + k, err := repository.ParseID(repoType, id) + if err != nil { + return "", err + } + // Provider namespaces are reserved server configuration, so the upstream + // identity remains known even when the cache filesystem is unavailable. + if k.Namespace != repository.HuggingFace { + return "", fmt.Errorf("repository has no Hugging Face upstream") + } + return k.Repo, nil +} diff --git a/internal/dao/repository_lifecycle.go b/internal/dao/repository_lifecycle.go new file mode 100644 index 0000000..48e4457 --- /dev/null +++ b/internal/dao/repository_lifecycle.go @@ -0,0 +1,96 @@ +package dao + +import ( + "fmt" + "os" + "path/filepath" + "strings" + + "dingospeed/pkg/config" + "dingospeed/pkg/inventory" + "dingospeed/pkg/repository" + "go.uber.org/zap" +) + +// Ordinary uploads share this lock; deletion waits for active requests to finish. +// Keep this outside repo/revision/blob locks to preserve their lock ordering. +var repositoryLifecycle = newKeyedRWMutex() + +func holdRepository(repoType, namespace, repo string) func() { + key := uploadRepoLockKey(repoType, namespace+"/"+repo) + repositoryLifecycle.RLock(key) + return func() { repositoryLifecycle.RUnlock(key) } +} + +type RepositoryDeleteResult struct { + Deleted bool `json:"deleted"` +} + +// DeleteHostedRepository completes the existing unlink/purge flow. It refuses +// live references, removes remaining staged/orphan data through the existing +// physical purge primitive, then removes metadata and wakes inventory reporting. +func (u *UploadDao) DeleteHostedRepository(k repository.RepoKey) (*RepositoryDeleteResult, error) { + if err := k.Validate(); err != nil { + return nil, err + } + if k.Namespace == repository.HuggingFace || k.Namespace == repository.ModelScope { + return nil, localUploadError{status: 400, code: "HOSTED_REPOSITORY_REQUIRED", msg: "remote caches cannot be deleted as hosted repositories"} + } + lock := uploadRepoLockKey(k.RepoType, k.ID()) + repositoryLifecycle.Lock(lock) + defer repositoryLifecycle.Unlock(lock) + uploadRepoLocks.Lock(lock) + defer uploadRepoLocks.Unlock(lock) + root := config.SysConfig.Repos() + if err := u.recoverInventoryLocked(k); err != nil { + return nil, err + } + d, err := repository.Read(root, k) + if os.IsNotExist(err) { + return &RepositoryDeleteResult{}, nil + } + if err != nil { + return nil, err + } + if d.Source != "hosted" { + return nil, fmt.Errorf("repository is not hosted") + } + for _, target := range []string{k.FilesRoot(root), k.APIRoot(root)} { + if err = repository.SafePath(root, target); err != nil { + return nil, err + } + } + if err = validateRepositoryReferences(k.RepoType, k.ID()); err != nil { + return nil, err + } + if len(buildRepoIndex(k.RepoType, k.ID()).BySha) != 0 { + return nil, localUploadError{status: 409, code: "REPOSITORY_NOT_EMPTY", msg: "repository still has file references; retry file deletion first"} + } + blobs, err := os.ReadDir(filepath.Join(k.FilesRoot(root), "blobs")) + if err != nil && !os.IsNotExist(err) { + return nil, err + } + for _, blob := range blobs { + if blob.IsDir() { + return nil, fmt.Errorf("unexpected directory in repository blobs") + } + if err = reclaimBlobFile(k.RepoType, k.ID(), blob.Name()); err != nil { + return nil, err + } + } + if err = inventory.Run(root, inventoryKey(k), "delete-repository", k, func() error { return u.removeEmptyRepository(k) }); err != nil { + return nil, err + } + if u.fileDao.baseData != nil { + for key := range u.fileDao.baseData.Cache.Items() { + for _, prefix := range []string{"localManifest/", "meta/", "metadatareq/", "filePathInfo/"} { + if strings.HasPrefix(key, prefix+k.RepoType+"/"+k.ID()+"/") { + u.fileDao.baseData.Cache.Delete(key) + } + } + } + } + NotifyPublished() + zap.S().Infof("[REPOSITORY] permanently deleted %s/%s", k.RepoType, k.ID()) + return &RepositoryDeleteResult{Deleted: true}, nil +} diff --git a/internal/dao/repository_lifecycle_test.go b/internal/dao/repository_lifecycle_test.go new file mode 100644 index 0000000..fb6473d --- /dev/null +++ b/internal/dao/repository_lifecycle_test.go @@ -0,0 +1,140 @@ +package dao + +import ( + "bytes" + "os" + "path/filepath" + "sync" + "sync/atomic" + "testing" + + "dingospeed/pkg/repository" +) + +func TestRepositoryDeleteRecreatePurgesDataAndPreservesOtherRepositories(t *testing.T) { + for _, withFiles := range []bool{false, true} { + t.Run(map[bool]string{false: "empty", true: "files"}[withFiles], func(t *testing.T) { + u, root := newTestUploadDao(t) + k := repository.RepoKey{RepoType: "models", Namespace: "dingo-local", Repo: "demo"} + p := publishParam("main") + if withFiles { + body := []byte("original bytes") + staged := deferredParam("a.bin", body) + mustStage(t, u, staged, body) + p.Files = []LocalManifestFile{manifestItem("a.bin", body)} + } + old := mustPublish(t, u, p) + if _, err := u.fileDao.ReadLocalManifest(k.RepoType, k.ID(), old.Commit); err != nil { + t.Fatal(err) + } + other := publishParam("main") + other.Repo = "other" + mustPublish(t, u, other) + admin := NewCacheAdminDao(u.fileDao) + if withFiles { + items := []DeleteItem{{RepoType: k.RepoType, OrgRepo: k.ID(), Path: "a.bin", Sha: p.Files[0].Sha256}} + if _, err := u.DeleteHostedRepository(k); err == nil { + t.Fatal("finalization accepted live references") + } + rows, err := admin.SoftDelete(items) + if err != nil || rows[0].Status == "failed" { + t.Fatalf("unlink: %+v %v", rows, err) + } + rows, err = admin.PurgeOrphans(items) + if err != nil || rows[0].Status == "failed" { + t.Fatalf("purge: %+v %v", rows, err) + } + } + deleted, err := u.DeleteHostedRepository(k) + if err != nil || !deleted.Deleted { + t.Fatalf("delete=%+v %v", deleted, err) + } + for _, path := range []string{k.APIRoot(root), k.FilesRoot(root), filepath.Join(root, "repository-trash")} { + if _, err := os.Stat(path); !os.IsNotExist(err) { + t.Fatalf("data retained at %s: %v", path, err) + } + } + if _, err = repository.Read(root, k); !os.IsNotExist(err) { + t.Fatalf("still active: %v", err) + } + if _, err = u.fileDao.ReadLocalManifest(k.RepoType, k.ID(), old.Commit); !os.IsNotExist(err) { + t.Fatalf("old snapshot cache survived: %v", err) + } + if result, err := u.DeleteHostedRepository(k); err != nil || result.Deleted { + t.Fatalf("retry=%+v %v", result, err) + } + p.Files = nil + p.CreateOnly = true + recreated := mustPublish(t, u, p) + files, err := u.fileDao.ReadLocalManifest(k.RepoType, k.ID(), recreated.Commit) + if err != nil || len(files) != 0 { + t.Fatalf("recreated=%+v %v", files, err) + } + if _, err := u.PublishFiles(p); err == nil { + t.Fatal("duplicate empty repository accepted") + } + if _, err := repository.Read(root, repository.RepoKey{RepoType: "models", Namespace: "dingo-local", Repo: "other"}); err != nil { + t.Fatal(err) + } + if withFiles { + body := []byte("original bytes") + staged := deferredParam("a.bin", body) + result, err := u.UploadWholeFile(staged, bytes.NewReader(body)) + if err != nil || result.BlobReused { + t.Fatalf("deleted blob reused: %+v %v", result, err) + } + } + }) + } +} + +func TestRepositoryCreateOnlyHasOneWinner(t *testing.T) { + u, _ := newTestUploadDao(t) + var winners atomic.Int32 + var wg sync.WaitGroup + for i := 0; i < 12; i++ { + wg.Add(1) + go func() { + defer wg.Done() + p := publishParam("main") + p.CreateOnly = true + if _, err := u.PublishFiles(p); err == nil { + winners.Add(1) + } + }() + } + wg.Wait() + if winners.Load() != 1 { + t.Fatalf("winners=%d", winners.Load()) + } +} + +func TestRepositoryDeleteFailureAndRemoteIsolation(t *testing.T) { + u, root := newTestUploadDao(t) + mustPublish(t, u, publishParam("main")) + k := repository.RepoKey{RepoType: "models", Namespace: "dingo-local", Repo: "demo"} + _ = root + body := []byte("still live") + mustStage(t, u, deferredParam("live.bin", body), body) + p := publishParam("main") + p.Files = []LocalManifestFile{manifestItem("live.bin", body)} + mustPublish(t, u, p) + if _, err := u.DeleteHostedRepository(k); err == nil { + t.Fatal("live repository deleted without unlink") + } + if _, err := repository.Read(root, k); err != nil { + t.Fatal(err) + } + for _, ns := range []string{repository.HuggingFace, repository.ModelScope} { + key := repository.RepoKey{RepoType: "models", Namespace: ns, Repo: "org/demo"} + if err := repository.Register(root, repository.Remote(key)); err != nil { + t.Fatal(err) + } + if _, err := u.DeleteHostedRepository(key); err == nil { + t.Fatal("remote accepted") + } + if _, err := repository.Read(root, key); err != nil { + t.Fatal(err) + } + } +} diff --git a/internal/dao/rpc_storage_test.go b/internal/dao/rpc_storage_test.go new file mode 100644 index 0000000..80c32f4 --- /dev/null +++ b/internal/dao/rpc_storage_test.go @@ -0,0 +1,55 @@ +package dao + +import ( + "context" + pb "dingospeed/pkg/proto/manager" + "google.golang.org/grpc" + "google.golang.org/protobuf/types/known/emptypb" + "testing" +) + +type storageClient struct { + pb.ManagerClient + file *pb.SchedulerFileRequest + sync *pb.SyncFileProcessReq +} + +func (c *storageClient) SchedulerFile(_ context.Context, r *pb.SchedulerFileRequest, _ ...grpc.CallOption) (*pb.SchedulerFileResponse, error) { + c.file = r + return &pb.SchedulerFileResponse{ProcessId: 41}, nil +} +func (c *storageClient) SyncFileProcess(_ context.Context, r *pb.SyncFileProcessReq, _ ...grpc.CallOption) (*emptypb.Empty, error) { + c.sync = r + return &emptypb.Empty{}, nil +} +func TestRPCStorageBoundaryDoesNotMutateRetryMessage(t *testing.T) { + c := &storageClient{} + d := &SchedulerDao{Client: c} + for _, ns := range []string{"huggingface", "datacanvas", "alice"} { + repo := "demo" + org := "dingo-local/" + ns + if ns == "huggingface" { + repo = "Qwen/demo" + org = "Qwen" + } + req := &pb.SchedulerFileRequest{DataType: "models", Org: ns, Repo: repo, Name: "weights.bin", Etag: "hash"} + for i := 0; i < 2; i++ { + if _, e := d.SchedulerFile(req); e != nil { + t.Fatal(e) + } + if c.file.Org != org || c.file.Repo != "demo" { + t.Fatalf("wrong SQL coordinates: %v", c.file) + } + if req.Org != ns || req.Repo != repo { + t.Fatal("retry message mutated") + } + } + sync := &pb.SyncFileProcessReq{FileProcessEntries: []*pb.FileProcessEntry{{DataType: "models", Org: ns, Repo: repo}}} + if e := d.SyncFileProcess(sync); e != nil { + t.Fatal(e) + } + if c.sync.FileProcessEntries[0].Org != org || sync.FileProcessEntries[0].Org != ns { + t.Fatal("wrong sync boundary") + } + } +} diff --git a/internal/dao/scheduler_dao.go b/internal/dao/scheduler_dao.go index cd715b4..ea48d7d 100644 --- a/internal/dao/scheduler_dao.go +++ b/internal/dao/scheduler_dao.go @@ -2,10 +2,14 @@ package dao import ( "context" + "dingospeed/pkg/repository" + "google.golang.org/protobuf/proto" + "time" "dingospeed/internal/data" "dingospeed/pkg/config" "dingospeed/pkg/consts" + "dingospeed/pkg/dependency" "dingospeed/pkg/proto/manager" "go.uber.org/zap" @@ -20,14 +24,18 @@ func NewSchedulerDao() *SchedulerDao { } func (s *SchedulerDao) Register() (*manager.RegisterResponse, error) { + managementURL, downloadURL := config.SysConfig.RegistrationEndpoints() ctx, cancel := context.WithTimeout(context.Background(), consts.RpcRequestTimeout) defer cancel() r, err := s.Client.Register(ctx, &manager.RegisterRequest{ - InstanceId: config.SysConfig.Scheduler.Discovery.InstanceId, - Host: config.SysConfig.Scheduler.Discovery.Host, - Port: int32(config.SysConfig.Scheduler.Discovery.Port), - Online: config.SysConfig.Server.Online, + InstanceId: config.SysConfig.Registration().NodeID, + Host: config.SysConfig.Registration().Host, + Port: int32(config.SysConfig.Registration().Port), + Online: config.SysConfig.Server.Online, + ManagementUrl: managementURL, + DownloadUrl: downloadURL, }) + dependency.Default.Observe(dependency.Scheduler, err == nil, time.Now()) if err != nil { zap.S().Errorf("speed register fail.%v", err) return nil, err @@ -39,14 +47,25 @@ func (s *SchedulerDao) Heartbeat() error { ctx, cancel := context.WithTimeout(context.Background(), consts.RpcRequestTimeout) defer cancel() req := &manager.HeartbeatRequest{ - Id: config.SysConfig.Id, - InstanceId: config.SysConfig.Scheduler.Discovery.InstanceId, - Online: config.SysConfig.Server.Online} + Id: config.SysConfig.SchedulerID(), + InstanceId: config.SysConfig.Registration().NodeID, + Online: config.SysConfig.Server.Online, + Health: heartbeatHealth(time.Now())} _, err := s.Client.Heartbeat(ctx, req) + dependency.Default.Observe(dependency.Scheduler, err == nil, time.Now()) return err } func (s *SchedulerDao) SchedulerFile(req *manager.SchedulerFileRequest) (*manager.SchedulerFileResponse, error) { + req = proto.Clone(req).(*manager.SchedulerFileRequest) + if req.Org != "" || req.Repo != "" { + org, repo, err := (repository.RepoKey{Namespace: req.Org, RepoType: req.DataType, Repo: req.Repo}).Storage() + if err != nil { + return nil, err + } + req.Org, req.Repo = org, repo + } + ctx, cancel := context.WithTimeout(context.Background(), consts.RpcRequestTimeout) defer cancel() resp, err := s.Client.SchedulerFile(ctx, req) @@ -54,6 +73,15 @@ func (s *SchedulerDao) SchedulerFile(req *manager.SchedulerFileRequest) (*manage } func (s *SchedulerDao) SyncFileProcess(req *manager.SyncFileProcessReq) error { + req = proto.Clone(req).(*manager.SyncFileProcessReq) + for _, e := range req.FileProcessEntries { + org, repo, err := (repository.RepoKey{Namespace: e.Org, RepoType: e.DataType, Repo: e.Repo}).Storage() + if err != nil { + return err + } + e.Org, e.Repo = org, repo + } + ctx, cancel := context.WithTimeout(context.Background(), consts.RpcRequestTimeout) defer cancel() _, err := s.Client.SyncFileProcess(ctx, req) @@ -71,6 +99,15 @@ func (s *SchedulerDao) ReportFileProcess(request *manager.FileProcessRequest) er } func (s *SchedulerDao) DeleteByEtagsAndFields(request *manager.DeleteByEtagsAndFieldsRequest) { + request = proto.Clone(request).(*manager.DeleteByEtagsAndFieldsRequest) + if request.Org != "" || request.Repo != "" { + org, repo, err := (repository.RepoKey{Namespace: request.Org, RepoType: request.Datatype, Repo: request.Repo}).Storage() + if err != nil { + return + } + request.Org, request.Repo = org, repo + } + ctx, cancel := context.WithTimeout(context.Background(), consts.RpcRequestTimeout) defer cancel() _, err := s.Client.DeleteByEtagsAndFields(ctx, request) @@ -81,6 +118,15 @@ func (s *SchedulerDao) DeleteByEtagsAndFields(request *manager.DeleteByEtagsAndF } func (s *SchedulerDao) CreateCacheJob(request *manager.CreateCacheJobReq) (*manager.CreateCacheJobResp, error) { + request = proto.Clone(request).(*manager.CreateCacheJobReq) + if request.Org != "" || request.Repo != "" { + org, repo, err := (repository.RepoKey{Namespace: request.Org, RepoType: request.Datatype, Repo: request.Repo}).Storage() + if err != nil { + return nil, err + } + request.Org, request.Repo = org, repo + } + ctx, cancel := context.WithTimeout(context.Background(), consts.RpcRequestTimeout) defer cancel() resp, err := s.Client.CreateCacheJob(ctx, request) @@ -92,6 +138,15 @@ func (s *SchedulerDao) CreateCacheJob(request *manager.CreateCacheJobReq) (*mana } func (s *SchedulerDao) UpdateCacheJobStatus(request *manager.UpdateCacheJobStatusReq) error { + request = proto.Clone(request).(*manager.UpdateCacheJobStatusReq) + if request.Org != "" || request.Repo != "" { + org, repo, err := (repository.RepoKey{Namespace: request.Org, RepoType: "models", Repo: request.Repo}).Storage() + if err != nil { + return err + } + request.Org, request.Repo = org, repo + } + zap.S().Infof("update status cacheJobId:%d, status:%d, %s", request.Id, request.Status, request.ErrorMsg) ctx, cancel := context.WithTimeout(context.Background(), consts.RpcRequestTimeout) defer cancel() diff --git a/internal/dao/scheduler_observation_test.go b/internal/dao/scheduler_observation_test.go new file mode 100644 index 0000000..3a088d1 --- /dev/null +++ b/internal/dao/scheduler_observation_test.go @@ -0,0 +1,54 @@ +package dao + +import ( + "context" + "dingospeed/pkg/config" + "dingospeed/pkg/dependency" + "dingospeed/pkg/proto/manager" + "errors" + "google.golang.org/grpc" + "google.golang.org/protobuf/types/known/emptypb" + "testing" + "time" +) + +type observationClient struct { + manager.ManagerClient + err error + response *manager.RegisterResponse + calls int +} + +func (c *observationClient) Register(context.Context, *manager.RegisterRequest, ...grpc.CallOption) (*manager.RegisterResponse, error) { + c.calls++ + return c.response, c.err +} +func (c *observationClient) Heartbeat(context.Context, *manager.HeartbeatRequest, ...grpc.CallOption) (*emptypb.Empty, error) { + c.calls++ + return nil, c.err +} + +func TestSchedulerObservationPreservesRPCResults(t *testing.T) { + old := config.SysConfig + config.SysConfig = &config.Config{} + defer func() { config.SysConfig = old }() + client := &observationClient{response: &manager.RegisterResponse{}} + dao := NewSchedulerDao() + dao.Client = client + response, err := dao.Register() + if err != nil || response != client.response || client.calls != 1 { + t.Fatal("successful registration changed") + } + want := errors.New("simulated unavailable") + client.err = want + _, err = dao.Register() + if err != want || client.calls != 2 { + t.Fatal("failure changed or retried") + } + if err := dao.Heartbeat(); err != want || client.calls != 3 { + t.Fatal("heartbeat failure changed or retried") + } + if !dependency.Default.Snapshot(dependency.Scheduler, time.Now()).Unresolved { + t.Fatal("RPC failure was not observed") + } +} diff --git a/internal/dao/upload_chunk_dao.go b/internal/dao/upload_chunk_dao.go index 170dc16..a5eafde 100644 --- a/internal/dao/upload_chunk_dao.go +++ b/internal/dao/upload_chunk_dao.go @@ -27,11 +27,11 @@ import ( // 而 publish 时 verifyPublishContent → inspectCompleteBlob 会检查“位图无空洞 + // size 匹配”。因此不完整的 blob 对下载侧天然不可见，分块上传不需要自己的生效开关。 type LocalChunkUploadParam struct { - RepoType string - Org string - Repo string - Revision string - FilePath string + RepoType string + Namespace string + Repo string + Revision string + FilePath string // Sha256 是整文件摘要，也就是 blob 文件名；本接口不校验它。 Sha256 string // Size 是整文件总字节数，用于绑定 blob 的 header。 @@ -45,6 +45,7 @@ type LocalChunkUploadParam struct { } type LocalChunkUploadResult struct { + Namespace string `json:"namespace"` RepoType string `json:"repoType"` Repo string `json:"repo"` Revision string `json:"revision"` @@ -72,7 +73,12 @@ func uploadBlobLockKey(repoType, orgRepo, sha string) string { // blob 读锁（挡住老接口的 rename 与回收的 remove）、DingCacheManager 的进程内唯一 // 句柄（位图不丢更新、header 不撕裂）、以及 WriteBlock 内部 fileLock 串行的置位与刷盘。 func (u *UploadDao) UploadChunk(param LocalChunkUploadParam, body io.Reader) (*LocalChunkUploadResult, error) { - orgRepo := util.GetOrgRepo(param.Org, param.Repo) + release := holdRepository(param.RepoType, param.Namespace, param.Repo) + defer release() + if err := RegisterHosted(param.RepoType, param.Namespace, param.Repo); err != nil { + return nil, localUploadError{status: 409, code: "REPOSITORY_REGISTRATION_CONFLICT", msg: err.Error()} + } + orgRepo := util.GetOrgRepo(param.Namespace, param.Repo) repos := config.SysConfig.Repos() blobPath := localBlobPath(param.RepoType, orgRepo, param.Sha256) if err := ensureLocalUploadPathSafe(repos, blobPath); err != nil { @@ -118,7 +124,8 @@ func (u *UploadDao) UploadChunk(param LocalChunkUploadParam, body io.Reader) (*L result := &LocalChunkUploadResult{ RepoType: param.RepoType, - Repo: orgRepo, + Namespace: param.Namespace, + Repo: param.Repo, Revision: param.Revision, FilePath: param.FilePath, Sha256: param.Sha256, diff --git a/internal/dao/upload_chunk_dao_test.go b/internal/dao/upload_chunk_dao_test.go index c435b82..64b6773 100644 --- a/internal/dao/upload_chunk_dao_test.go +++ b/internal/dao/upload_chunk_dao_test.go @@ -12,13 +12,13 @@ import ( func chunkParam(filePath string, content []byte) LocalChunkUploadParam { return LocalChunkUploadParam{ - RepoType: "models", - Org: "dingo-local", - Repo: "demo", - Revision: "main", - FilePath: filePath, - Sha256: sha256Hex(content), - Size: int64(len(content)), + RepoType: "models", + Namespace: "dingo-local", + Repo: "demo", + Revision: "main", + FilePath: filePath, + Sha256: sha256Hex(content), + Size: int64(len(content)), } } @@ -57,12 +57,12 @@ func chunkPlan(size, chunkLen int64) [][2]int64 { func progressOf(t *testing.T, u *UploadDao, param LocalChunkUploadParam) *LocalUploadProgress { t.Helper() progress, err := u.QueryProgress(LocalUploadParam{ - RepoType: param.RepoType, - Org: param.Org, - Repo: param.Repo, - Revision: param.Revision, - FilePath: param.FilePath, - Sha256: param.Sha256, + RepoType: param.RepoType, + Namespace: param.Namespace, + Repo: param.Repo, + Revision: param.Revision, + FilePath: param.FilePath, + Sha256: param.Sha256, }) if err != nil { t.Fatalf("query progress failed: %v", err) diff --git a/internal/dao/upload_cleanup_scope_test.go b/internal/dao/upload_cleanup_scope_test.go new file mode 100644 index 0000000..7b3d051 --- /dev/null +++ b/internal/dao/upload_cleanup_scope_test.go @@ -0,0 +1,79 @@ +package dao + +import ( + "os" + "path/filepath" + "testing" + "time" + + "dingospeed/pkg/repository" +) + +func TestUploadCleanupIgnoresBrokenRemoteRegistry(t *testing.T) { + u, root := newTestUploadDao(t) + k := repository.RepoKey{Namespace: "team", RepoType: "models", Repo: "nested/model"} + if err := RegisterHosted(k.RepoType, k.Namespace, k.Repo); err != nil { + t.Fatal(err) + } + blobs := filepath.Join(k.FilesRoot(root), "blobs") + if err := os.MkdirAll(blobs, 0755); err != nil { + t.Fatal(err) + } + stage := filepath.Join(blobs, "pending"+localUploadStageSuffix) + final := filepath.Join(blobs, "kept") + for _, p := range []string{stage, final} { + if err := os.WriteFile(p, []byte("fixture"), 0644); err != nil { + t.Fatal(err) + } + } + old := time.Now().Add(-48 * time.Hour) + if err := os.Chtimes(stage, old, old); err != nil { + t.Fatal(err) + } + // Any accidental full registry scan (including per-file Locate) fails here. + remote := filepath.Join(root, "api", "models", "remote", "model", repository.Marker) + if err := os.MkdirAll(filepath.Dir(remote), 0755); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(remote, []byte("invalid descriptor"), 0644); err != nil { + t.Fatal(err) + } + if n, err := u.CleanupExpiredStagedUploads(24 * time.Hour); err != nil || n != 1 { + t.Fatalf("cleanup: removed=%d err=%v", n, err) + } + if _, err := os.Stat(stage); !os.IsNotExist(err) { + t.Fatalf("expired stage remains: %v", err) + } + found, err := collectLocalBlobs(filepath.Join(root, "files")) + if err != nil { + t.Fatal(err) + } + got := found[localRepoKey{k.RepoType, k.ID()}] + if len(got) != 1 || got[0] != "kept" { + t.Fatalf("unexpected blobs: %v", found) + } + if b, err := os.ReadFile(remote); err != nil || string(b) != "invalid descriptor" { + t.Fatalf("remote fixture changed: %q %v", b, err) + } +} + +func TestStagedBlobIdentityStaysWithinKnownRepository(t *testing.T) { + root := t.TempDir() + d := repository.Hosted(repository.RepoKey{Namespace: "team", RepoType: "models", Repo: "nested/model"}) + base := d.FilesRoot(root) + for _, tc := range []struct { + path string + valid bool + }{ + {filepath.Join(base, "blobs", "hash"+localUploadStageSuffix), true}, + {filepath.Join(base, "blobs", "nested", "hash"+localUploadStageSuffix), false}, + {filepath.Join(base, "blobs", localUploadStageSuffix), false}, + {filepath.Join(base, "resolve", "hash"+localUploadStageSuffix), false}, + {filepath.Join(base+"-other", "blobs", "hash"+localUploadStageSuffix), false}, + } { + typ, id, sha, ok := stagedBlobIdentity(root, tc.path, d) + if ok != tc.valid || (ok && (typ != d.RepoType || id != d.ID() || sha != "hash")) { + t.Fatalf("%s: got %s %s %s %v", tc.path, typ, id, sha, ok) + } + } +} diff --git a/internal/dao/upload_dao.go b/internal/dao/upload_dao.go index 31cf709..4b487a8 100644 --- a/internal/dao/upload_dao.go +++ b/internal/dao/upload_dao.go @@ -17,6 +17,8 @@ import ( "dingospeed/internal/downloader" "dingospeed/pkg/config" + "dingospeed/pkg/inventory" + "dingospeed/pkg/repository" "dingospeed/pkg/util" "github.com/bytedance/sonic" @@ -30,7 +32,7 @@ type UploadDao struct { type LocalUploadParam struct { RepoType string - Org string + Namespace string Repo string Revision string FilePath string @@ -47,19 +49,21 @@ type LocalUploadParam struct { // 服务端不记忆“哪些上传属于哪一批”——理由与“不另建进度状态”完全一致： // 独立状态与实际落盘内容之间必然出现不一致，而清单三个字段本来就是上传时的必填项。 type LocalPublishParam struct { - RepoType string - Org string - Repo string - Revision string - Overwrite bool - Files []LocalManifestFile + CreateOnly bool + RepoType string + Namespace string + Repo string + Revision string + Overwrite bool + Files []LocalManifestFile } type LocalPublishResult struct { - RepoType string `json:"repoType"` - Repo string `json:"repo"` - Revision string `json:"revision"` - Commit string `json:"commit"` + Namespace string `json:"namespace"` + RepoType string `json:"repoType"` + Repo string `json:"repo"` + Revision string `json:"revision"` + Commit string `json:"commit"` // Published 是本批次声明的条目数，FileCount 是合并后该版本的文件总数。 Published int `json:"published"` FileCount int `json:"fileCount"` @@ -72,25 +76,27 @@ type LocalPublishResult struct { } type LocalUploadResult struct { - RepoType string `json:"repoType"` - Repo string `json:"repo"` - Revision string `json:"revision"` - Commit string `json:"commit"` - FilePath string `json:"filePath"` - Size int64 `json:"size"` - Sha256 string `json:"sha256"` - Status string `json:"status"` + Namespace string `json:"namespace"` + RepoType string `json:"repoType"` + Repo string `json:"repo"` + Revision string `json:"revision"` + Commit string `json:"commit"` + FilePath string `json:"filePath"` + Size int64 `json:"size"` + Sha256 string `json:"sha256"` + Status string `json:"status"` // BlobReused 为 true 表示本次未重复写入内容，走的是 FR-2 6.2.5 的幂等快路径。 BlobReused bool `json:"blobReused"` } type LocalUploadProgress struct { - RepoType string `json:"repoType"` - Repo string `json:"repo"` - Revision string `json:"revision"` - FilePath string `json:"filePath"` - Size int64 `json:"size"` - Sha256 string `json:"sha256"` + Namespace string `json:"namespace"` + RepoType string `json:"repoType"` + Repo string `json:"repo"` + Revision string `json:"revision"` + FilePath string `json:"filePath"` + Size int64 `json:"size"` + Sha256 string `json:"sha256"` // ResumeOffset 是第一个空洞的位置，也就是顺序续传（老接口的 start）该从哪里接上。 // 乱序并发的分块上传下它没有意义，用 MissingRanges。 ResumeOffset int64 `json:"resumeOffset"` @@ -122,7 +128,7 @@ type LocalManifestFile struct { const localManifestFileName = "dingo-local-manifest.json" func LocalManifestPath(repoType, orgRepo, commit string) string { - return filepath.Join(config.SysConfig.Repos(), "api", repoType, orgRepo, "revision", commit, localManifestFileName) + return filepath.Join(RepositoryKey(repoType, orgRepo).Revision(config.SysConfig.Repos(), commit), localManifestFileName) } func NewUploadDao(fileDao *FileDao, lockDao *LockDao) *UploadDao { @@ -266,11 +272,32 @@ func IsLocalOrgRepo(orgRepo string) bool { if config.SysConfig != nil && config.SysConfig.Upload.Namespace != "" { namespace = config.SysConfig.Upload.Namespace } - return orgRepo == namespace || strings.HasPrefix(orgRepo, namespace+"/") + if orgRepo == namespace || strings.HasPrefix(orgRepo, namespace+"/") { + return true + } + if config.SysConfig == nil { + return false + } + for _, typ := range []string{"models", "datasets", "spaces"} { + k, err := repository.ParseID(typ, orgRepo) + if err != nil { + return false + } + d, err := repository.Read(config.SysConfig.Repos(), k) + if err == nil && d.Source == "hosted" && d.Persistent { + return true + } + } + return false } func (u *UploadDao) UploadWholeFile(param LocalUploadParam, body io.Reader) (*LocalUploadResult, error) { - orgRepo := util.GetOrgRepo(param.Org, param.Repo) + release := holdRepository(param.RepoType, param.Namespace, param.Repo) + defer release() + if err := RegisterHosted(param.RepoType, param.Namespace, param.Repo); err != nil { + return nil, localUploadError{status: 409, code: "REPOSITORY_REGISTRATION_CONFLICT", msg: err.Error()} + } + orgRepo := util.GetOrgRepo(param.Namespace, param.Repo) fileKey := fmt.Sprintf("%s:%s:%s:%s", param.RepoType, orgRepo, param.Revision, param.FilePath) if !tryEnterLocalUpload(fileKey) { return nil, localUploadError{status: http.StatusConflict, code: "UPLOAD_FILE_BUSY", msg: "same file is already being uploaded"} @@ -289,7 +316,8 @@ func (u *UploadDao) UploadWholeFile(param LocalUploadParam, body io.Reader) (*Lo } return &LocalUploadResult{ RepoType: param.RepoType, - Repo: orgRepo, + Namespace: param.Namespace, + Repo: param.Repo, Revision: param.Revision, FilePath: param.FilePath, Size: param.Size, @@ -321,6 +349,9 @@ func (u *UploadDao) UploadWholeFile(param LocalUploadParam, body io.Reader) (*Lo repoLockKey := uploadRepoLockKey(param.RepoType, orgRepo) uploadRepoLocks.Lock(repoLockKey) defer uploadRepoLocks.Unlock(repoLockKey) + if err := u.recoverInventoryLocked(RepositoryKey(param.RepoType, orgRepo)); err != nil { + return nil, err + } uploadRevisionLocks.Lock(revisionLockKey) defer uploadRevisionLocks.Unlock(revisionLockKey) @@ -346,7 +377,8 @@ func (u *UploadDao) UploadWholeFile(param LocalUploadParam, body io.Reader) (*Lo } return &LocalUploadResult{ RepoType: param.RepoType, - Repo: orgRepo, + Namespace: param.Namespace, + Repo: param.Repo, Revision: param.Revision, Commit: commit, FilePath: param.FilePath, @@ -363,7 +395,9 @@ func (u *UploadDao) UploadWholeFile(param LocalUploadParam, body io.Reader) (*Lo // 只要合并后的清单相同，快照标识就必然逐字符相同，两条路径进系统的内容不会被 // 客户端当成两个不同的版本。 func (u *UploadDao) PublishFiles(param LocalPublishParam) (*LocalPublishResult, error) { - orgRepo := util.GetOrgRepo(param.Org, param.Repo) + release := holdRepository(param.RepoType, param.Namespace, param.Repo) + defer release() + orgRepo := util.GetOrgRepo(param.Namespace, param.Repo) // 发布之间必须明确拒绝而不是排队（BR-4.2.1），所以这里用 try-enter 而不是版本锁。 publishKey := fmt.Sprintf("upload-publish:%s:%s:%s", param.RepoType, orgRepo, param.Revision) @@ -381,8 +415,23 @@ func (u *UploadDao) PublishFiles(param LocalPublishParam) (*LocalPublishResult, repoLockKey := uploadRepoLockKey(param.RepoType, orgRepo) uploadRepoLocks.Lock(repoLockKey) defer uploadRepoLocks.Unlock(repoLockKey) + if err := u.recoverInventoryLocked(RepositoryKey(param.RepoType, orgRepo)); err != nil { + return nil, err + } // 版本锁与即时生效上传共用同一把：两者都要“读清单 → 合并 → 写回”， + if param.CreateOnly { + _, err := repository.Read(config.SysConfig.Repos(), RepositoryKey(param.RepoType, orgRepo)) + if err == nil { + return nil, localUploadError{status: 409, code: "REPOSITORY_EXISTS", msg: "repository already exists on this node"} + } + if !os.IsNotExist(err) { + return nil, err + } + } + if err := RegisterHosted(param.RepoType, param.Namespace, param.Repo); err != nil { + return nil, localUploadError{status: 409, code: "REPOSITORY_REGISTRATION_CONFLICT", msg: err.Error()} + } // 不共用就会各自基于陈旧清单重建，后写的覆盖先写的。 revisionLockKey := fmt.Sprintf("upload-revision:%s:%s:%s", param.RepoType, orgRepo, param.Revision) uploadRevisionLocks.Lock(revisionLockKey) @@ -405,7 +454,8 @@ func (u *UploadDao) PublishFiles(param LocalPublishParam) (*LocalPublishResult, result := &LocalPublishResult{ RepoType: param.RepoType, - Repo: orgRepo, + Namespace: param.Namespace, + Repo: param.Repo, Revision: param.Revision, Commit: commit, Published: len(param.Files), @@ -531,7 +581,8 @@ func (u *UploadDao) precheck(revisionLockKey string, param LocalUploadParam, org } return &LocalUploadResult{ RepoType: param.RepoType, - Repo: orgRepo, + Namespace: param.Namespace, + Repo: param.Repo, Revision: param.Revision, Commit: currentCommit, FilePath: param.FilePath, @@ -799,14 +850,15 @@ func hashDingCachePayload(path string, size int64) (string, error) { } func (u *UploadDao) QueryProgress(param LocalUploadParam) (*LocalUploadProgress, error) { - orgRepo := util.GetOrgRepo(param.Org, param.Repo) + orgRepo := util.GetOrgRepo(param.Namespace, param.Repo) blobPath := localBlobPath(param.RepoType, orgRepo, param.Sha256) if err := ensureLocalUploadPathSafe(config.SysConfig.Repos(), blobPath); err != nil { return nil, err } result := &LocalUploadProgress{ RepoType: param.RepoType, - Repo: orgRepo, + Namespace: param.Namespace, + Repo: param.Repo, Revision: param.Revision, FilePath: param.FilePath, Sha256: param.Sha256, @@ -958,13 +1010,16 @@ func (u *UploadDao) CleanupExpiredStagedUploads(retention time.Duration) (int, e } cutoff := time.Now().Add(-retention) removed := 0 - namespace := "dingo-local" - if config.SysConfig != nil && config.SysConfig.Upload.Namespace != "" { - namespace = config.SysConfig.Upload.Namespace + all, listErr := repository.ListHosted(config.SysConfig.Repos()) + if listErr != nil { + return 0, listErr } - filesRoot := filepath.Join(config.SysConfig.Repos(), "files") - for _, repoType := range []string{"models", "datasets"} { - root := filepath.Join(filesRoot, repoType, namespace) + for _, descriptor := range all { + if descriptor.Source != "hosted" || !descriptor.Persistent { + continue + } + repoType := descriptor.RepoType + root := filepath.Join(descriptor.FilesRoot(config.SysConfig.Repos()), "blobs") err := filepath.WalkDir(root, func(path string, entry fs.DirEntry, err error) error { if err != nil { if os.IsNotExist(err) { @@ -975,7 +1030,7 @@ func (u *UploadDao) CleanupExpiredStagedUploads(retention time.Duration) (int, e if entry.IsDir() || !strings.HasSuffix(entry.Name(), localUploadStageSuffix) { return nil } - parsedRepoType, orgRepo, sha, ok := stagedBlobIdentity(filesRoot, path) + parsedRepoType, orgRepo, sha, ok := stagedBlobIdentity(config.SysConfig.Repos(), path, descriptor) if !ok || parsedRepoType != repoType || !IsLocalOrgRepo(orgRepo) { return nil } @@ -1059,31 +1114,47 @@ type localRepoKey struct { // 那是在改动磁盘清理对公开模型的既有行为。 func collectLocalBlobs(root string) (map[localRepoKey][]string, error) { result := make(map[localRepoKey][]string) - err := filepath.WalkDir(root, func(path string, entry fs.DirEntry, err error) error { - if err != nil { - if os.IsNotExist(err) { - return nil - } - return err + repos := filepath.Dir(root) + all, err := repository.ListHosted(repos) + if err != nil { + return nil, err + } + for _, d := range all { + if d.Source != "hosted" || !d.Persistent { + continue } - if entry.IsDir() || strings.HasSuffix(entry.Name(), localUploadStageSuffix) { - return nil + blobRoot := filepath.Join(d.FilesRoot(repos), "blobs") + if err := repository.SafePath(repos, blobRoot); err != nil { + return nil, err } - repoType, orgRepo, sha, ok := stagedBlobIdentity(root, path) - if !ok || !IsLocalOrgRepo(orgRepo) { - return nil + entries, err := os.ReadDir(blobRoot) + if os.IsNotExist(err) { + continue + } + if err != nil { + return nil, err + } + for _, entry := range entries { + if entry.IsDir() || strings.HasSuffix(entry.Name(), localUploadStageSuffix) { + continue + } + if err := repository.Segment(entry.Name()); err != nil { + return nil, err + } + if entry.Type()&os.ModeSymlink != 0 { + return nil, fmt.Errorf("symlink in hosted blob directory") + } + key := localRepoKey{repoType: d.RepoType, orgRepo: d.ID()} + result[key] = append(result[key], entry.Name()) } - key := localRepoKey{repoType: repoType, orgRepo: orgRepo} - result[key] = append(result[key], sha) - return nil - }) - if os.IsNotExist(err) { - return result, nil } - return result, err + return result, nil } func (u *UploadDao) cleanupRepoBlobs(repos, repoType, orgRepo string, blobs []string, cutoff time.Time) (int, error) { + if err := validateRepositoryReferences(repoType, orgRepo); err != nil { + return 0, err + } // 仓库锁挡住并发的发布：引用集合在这把锁下算出来之后就不会再被追加引用。 repoLockKey := uploadRepoLockKey(repoType, orgRepo) uploadRepoLocks.Lock(repoLockKey) @@ -1172,6 +1243,9 @@ func (u *UploadDao) CleanupRecycledBlobs(retention time.Duration) (int, error) { } func (u *UploadDao) cleanupRecycledRepo(repoType, orgRepo string, cutoff time.Time) int { + if err := validateRepositoryReferences(repoType, orgRepo); err != nil { + return 0 + } tombstones := readRecycleEntries(repoType, orgRepo) if len(tombstones) == 0 { return 0 @@ -1209,7 +1283,7 @@ func (u *UploadDao) cleanupRecycledRepo(repoType, orgRepo string, cutoff time.Ti // 扫的是全部 commit 清单，不是版本标签当前指向的那一份。 func (u *UploadDao) referencedShas(repos, repoType, orgRepo string) (map[string]struct{}, error) { referenced := make(map[string]struct{}) - revisionRoot := filepath.Join(repos, "api", repoType, filepath.FromSlash(orgRepo), "revision") + revisionRoot := filepath.Join(RepositoryKey(repoType, orgRepo).APIRoot(repos), "revision") entries, err := os.ReadDir(revisionRoot) if err != nil { if os.IsNotExist(err) { @@ -1221,7 +1295,17 @@ func (u *UploadDao) referencedShas(repos, repoType, orgRepo string) (map[string] if !entry.IsDir() { continue } - for _, item := range u.readManifest(repoType, orgRepo, entry.Name()) { + manifest, readErr := readManifestFile(LocalManifestPath(repoType, orgRepo, entry.Name())) + if os.IsNotExist(readErr) { + commit := readMetaSha(filepath.Join(revisionRoot, entry.Name(), "meta_get.json")) + if commit != "" && commit != entry.Name() { + manifest, readErr = readManifestFile(LocalManifestPath(repoType, orgRepo, commit)) + } + } + if readErr != nil { + return nil, readErr + } + for _, item := range manifest { referenced[item.Sha256] = struct{}{} } } @@ -1250,32 +1334,26 @@ func (u *UploadDao) RunStagedUploadCleanup(ctx context.Context) { } } -func stagedBlobIdentity(root, stagePath string) (repoType, orgRepo, sha string, ok bool) { - rel, err := filepath.Rel(root, stagePath) - if err != nil || strings.HasPrefix(rel, "..") || filepath.IsAbs(rel) { +func stagedBlobIdentity(root, stagePath string, descriptor repository.Descriptor) (repoType, orgRepo, sha string, ok bool) { + // The cleanup already enumerated this repository; reuse its identity instead + // of scanning the registry again for every staged file. + d, rel, err := repository.LocateRegistered(root, stagePath, []repository.Descriptor{descriptor}) + if err != nil || d.Source != "hosted" || !d.Persistent { return "", "", "", false } - parts := strings.Split(filepath.ToSlash(rel), "/") - if len(parts) < 5 || parts[len(parts)-2] != "blobs" { + parts := strings.Split(rel, "/") + if len(parts) != 2 || parts[0] != "blobs" { return "", "", "", false } - repoType = parts[0] - if repoType != "models" && repoType != "datasets" { - return "", "", "", false - } - sha = strings.TrimSuffix(parts[len(parts)-1], localUploadStageSuffix) + sha = strings.TrimSuffix(parts[1], localUploadStageSuffix) if sha == "" { return "", "", "", false } - orgRepo = strings.Join(parts[1:len(parts)-2], "/") - if orgRepo == "" { - return "", "", "", false - } - return repoType, orgRepo, sha, true + return d.RepoType, d.ID(), sha, true } func (u *UploadDao) currentManifest(param LocalUploadParam) (string, []LocalManifestFile) { - return u.currentManifestOf(param.RepoType, util.GetOrgRepo(param.Org, param.Repo), param.Revision) + return u.currentManifestOf(param.RepoType, util.GetOrgRepo(param.Namespace, param.Repo), param.Revision) } func (u *UploadDao) currentManifestOf(repoType, orgRepo, revision string) (string, []LocalManifestFile) { @@ -1362,7 +1440,7 @@ func (u *UploadDao) readManifest(repoType, orgRepo, commit string) []LocalManife // 每次上传只写固定的 3 个文件，与该版本已有多少文件无关。 // 逐条落 paths-info、逐条建 resolve 链接会让逐个上传 N 个文件的代价变成 O(N²) // 个文件系统对象（N=1000 时约一百万个），且其中绝大多数永远不会被读到。 -func (u *UploadDao) writeEffectiveMetadata(repoType, orgRepo, revision, commit string, manifest []LocalManifestFile) error { +func (u *UploadDao) writeEffectiveMetadataRaw(repoType, orgRepo, revision, commit string, manifest []LocalManifestFile) error { manifestPath := LocalManifestPath(repoType, orgRepo, commit) if err := ensureLocalUploadPathSafe(config.SysConfig.Repos(), manifestPath); err != nil { return err @@ -1376,7 +1454,15 @@ func (u *UploadDao) writeEffectiveMetadata(repoType, orgRepo, revision, commit s if err := u.writeMeta(repoType, orgRepo, commit, commit, manifest); err != nil { return err } - return u.writeMeta(repoType, orgRepo, revision, commit, manifest) + if err := u.writeMeta(repoType, orgRepo, revision, commit, manifest); err != nil { + return err + } + for _, dir := range []string{filepath.Dir(manifestPath), RepositoryKey(repoType, orgRepo).Revision(config.SysConfig.Repos(), revision)} { + if err := inventory.SyncParents(config.SysConfig.Repos(), filepath.Join(dir, "meta_get.json")); err != nil { + return err + } + } + return nil } func (u *UploadDao) writeMeta(repoType, orgRepo, revision, commit string, manifest []LocalManifestFile) error { @@ -1400,7 +1486,7 @@ func (u *UploadDao) writeMeta(repoType, orgRepo, revision, commit string, manife "content-length": fmt.Sprintf("%d", len(body)), "x-repo-commit": commit, } - apiDir := filepath.Join(config.SysConfig.Repos(), "api", repoType, orgRepo, "revision", revision) + apiDir := RepositoryKey(repoType, orgRepo).Revision(config.SysConfig.Repos(), revision) metaGetPath := filepath.Join(apiDir, "meta_get.json") if err = ensureLocalUploadPathSafe(config.SysConfig.Repos(), metaGetPath); err != nil { return err @@ -1481,3 +1567,24 @@ func ensureLocalUploadPathSafe(root, target string) error { parent = filepath.Dir(parent) } } + +// VerifyPublishedFiles validates the manifest digest and actual local contents. +func VerifyPublishedFiles(repoType, id, commit string, files []LocalManifestFile) error { + computed, err := manifestCommit(files) + if err != nil { + return err + } + if computed != commit { + return fmt.Errorf("manifest digest mismatch") + } + return verifyManifestContent(repoType, id, files) +} + +var PublishedChanges = make(chan struct{}, 1) + +func NotifyPublished() { + select { + case PublishedChanges <- struct{}{}: + default: + } +} diff --git a/internal/dao/upload_dao_test.go b/internal/dao/upload_dao_test.go index f258345..28f72b4 100644 --- a/internal/dao/upload_dao_test.go +++ b/internal/dao/upload_dao_test.go @@ -45,13 +45,13 @@ func sha256Hex(b []byte) string { func uploadParam(filePath string, content []byte) LocalUploadParam { return LocalUploadParam{ - RepoType: "models", - Org: "dingo-local", - Repo: "demo", - Revision: "main", - FilePath: filePath, - Size: int64(len(content)), - Sha256: sha256Hex(content), + RepoType: "models", + Namespace: "dingo-local", + Repo: "demo", + Revision: "main", + FilePath: filePath, + Size: int64(len(content)), + Sha256: sha256Hex(content), } } @@ -473,6 +473,12 @@ func TestStagedUploadRetentionAndCleanup(t *testing.T) { func TestStagedUploadCleanupOnlyScansLocalNamespace(t *testing.T) { u, repos := newTestUploadDao(t) + if err := RegisterHosted("models", "dingo-local", "local-repo"); err != nil { + t.Fatal(err) + } + if err := RegisterHosted("datasets", "dingo-local", "local-dataset"); err != nil { + t.Fatal(err) + } stale := time.Now().Add(-48 * time.Hour) localPath := stagedBlobPath(repos, "models", "dingo-local/local-repo", "local-sha") @@ -819,7 +825,8 @@ func TestUploadMetadataFootprintStaysLinear(t *testing.T) { // 每次上传固定写 1 个 blob + 1 份清单 + 2 份 commit 元数据， // 外加全程复用的 2 份版本标签元数据。逐条落 paths-info / 建 resolve 链接的 // 老写法在这个规模下会产生一万多个文件。 - if want := 4*fileCount + 2; objects != want { + // One durable inventory journal outside the repository is constant overhead. + if want := 4*fileCount + 5; objects != want { t.Fatalf("metadata footprint is %d files, want %d (quadratic growth regression?)", objects, want) } } diff --git a/internal/dao/upload_publish_scale_test.go b/internal/dao/upload_publish_scale_test.go index 7bd87cf..326430e 100644 --- a/internal/dao/upload_publish_scale_test.go +++ b/internal/dao/upload_publish_scale_test.go @@ -232,7 +232,7 @@ func TestPublishScaleComparedToSequentialUploads(t *testing.T) { t.Fatalf("published manifest has %d files, want %d", published.FileCount, fileCount) } // 一次发布只写 1 份清单 + 2 份 commit 元数据 + 2 份版本标签元数据，与 N 无关。 - if batchObjects != 5 { + if batchObjects != 7 { t.Fatalf("batch publish wrote %d metadata objects, want 5 regardless of file count", batchObjects) } // 逐个上传每次重写全量清单，元数据字节数必然远大于批量。 diff --git a/internal/dao/upload_publish_test.go b/internal/dao/upload_publish_test.go index 16669ee..44da944 100644 --- a/internal/dao/upload_publish_test.go +++ b/internal/dao/upload_publish_test.go @@ -37,11 +37,11 @@ func mustStage(t *testing.T, u *UploadDao, param LocalUploadParam, content []byt func publishParam(revision string, files ...LocalManifestFile) LocalPublishParam { return LocalPublishParam{ - RepoType: "models", - Org: "dingo-local", - Repo: "demo", - Revision: revision, - Files: files, + RepoType: "models", + Namespace: "dingo-local", + Repo: "demo", + Revision: revision, + Files: files, } } @@ -666,8 +666,8 @@ func TestPublishMetadataFootprintIsConstant(t *testing.T) { } // 发布之前，磁盘上只有内容，没有任何元数据。 - if got := countFiles(t, repos); got != fileCount { - t.Fatalf("staging wrote %d files, want %d blobs and no metadata", got, fileCount) + if got := countFiles(t, repos); got != fileCount+2 { + t.Fatalf("staging wrote %d files, want %d blobs plus repository registration", got, fileCount) } result := mustPublish(t, u, publishParam("main", batch...)) @@ -675,7 +675,8 @@ func TestPublishMetadataFootprintIsConstant(t *testing.T) { t.Fatalf("manifest has %d files, want %d", result.FileCount, fileCount) } // N 个 blob + 1 份清单 + 2 份 commit 元数据 + 2 份版本标签元数据。 - if want := fileCount + 5; countFiles(t, repos) != want { + // The durable inventory journal adds one constant file per node. + if want := fileCount + 8; countFiles(t, repos) != want { t.Fatalf("metadata footprint is %d files, want %d", countFiles(t, repos), want) } if got := commitDirs(t, repos, "models", orgRepo); len(got) != 1 { diff --git a/internal/dao/upload_tree_dao.go b/internal/dao/upload_tree_dao.go index 7910871..28f2e2f 100644 --- a/internal/dao/upload_tree_dao.go +++ b/internal/dao/upload_tree_dao.go @@ -24,7 +24,7 @@ import ( // revision 的当前指向，否则说明别人已经改过，静默覆盖会丢掉对方的提交。 type LocalPublishTreeParam struct { RepoType string - Org string + Namespace string Repo string Revision string BaseCommit string @@ -32,10 +32,11 @@ type LocalPublishTreeParam struct { } type LocalPublishTreeResult struct { - RepoType string `json:"repoType"` - Repo string `json:"repo"` - Revision string `json:"revision"` - Commit string `json:"commit"` + Namespace string `json:"namespace"` + RepoType string `json:"repoType"` + Repo string `json:"repo"` + Revision string `json:"revision"` + Commit string `json:"commit"` // PreviousCommit 是本次提交之前 revision 的指向。 // // 调用方判断“改动是否生效”必须看它与 Commit 是否不同，不能看 Commit 是否变化： @@ -75,7 +76,12 @@ func supersededMarkerPath(repoType, orgRepo, commit string) string { // 与 PublishFiles 共用 manifestCommit 与 writeEffectiveMetadata：两条路径算出的 // 标识必须来自同一份序列化，否则同样的内容会得到两个标识，客户端会白白重下一遍。 func (u *UploadDao) PublishTree(param LocalPublishTreeParam) (*LocalPublishTreeResult, error) { - orgRepo := util.GetOrgRepo(param.Org, param.Repo) + release := holdRepository(param.RepoType, param.Namespace, param.Repo) + defer release() + if err := RegisterHosted(param.RepoType, param.Namespace, param.Repo); err != nil { + return nil, localUploadError{status: 409, code: "REPOSITORY_REGISTRATION_CONFLICT", msg: err.Error()} + } + orgRepo := util.GetOrgRepo(param.Namespace, param.Repo) // 锁序与 PublishFiles 完全一致：publish try-enter → 仓库锁 → 版本锁。 // 少一把或换个顺序都会与即时生效上传、回收任务交叉成死锁或脏读。 @@ -92,6 +98,9 @@ func (u *UploadDao) PublishTree(param LocalPublishTreeParam) (*LocalPublishTreeR repoLockKey := uploadRepoLockKey(param.RepoType, orgRepo) uploadRepoLocks.Lock(repoLockKey) defer uploadRepoLocks.Unlock(repoLockKey) + if err := u.recoverInventoryLocked(RepositoryKey(param.RepoType, orgRepo)); err != nil { + return nil, err + } revisionLockKey := fmt.Sprintf("upload-revision:%s:%s:%s", param.RepoType, orgRepo, param.Revision) uploadRevisionLocks.Lock(revisionLockKey) @@ -132,7 +141,8 @@ func (u *UploadDao) PublishTree(param LocalPublishTreeParam) (*LocalPublishTreeR added, replaced, unchanged, removed := diffManifest(currentManifest, manifest) result := &LocalPublishTreeResult{ RepoType: param.RepoType, - Repo: orgRepo, + Namespace: param.Namespace, + Repo: param.Repo, Revision: param.Revision, Commit: commit, PreviousCommit: currentCommit, diff --git a/internal/dao/upload_tree_dao_test.go b/internal/dao/upload_tree_dao_test.go index 6299f0d..43469f0 100644 --- a/internal/dao/upload_tree_dao_test.go +++ b/internal/dao/upload_tree_dao_test.go @@ -15,7 +15,7 @@ import ( func treeParam(revision, baseCommit string, files ...LocalManifestFile) LocalPublishTreeParam { return LocalPublishTreeParam{ RepoType: "models", - Org: "dingo-local", + Namespace: "dingo-local", Repo: "demo", Revision: revision, BaseCommit: baseCommit, diff --git a/internal/data/global.go b/internal/data/global.go index b80f753..ea2b05f 100644 --- a/internal/data/global.go +++ b/internal/data/global.go @@ -12,8 +12,9 @@ import ( ) type LocalOperation struct { - Type int `json:"type"` - Body string `json:"body"` + IdentityVersion int `json:"identityVersion,omitempty"` + Type int `json:"type"` + Body string `json:"body"` } type FileProcessParam struct { @@ -73,8 +74,9 @@ func WriteLocalOperationChan(operationType int, body interface{}) { return } opera := &LocalOperation{ - Type: operationType, - Body: string(marshal), + IdentityVersion: 2, + Type: operationType, + Body: string(marshal), } select { case localOperationChan <- opera: diff --git a/internal/downloader/cache_completion_test.go b/internal/downloader/cache_completion_test.go new file mode 100644 index 0000000..cb03981 --- /dev/null +++ b/internal/downloader/cache_completion_test.go @@ -0,0 +1,59 @@ +package downloader + +import ( + "context" + "dingospeed/pkg/config" + "path/filepath" + "testing" +) + +func TestCacheCompletionRequiresWholeReadableFile(t *testing.T) { + old := config.SysConfig + defer func() { config.SysConfig = old }() + config.SysConfig = &config.Config{} + for _, tc := range []struct { + name string + start, end int64 + missing, canceled bool + want int + }{ + {name: "complete", end: 7, want: 1}, + {name: "prefix", end: 4}, + {name: "suffix", start: 4, end: 7}, + {name: "missing_tail", end: 7, missing: true}, + {name: "canceled", end: 7, canceled: true}, + } { + t.Run(tc.name, func(t *testing.T) { + cache, err := NewDingCache(filepath.Join(t.TempDir(), "blob"), 4) + if err != nil { + t.Fatal(err) + } + defer cache.Close() + if err = cache.Resize(7); err != nil { + t.Fatal(err) + } + if err = cache.WriteBlock(0, []byte("abcd")); err != nil { + t.Fatal(err) + } + if !tc.missing { + if err = cache.WriteBlock(1, []byte("efg\x00")); err != nil { + t.Fatal(err) + } + } + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + if tc.canceled { + cancel() + } + task := NewCacheFileTask(0, tc.start, tc.end) + task.Context, task.DingFile = ctx, cache + task.ResponseChan = make(chan []byte, 4) + calls := 0 + task.OnComplete = func() { calls++ } + task.OutResult() + if calls != tc.want { + t.Fatalf("completion calls=%d want=%d", calls, tc.want) + } + }) + } +} diff --git a/internal/downloader/cache_task.go b/internal/downloader/cache_task.go index eeedf74..240a2c6 100644 --- a/internal/downloader/cache_task.go +++ b/internal/downloader/cache_task.go @@ -16,29 +16,46 @@ package downloader import ( "context" + "dingospeed/pkg/repository" + "net/http" "go.uber.org/zap" ) +// RemoteSource supplies provider-specific transport; the cache/task machinery is shared. +type RemoteSource struct { + Domain string + Headers map[string]string + Fetch func(context.Context, string, string, map[string]string, func(*http.Response) error) error +} + type TaskParam struct { - Context context.Context - DingFile *DingCache - TaskNo int - TaskSize int - BlobsFile string - FileName string - FileSize int64 - ResponseChan chan []byte - OrgRepo string - Authorization string - Domain string - Uri string - DataType string - Etag string - Cancel context.CancelFunc + Source *RemoteSource + LocalOnly bool // Skip scheduler and reporting until this provider's cluster integration is enabled. + CacheResult chan error // Optional buffered terminal result, sent after cache writes finish. + OnCacheComplete func() + Peer bool + Revision string + RepoKey repository.RepoKey + Context context.Context + DingFile *DingCache + TaskNo int + TaskSize int + BlobsFile string + FileName string + FileSize int64 + ResponseChan chan []byte + OrgRepo string + Authorization string + Domain string + Uri string + DataType string + Etag string + Cancel context.CancelFunc } type DownloadTask struct { + RepoKey repository.RepoKey TaskNo int RangeStartPos int64 RangeEndPos int64 @@ -65,6 +82,7 @@ func (d *DownloadTask) GetCancelFun() context.CancelFunc { type CacheFileTask struct { *DownloadTask + OnComplete func() } func NewCacheFileTask(taskNo int, rangeStartPos int64, rangeEndPos int64) *CacheFileTask { @@ -128,6 +146,12 @@ func (c *CacheFileTask) OutResult() { } if curPos != c.RangeEndPos { zap.S().Errorf("file:%s, cache range from %d to %d is incomplete.", c.FileName, c.RangeStartPos, c.RangeEndPos) + return + } + // Reconcile only after reading the whole cached file successfully. A cached + // subrange is not evidence that the file's missing prefix or suffix exists. + if c.RangeStartPos == 0 && c.RangeEndPos == c.DingFile.GetFileSize() && c.OnComplete != nil { + c.OnComplete() } zap.S().Infof("cache out:%s/%s, taskNo:%d, size:%d, startPos:%d, endPos:%d", c.OrgRepo, c.FileName, c.TaskNo, c.TaskSize, c.RangeStartPos, c.RangeEndPos) } diff --git a/internal/downloader/file.go b/internal/downloader/file.go index 1d6cf1d..20cb13e 100644 --- a/internal/downloader/file.go +++ b/internal/downloader/file.go @@ -335,10 +335,23 @@ func (c *DingCache) WriteBlock(blockIndex int64, blockBytes []byte) error { // 加锁顺序固定为 fileLock → headerLock，与 Resize 一致，不会反向获取。 c.headerLock.Lock() defer c.headerLock.Unlock() + return c.commitHeaderBlock(blockIndex) +} + +// Caller holds fileLock and headerLock. A failed header write must not leave +// a newly cached block visible to subsequent readers in memory. +func (c *DingCache) commitHeaderBlock(blockIndex int64) error { + wasSet, err := c.header.BlockMask.Test(uint64(blockIndex)) + if err != nil { + return err + } if err = c.setHeaderBlock(blockIndex); err != nil { return err } if err = c.flushHeader(); err != nil { + if !wasSet { + _ = c.header.BlockMask.Clear(uint64(blockIndex)) + } return err } // key := c.getBlockKey(blockIndex) 不需要删除，本来就没有 diff --git a/internal/downloader/remote_task.go b/internal/downloader/remote_task.go index 7655547..18015d2 100644 --- a/internal/downloader/remote_task.go +++ b/internal/downloader/remote_task.go @@ -29,12 +29,17 @@ import ( "dingospeed/pkg/consts" myerr "dingospeed/pkg/error" "dingospeed/pkg/prom" + "dingospeed/pkg/repository" "dingospeed/pkg/util" "go.uber.org/zap" ) type RemoteFileTask struct { + Source *RemoteSource + LocalOnly bool + Peer bool + UpstreamURI string *DownloadTask Authorization string Domain string @@ -43,6 +48,7 @@ type RemoteFileTask struct { Etag string Queue chan []byte `json:"-"` Cancel context.CancelFunc + ResultError error } func NewRemoteFileTask(taskNo int, rangeStartPos int64, rangeEndPos int64) *RemoteFileTask { @@ -60,6 +66,8 @@ func (r *RemoteFileTask) DoTask() { curBlock int64 wg sync.WaitGroup streamCache = bytes.Buffer{} + fetchErr error + cacheErr error ) contentChan := make(chan []byte, consts.RespChanSize) rangeStartPos, rangeEndPos := r.RangeStartPos, r.RangeEndPos @@ -67,11 +75,11 @@ func (r *RemoteFileTask) DoTask() { wg.Add(2) go func() { defer wg.Done() - if err := r.getFileRangeFromRemote(rangeStartPos, rangeEndPos, contentChan); err != nil { + defer close(contentChan) + fetchErr = r.getFileRangeFromRemote(rangeStartPos, rangeEndPos, contentChan) + if err := fetchErr; err != nil { zap.S().Errorf("getFileRangeFromRemote err.%v", err) r.Cancel() - } else { - close(contentChan) } }() curPos, lastReportPos := rangeStartPos, rangeStartPos @@ -94,7 +102,6 @@ func (r *RemoteFileTask) DoTask() { case r.Queue <- chunk: case <-r.Context.Done(): zap.S().Warnf("send chunk err:%s/%s, task %d, ctx done, DoTask exit.", r.OrgRepo, r.FileName, r.TaskNo) - data.ReportFileProcess(r.Context, r.constructFileProcessParam(lastReportPos, lastBlockEndPos, consts.StatusDownloadBreak)) return } chunkLen := int64(len(chunk)) @@ -105,12 +112,15 @@ func (r *RemoteFileTask) DoTask() { prom.PromRequestByteCounter(prom.RequestRemoteByte, source, r.OrgRepo, r.Domain, chunkLen) } + if cacheErr != nil { + continue // Cache failure must not cancel delivery of upstream bytes. + } if len(chunk) != 0 { streamCache.Write(chunk) } curBlock = curPos / r.DingFile.GetBlockSize() // 若是一个新的数据块，则将上一个数据块持久化。 - if curBlock != lastBlock { + for curBlock > lastBlock && (lastBlock+1)*r.DingFile.GetBlockSize() <= r.DingFile.GetFileSize() && cacheErr == nil { splitPos := lastBlockEndPos - max(lastBlockStartPos, rangeStartPos) cacheLen := int64(streamCache.Len()) if splitPos > cacheLen { @@ -124,13 +134,16 @@ func (r *RemoteFileTask) DoTask() { hasBlockBool, err := r.DingFile.HasBlock(lastBlock) if err != nil { zap.S().Errorf("HasBlock err.%v", err) + cacheErr = err } if err == nil && !hasBlockBool { - if err = r.DingFile.WriteBlock(lastBlock, rawBlock); err != nil { + if err = r.writeCacheBlock(lastBlock, rawBlock); err != nil { zap.S().Errorf("writeBlock err.%v", err) + cacheErr = err + continue } zap.S().Debugf("from:%s, %s/%s, taskNo:%d, block：%d(%d)write done, range：%d-%d.", r.Domain, r.OrgRepo, r.FileName, r.TaskNo, lastBlock, blockNumber, lastBlockStartPos, lastBlockEndPos) - if interval == config.SysConfig.GetSyncProcessInterval() { + if !r.LocalOnly && interval == config.SysConfig.GetSyncProcessInterval() { data.ReportFileProcess(r.Context, r.constructFileProcessParam(lastReportPos, lastBlockEndPos, consts.StatusDownloading)) lastReportPos = lastBlockEndPos interval = 1 @@ -142,17 +155,37 @@ func (r *RemoteFileTask) DoTask() { nextBlock := streamCacheBytes[splitPos:] // 下一个块的数据 streamCache.Truncate(0) streamCache.Write(nextBlock) - lastBlock, lastBlockStartPos, lastBlockEndPos = GetBlockInfo(curPos, r.DingFile.GetBlockSize(), r.DingFile.GetFileSize()) + lastBlock, lastBlockStartPos, lastBlockEndPos = GetBlockInfo(lastBlockEndPos, r.DingFile.GetBlockSize(), r.DingFile.GetFileSize()) } } case <-r.Context.Done(): zap.S().Warnf("file:%s/%s taskNo:%d ctx done, DoTask exit.", r.OrgRepo, r.FileName, r.TaskNo) - data.ReportFileProcess(r.Context, r.constructFileProcessParam(lastReportPos, lastBlockEndPos, consts.StatusDownloadBreak)) return } } }() wg.Wait() + defer func() { + if cacheErr != nil { + r.ResultError = cacheErr + } else if fetchErr != nil { + r.ResultError = fetchErr + } + }() + // Report exactly one terminal outcome, after both transfer and cache writes finish. + completed := false + defer func() { + status := int32(consts.StatusDownloadBreak) + if completed { + status = consts.StatusDownloaded + } + if !r.LocalOnly { + data.ReportFileProcess(r.Context, r.constructFileProcessParam(lastReportPos, rangeEndPos, status)) + } + }() + if fetchErr != nil || cacheErr != nil || curPos != rangeEndPos { + return + } rawBlock := streamCache.Bytes() if curBlock == r.DingFile.getBlockNumber()-1 { // 对不足一个block的数据做补全 @@ -162,32 +195,40 @@ func (r *RemoteFileTask) DoTask() { } lastBlock = curBlock } - // 一个空文件，或文件刚好为blocksize的整数倍，直接标记为完成 - if len(rawBlock) == 0 { - data.ReportFileProcess(r.Context, r.constructFileProcessParam(lastReportPos, curPos, consts.StatusDownloaded)) - } else if int64(len(rawBlock)) == r.DingFile.GetBlockSize() { + // Commit the remaining block before checking cache coverage. + if int64(len(rawBlock)) == r.DingFile.GetBlockSize() { hasBlockBool, err := r.DingFile.HasBlock(lastBlock) if err != nil { + cacheErr = err zap.S().Errorf("HasBlock err.%v", err) return } if !hasBlockBool { - if err = r.DingFile.WriteBlock(lastBlock, rawBlock); err != nil { + if err = r.writeCacheBlock(lastBlock, rawBlock); err != nil { + cacheErr = err zap.S().Errorf("last writeBlock err.%v", err) + return } zap.S().Debugf("from:%s, %s/%s, taskNo:%d, last block：%d(%d)write done, range：%d-%d.", r.Domain, r.OrgRepo, r.FileName, r.TaskNo, lastBlock, blockNumber, lastBlockStartPos, lastBlockEndPos) - data.ReportFileProcess(r.Context, r.constructFileProcessParam(lastReportPos, curPos, consts.StatusDownloaded)) } } - if curPos != rangeEndPos { - zap.S().Errorf("file:%s/%s, taskNo:%d, remote range (%d) is different from sent size (%d).", r.OrgRepo, r.FileName, r.TaskNo, rangeEndPos-rangeStartPos, curPos-rangeStartPos) - return + // Receiving all bytes alone does not establish that the requested cache range exists. + for block := rangeStartPos / r.DingFile.GetBlockSize(); block*r.DingFile.GetBlockSize() < rangeEndPos; block++ { + has, err := r.DingFile.HasBlock(block) + if err != nil || !has { + zap.S().Errorf("remote cache incomplete: block=%d, err=%v", block, err) + return + } } + // A client may close its connection after receiving Content-Length bytes while + // the final cache write is still running. Completed cache coverage survives + // that cancellation; incomplete transfers and failed writes are rejected above. + completed = true zap.S().Infof("end remote dotask:%s/%s, taskNo:%d, size:%d, domain:%s, startPos:%d, endPos:%d", r.OrgRepo, r.FileName, r.TaskNo, r.TaskSize, r.Domain, rangeStartPos, rangeEndPos) } func (r *RemoteFileTask) constructFileProcessParam(startPos, endPos int64, status int32) *data.FileProcessParam { - org, repo := util.SplitOrgRepo(r.OrgRepo) + org, repo := r.RepoKey.Namespace, r.RepoKey.Repo return &data.FileProcessParam{ Datatype: r.DataType, Org: org, @@ -236,6 +277,11 @@ func (r *RemoteFileTask) getFileRangeFromRemote(startPos, endPos int64, contentC n int headers = make(map[string]string) ) + if r.Source != nil { + for k, v := range r.Source.Headers { + headers[k] = v + } + } if r.Authorization != "" { headers["authorization"] = r.Authorization } @@ -243,8 +289,29 @@ func (r *RemoteFileTask) getFileRangeFromRemote(startPos, endPos int64, contentC headers["range"] = fmt.Sprintf("bytes=%d-%d", startPos, endPos-1) } for i := 0; i < attempts; { - if _, err = util.RetryRequest(func() (*common.Response, error) { - err = util.GetStream(r.Domain, r.Uri, headers, func(resp *http.Response) error { + if err = util.RetryDownloadContext(r.Context, func() error { + trace := common.CacheTrace(r.Context) + if !trace.Begin() { + return context.Canceled + } + defer func() { trace.End(r.Context.Err() != nil) }() + + fetch := func(domain, uri string, headers map[string]string, cb func(*http.Response) error) error { + return util.GetStreamContext(r.Context, domain, uri, headers, cb) + } + if r.Peer { + fetch = func(domain, uri string, headers map[string]string, cb func(*http.Response) error) error { + return util.GetPeerStreamContext(r.Context, domain, uri, headers, cb) + } + } else if r.Source != nil && r.Source.Fetch != nil { + fetch = func(domain, uri string, headers map[string]string, cb func(*http.Response) error) error { + return r.Source.Fetch(r.Context, domain, uri, headers, cb) + } + } + err = fetch(r.Domain, r.Uri, headers, func(resp *http.Response) error { + if headers["range"] != "" && resp.StatusCode == http.StatusOK { + return fmt.Errorf("source ignored Range request") + } contentEncoding = resp.Header.Get("content-encoding") code := resp.StatusCode if code != http.StatusOK && code != http.StatusPartialContent { @@ -255,7 +322,7 @@ func (r *RemoteFileTask) getFileRangeFromRemote(startPos, endPos int64, contentC } else { zap.S().Errorf("Failed resource request.(%d) %s", code, r.OrgRepo) } - return nil + return myerr.NewAppendCode(code, fmt.Sprintf("upstream returned HTTP %d", code)) } for { select { @@ -285,7 +352,7 @@ func (r *RemoteFileTask) getFileRangeFromRemote(startPos, endPos int64, contentC // 数据不完整，将EOF视为读取错误以触发重试/断点续传 zap.S().Errorf("file:%s/%s, taskNo:%d, premature EOF: expected %d bytes, got %d", r.OrgRepo, r.FileName, r.TaskNo, endPos-startPos, chunkByteLen) headers["range"] = fmt.Sprintf("bytes=%d-%d", startPos+int64(chunkByteLen), endPos-1) - return fmt.Errorf("premature EOF: expected %d bytes, got %d", endPos-startPos, chunkByteLen) + return fmt.Errorf("premature EOF: expected %d bytes, got %d: %w", endPos-startPos, chunkByteLen, io.ErrUnexpectedEOF) } return nil } @@ -298,17 +365,25 @@ func (r *RemoteFileTask) getFileRangeFromRemote(startPos, endPos int64, contentC } } }) - return nil, err + return err }); err != nil { + if r.Context.Err() != nil { + return r.Context.Err() + } var t myerr.Error if errors.As(err, &t) { break } - // 若从内部其他节点获取数据出现异常，则切换到官网获取。 - if config.SysConfig.IsCluster() && util.IsInnerDomain(r.Domain) { + // 若从内部其他节点获取数据出现异常，则切换到对应 provider 的官网获取。 + if r.Peer && (r.RepoKey.Namespace == repository.HuggingFace || r.Source != nil) { officialDomain := config.SysConfig.GetHFURLBase() + if r.Source != nil { + officialDomain = r.Source.Domain + } zap.S().Infof("request fail %s/%s req from %s to %s", r.OrgRepo, r.FileName, r.Domain, officialDomain) r.Domain = officialDomain + r.Uri = r.UpstreamURI + r.Peer = false if chunkByteLen > 0 { headers["range"] = fmt.Sprintf("bytes=%d-%d", startPos+int64(chunkByteLen), endPos-1) } @@ -321,7 +396,7 @@ func (r *RemoteFileTask) getFileRangeFromRemote(startPos, endPos int64, contentC } } if err != nil { - return fmt.Errorf("GetStream err.%v", err) + return fmt.Errorf("GetStream: %w", err) } if contentEncoding != "" { // 这里需要实现解压缩逻辑 @@ -330,7 +405,11 @@ func (r *RemoteFileTask) getFileRangeFromRemote(startPos, endPos int64, contentC zap.S().Errorf("DecompressData err.%v", err) return err } - contentChan <- finalData // 返回解码后的数据流 + select { + case contentChan <- finalData: + case <-r.Context.Done(): + return r.Context.Err() + } chunkByteLen = len(finalData) // 将解码后的长度复制为原始的chunkByteLen } expectedLength := endPos - startPos @@ -339,3 +418,11 @@ func (r *RemoteFileTask) getFileRangeFromRemote(startPos, endPos int64, contentC } return nil } + +func (r *RemoteFileTask) writeCacheBlock(block int64, data []byte) error { + if err := r.DingFile.WriteBlock(block, data); err != nil { + return err + } + common.CacheTrace(r.Context).Wrote(int64(len(data))) + return nil +} diff --git a/internal/downloader/remote_task_completion_test.go b/internal/downloader/remote_task_completion_test.go new file mode 100644 index 0000000..0860d4a --- /dev/null +++ b/internal/downloader/remote_task_completion_test.go @@ -0,0 +1,174 @@ +package downloader + +import ( + "context" + "net/http" + "net/http/httptest" + "path/filepath" + "testing" + "time" + + "dingospeed/internal/data" + "dingospeed/pkg/config" + "dingospeed/pkg/consts" +) + +func TestRemoteTaskCacheCompletion(t *testing.T) { + old := config.SysConfig + t.Cleanup(func() { config.SysConfig = old }) + for _, tc := range []struct { + name string + size int64 + chunk int64 + failWrite, truncated, canceled, cached bool + failLookup, oversized bool + }{ + {name: "full_blocks", size: 8, chunk: 4}, + {name: "tail", size: 7, chunk: 4}, + {name: "multiple_blocks_per_chunk", size: 15, chunk: 32}, + {name: "already_cached", size: 7, chunk: 4, cached: true}, + {name: "full_block_write_failure", size: 8, chunk: 4, failWrite: true}, + {name: "tail_write_failure", size: 3, chunk: 4, failWrite: true}, + {name: "block_lookup_failure", size: 8, chunk: 4, failLookup: true}, + {name: "oversized", size: 7, chunk: 32, oversized: true}, + {name: "truncated", size: 8, chunk: 4, truncated: true}, + {name: "canceled", size: 8, chunk: 4, canceled: true}, + } { + t.Run(tc.name, func(t *testing.T) { + config.SysConfig = &config.Config{} + config.SysConfig.Download.RespChunkSize = tc.chunk + config.SysConfig.Retry.Attempts = 1 + config.SysConfig.Scheduler.Mode = consts.SchedulerModeCluster + config.SysConfig.Scheduler.OriginMode = consts.SchedulerModeCluster + config.SysConfig.Scheduler.Strategy.SyncProcessInterval = 1 + data.NewBaseData() // Initializes in-memory reporting only; no scheduler connection. + cache, err := NewDingCache(filepath.Join(t.TempDir(), "cache"), 4) + if err != nil { + t.Fatal(err) + } + defer cache.Close() + if err = cache.Resize(tc.size); err != nil { + t.Fatal(err) + } + payload := "abcdefghijklmno"[:tc.size] + if tc.cached { + for b := int64(0); b < cache.getBlockNumber(); b++ { + block := make([]byte, 4) + copy(block, payload[b*4:min((b+1)*4, tc.size)]) + if err = cache.WriteBlock(b, block); err != nil { + t.Fatal(err) + } + } + } + if tc.failWrite { + cache.path = filepath.Join(t.TempDir(), "missing", "cache") + } + if tc.failLookup { + cache.header.BlockMask = NewBitset(0) + } + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) { + if tc.oversized { + _, _ = w.Write([]byte(payload + "extra")) + return + } + if tc.truncated { + w.Header().Set("Content-Length", "8") + _, _ = w.Write([]byte("abcd")) + return + } + _, _ = w.Write([]byte(payload)) + })) + defer server.Close() + ctx, cancel := context.WithCancel(context.WithValue(context.Background(), consts.KeyProcessId, int64(1))) + defer cancel() + if tc.canceled { + cancel() + } + task := NewRemoteFileTask(0, 0, tc.size) + task.Context, task.Cancel, task.DingFile = ctx, cancel, cache + task.Domain, task.Uri = server.URL, "/file" + task.Queue = make(chan []byte, 32) + done := make(chan struct{}) + go func() { task.DoTask(); close(done) }() + select { + case <-done: + case <-time.After(5 * time.Second): + t.Fatal("task did not terminate") + } + var received []byte + for chunk := range task.Queue { + received = append(received, chunk...) + } + if !tc.truncated && !tc.canceled && !tc.oversized { + if string(received) != payload { + t.Fatalf("delivery = %q, want %q", received, payload) + } + if ctx.Err() != nil { + t.Fatalf("cache failure canceled delivery: %v", ctx.Err()) + } + } + want := int32(consts.StatusDownloaded) + if tc.failWrite || tc.failLookup || tc.truncated || tc.canceled || tc.oversized { + want = consts.StatusDownloadBreak + } + terminal := 0 + for len(data.GetFileProcessChan()) > 0 { + report := <-data.GetFileProcessChan() + if report.Status != consts.StatusDownloading { + terminal++ + if report.Status != want { + t.Fatalf("terminal = %d, want %d", report.Status, want) + } + } else if tc.failWrite || tc.failLookup { + t.Fatal("failed write reported as cached progress") + } + } + if terminal != 1 { + t.Fatalf("terminal reports = %d", terminal) + } + if want == consts.StatusDownloaded { + for b := int64(0); b < cache.getBlockNumber(); b++ { + block, err := cache.ReadBlock(b) + if err != nil { + t.Fatal(err) + } + expected := payload[b*4 : min((b+1)*4, tc.size)] + if string(block[:len(expected)]) != expected { + t.Fatalf("block %d = %q", b, block) + } + } + } + }) + } +} + +func TestFailedHeaderCommitDoesNotPublishBlock(t *testing.T) { + cache, err := NewDingCache(filepath.Join(t.TempDir(), "cache"), 4) + if err != nil { + t.Fatal(err) + } + defer cache.Close() + if err = cache.Resize(4); err != nil { + t.Fatal(err) + } + originalPath := cache.path + cache.path = filepath.Join(t.TempDir(), "missing", "cache") + cache.fileLock.Lock() + cache.headerLock.Lock() + err = cache.commitHeaderBlock(0) + cache.headerLock.Unlock() + cache.fileLock.Unlock() + if err == nil { + t.Fatal("expected header write failure") + } + if has, err := cache.HasBlock(0); err != nil || has { + t.Fatalf("failed header commit published block: has=%v err=%v", has, err) + } + cache.path = originalPath + if err = cache.WriteBlock(0, []byte("abcd")); err != nil { + t.Fatal(err) + } + if has, err := cache.HasBlock(0); err != nil || !has { + t.Fatalf("retry did not publish block: has=%v err=%v", has, err) + } +} diff --git a/internal/downloader/remote_task_disconnect_test.go b/internal/downloader/remote_task_disconnect_test.go new file mode 100644 index 0000000..4933159 --- /dev/null +++ b/internal/downloader/remote_task_disconnect_test.go @@ -0,0 +1,118 @@ +package downloader + +import ( + "context" + "dingospeed/internal/data" + "dingospeed/pkg/config" + "dingospeed/pkg/consts" + "io" + "net/http" + "net/http/httptest" + "path/filepath" + "testing" + "time" +) + +// Regression test: a real HTTP client receives Content-Length bytes while +// the final cache write is held behind the existing file lock. No mocked context. +func TestHTTPDisconnectAfterFullDelivery(t *testing.T) { + old := config.SysConfig + defer func() { config.SysConfig = old }() + for _, disconnect := range []bool{false, true} { + name := "keep_connection" + if disconnect { + name = "close_after_full_body" + } + t.Run(name, func(t *testing.T) { + config.SysConfig = &config.Config{} + config.SysConfig.Download.RespChunkSize = 8 + config.SysConfig.Retry.Attempts = 1 + config.SysConfig.Scheduler.Mode = consts.SchedulerModeCluster + config.SysConfig.Scheduler.OriginMode = consts.SchedulerModeCluster + data.NewBaseData() + cache, err := NewDingCache(filepath.Join(t.TempDir(), "blob"), 16) + if err != nil { + t.Fatal(err) + } + defer cache.Close() + if err = cache.Resize(7); err != nil { + t.Fatal(err) + } + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { _, _ = w.Write([]byte("abcdefg")) })) + defer upstream.Close() + cache.fileLock.Lock() + locked := true + defer func() { + if locked { + cache.fileLock.Unlock() + } + }() + taskDone := make(chan struct{}) + requestCanceled := make(chan struct{}) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + ctx, cancel := context.WithCancel(context.WithValue(r.Context(), consts.KeyProcessId, int64(99))) + defer cancel() + go func() { <-ctx.Done(); close(requestCanceled) }() + task := NewRemoteFileTask(0, 0, 7) + task.Context = ctx + task.Cancel = cancel + task.DingFile = cache + task.Domain = upstream.URL + task.Uri = "/file" + task.Queue = make(chan []byte, 32) + go func() { task.DoTask(); close(taskDone) }() + w.Header().Set("Content-Length", "7") + for b := range task.Queue { + _, _ = w.Write(b) + w.(http.Flusher).Flush() + } + <-taskDone // handler deliberately stays alive while final cache commit waits + })) + defer server.Close() + tr := &http.Transport{DisableKeepAlives: disconnect} + defer tr.CloseIdleConnections() + client := &http.Client{Transport: tr, Timeout: 5 * time.Second} + resp, err := client.Get(server.URL) + if err != nil { + t.Fatal(err) + } + body, err := io.ReadAll(resp.Body) + resp.Body.Close() + if err != nil || string(body) != "abcdefg" { + t.Fatalf("body=%q err=%v", body, err) + } + if disconnect { + select { + case <-requestCanceled: + case <-time.After(3 * time.Second): + t.Fatal("client close did not cancel request") + } + } + cache.fileLock.Unlock() + locked = false + select { + case <-taskDone: + case <-time.After(3 * time.Second): + t.Fatal("task did not finish") + } + has, err := cache.HasBlock(0) + if err != nil || !has { + t.Fatalf("cache incomplete: %v %v", has, err) + } + block, err := cache.ReadBlock(0) + if err != nil || string(block[:7]) != "abcdefg" { + t.Fatalf("disk mismatch %v", err) + } + select { + case progress := <-data.GetFileProcessChan(): + want := int32(consts.StatusDownloaded) + t.Logf("all 7 response bytes verified; persisted cache verified; client_disconnect=%v; reported_status=%d", disconnect, progress.Status) + if progress.Status != want { + t.Fatalf("diagnosed status=%d want=%d", progress.Status, want) + } + case <-time.After(time.Second): + t.Fatal("no report") + } + }) + } +} diff --git a/internal/downloader/remote_task_peer_fallback_test.go b/internal/downloader/remote_task_peer_fallback_test.go new file mode 100644 index 0000000..bd668c0 --- /dev/null +++ b/internal/downloader/remote_task_peer_fallback_test.go @@ -0,0 +1,93 @@ +package downloader + +import ( + "context" + "net/http" + "net/http/httptest" + "path/filepath" + "testing" + + "dingospeed/pkg/config" + "dingospeed/pkg/repository" +) + +func TestModelScopePeerFailureFallsBackToOfficialFromDownloadedOffset(t *testing.T) { + old := config.SysConfig + defer func() { config.SysConfig = old }() + config.SysConfig = &config.Config{} + config.SysConfig.Download.RespChunkSize = 16 + config.SysConfig.Retry.Attempts = 1 + + cache, err := NewDingCache(filepath.Join(t.TempDir(), "blob"), 16) + if err != nil { + t.Fatal(err) + } + defer cache.Close() + if err = cache.Resize(10); err != nil { + t.Fatal(err) + } + + peer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Length", "10") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte("abc")) + })) + defer peer.Close() + + var gotRange string + official := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotRange = r.Header.Get("Range") + w.WriteHeader(http.StatusPartialContent) + _, _ = w.Write([]byte("defghij")) + })) + defer official.Close() + + source := &RemoteSource{ + Domain: official.URL, + Fetch: func(ctx context.Context, domain, uri string, headers map[string]string, consume func(*http.Response) error) error { + req, err := http.NewRequestWithContext(ctx, http.MethodGet, domain+uri, nil) + if err != nil { + return err + } + for key, value := range headers { + req.Header.Set(key, value) + } + resp, err := http.DefaultClient.Do(req) + if err != nil { + return err + } + defer resp.Body.Close() + return consume(resp) + }, + } + task := NewRemoteFileTask(0, 0, 10) + task.Context = context.Background() + task.DingFile = cache + task.RepoKey = repository.RepoKey{Namespace: repository.ModelScope, RepoType: "models", Repo: "Qwen/demo"} + task.OrgRepo = task.RepoKey.ID() + task.FileName = "model.bin" + task.Domain = peer.URL + task.Uri = "/api/repositories/models/modelscope/file" + task.UpstreamURI = "/api/v1/models/Qwen/demo/repo" + task.Peer = true + task.Source = source + + content := make(chan []byte, 2) + if err = task.getFileRangeFromRemote(0, 10, content); err != nil { + t.Fatal(err) + } + close(content) + var body []byte + for chunk := range content { + body = append(body, chunk...) + } + if string(body) != "abcdefghij" { + t.Fatalf("body=%q", body) + } + if gotRange != "bytes=3-9" { + t.Fatalf("official Range=%q, want %q", gotRange, "bytes=3-9") + } + if task.Domain != official.URL || task.Uri != task.UpstreamURI || task.Peer { + t.Fatalf("fallback state domain=%q uri=%q peer=%v", task.Domain, task.Uri, task.Peer) + } +} diff --git a/internal/downloader/remote_task_retry_test.go b/internal/downloader/remote_task_retry_test.go new file mode 100644 index 0000000..b43722c --- /dev/null +++ b/internal/downloader/remote_task_retry_test.go @@ -0,0 +1,86 @@ +package downloader + +import ( + "context" + "fmt" + "net/http" + "net/http/httptest" + "path/filepath" + "sync/atomic" + "testing" + "time" + + "dingospeed/pkg/config" +) + +func TestRemoteDownloadRetriesFromReceivedOffset(t *testing.T) { + old := config.SysConfig + t.Cleanup(func() { config.SysConfig = old }) + for _, failure := range []string{"timeout", "short-body", "service-unavailable"} { + t.Run(failure, func(t *testing.T) { + config.SysConfig = &config.Config{Retry: config.Retry{Attempts: 3}, Download: config.Download{RespChunkSize: 4}} + var calls atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + n := calls.Add(1) + if n == 1 && failure == "service-unavailable" { + w.WriteHeader(503) + return + } + if n == 1 { + w.Header().Set("Content-Length", "8") + fmt.Fprint(w, "abcd") + w.(http.Flusher).Flush() + if failure == "timeout" { + <-r.Context().Done() + } + return + } + if failure == "service-unavailable" { + fmt.Fprint(w, "abcdefgh") + return + } + if r.Header.Get("Range") != "bytes=4-7" { + t.Errorf("wrong retry Range: %q", r.Header.Get("Range")) + } + w.Header().Set("Content-Length", "4") + w.Header().Set("Content-Range", "bytes 4-7/8") + w.WriteHeader(206) + fmt.Fprint(w, "efgh") + })) + defer server.Close() + cache, err := NewDingCache(filepath.Join(t.TempDir(), "blob"), 4) + if err != nil { + t.Fatal(err) + } + defer cache.Close() + if err := cache.Resize(8); err != nil { + t.Fatal(err) + } + client := &http.Client{Timeout: 100 * time.Millisecond} + r := NewRemoteFileTask(0, 0, 8) + r.Context, r.DingFile, r.Domain, r.Uri = context.Background(), cache, server.URL, "/file" + r.Source = &RemoteSource{Fetch: func(ctx context.Context, domain, uri string, headers map[string]string, consume func(*http.Response) error) error { + req, _ := http.NewRequestWithContext(ctx, "GET", domain+uri, nil) + for k, v := range headers { + req.Header.Set(k, v) + } + resp, err := client.Do(req) + if err != nil { + return err + } + defer resp.Body.Close() + return consume(resp) + }} + chunks := make(chan []byte, 32) + err = r.getFileRangeFromRemote(0, 8, chunks) + close(chunks) + var body []byte + for chunk := range chunks { + body = append(body, chunk...) + } + if err != nil || string(body) != "abcdefgh" || calls.Load() != 2 { + t.Fatalf("calls=%d body=%q error=%v", calls.Load(), body, err) + } + }) + } +} diff --git a/internal/handler/cache_admin_handler.go b/internal/handler/cache_admin_handler.go index e1ab268..8b21cba 100644 --- a/internal/handler/cache_admin_handler.go +++ b/internal/handler/cache_admin_handler.go @@ -15,6 +15,8 @@ package handler import ( + "bytes" + "encoding/json" "fmt" "io" "net/http" @@ -22,8 +24,9 @@ import ( "dingospeed/internal/dao" "dingospeed/internal/service" + "dingospeed/pkg/repository" + "dingospeed/pkg/util" - "github.com/bytedance/sonic" "github.com/labstack/echo/v4" ) @@ -59,7 +62,20 @@ func (h *CacheAdminHandler) ListRepos(c echo.Context) error { if err != nil { return writeUploadError(c, "cache repo list failed", err) } - return c.JSON(http.StatusOK, map[string]interface{}{"repos": result}) + filtered := result[:0] + for _, repo := range result { + if ns := c.QueryParam("namespace"); ns != "" && repo.Namespace != ns { + continue + } + if typ := c.QueryParam("repoType"); typ != "" && repo.RepoType != typ { + continue + } + if name := c.QueryParam("repo"); name != "" && repo.Repo != name { + continue + } + filtered = append(filtered, repo) + } + return c.JSON(http.StatusOK, map[string]interface{}{"repos": filtered}) } func (h *CacheAdminHandler) ListFiles(c echo.Context) error { @@ -103,15 +119,20 @@ func (h *CacheAdminHandler) PurgeOrphans(c echo.Context) error { } func parseCacheQuery(c echo.Context) service.CacheQuery { + id := "" + if c.QueryParam("repo") != "" { + id = util.GetOrgRepo(c.QueryParam("namespace"), c.QueryParam("repo")) + } return service.CacheQuery{ - RepoType: c.QueryParam("repoType"), - OrgRepo: c.QueryParam("orgRepo"), - Source: c.QueryParam("source"), - Keyword: c.QueryParam("keyword"), - Sort: c.QueryParam("sort"), - Order: c.QueryParam("order"), - Page: atoiOrZero(c.QueryParam("page")), - PageSize: atoiOrZero(c.QueryParam("pageSize")), + Namespace: c.QueryParam("namespace"), + RepoType: c.QueryParam("repoType"), + OrgRepo: id, + Source: c.QueryParam("source"), + Keyword: c.QueryParam("keyword"), + Sort: c.QueryParam("sort"), + Order: c.QueryParam("order"), + Page: atoiOrZero(c.QueryParam("page")), + PageSize: atoiOrZero(c.QueryParam("pageSize")), } } @@ -136,7 +157,9 @@ func readDeleteItems(c echo.Context) ([]dao.DeleteItem, error) { return nil, fmt.Errorf("request body too large") } var req deleteRequest - if err = sonic.Unmarshal(body, &req); err != nil { + decoder := json.NewDecoder(bytes.NewReader(body)) + decoder.DisallowUnknownFields() + if err = decoder.Decode(&req); err != nil { return nil, fmt.Errorf("body is not valid json: %v", err) } if len(req.Items) == 0 { @@ -145,5 +168,13 @@ func readDeleteItems(c echo.Context) ([]dao.DeleteItem, error) { if len(req.Items) > 2000 { return nil, fmt.Errorf("too many items in one request: %d", len(req.Items)) } + for i := range req.Items { + item := &req.Items[i] + key := repository.RepoKey{Namespace: item.Namespace, RepoType: item.RepoType, Repo: item.Repo} + if err := key.Validate(); err != nil { + return nil, err + } + item.OrgRepo = key.ID() + } return req.Items, nil } diff --git a/internal/handler/cache_job_handler.go b/internal/handler/cache_job_handler.go index 816b3bf..ed678f9 100644 --- a/internal/handler/cache_job_handler.go +++ b/internal/handler/cache_job_handler.go @@ -1,7 +1,10 @@ package handler import ( + "errors" "net/http" + "os" + "strconv" "dingospeed/internal/model/query" "dingospeed/internal/service" @@ -12,6 +15,18 @@ import ( "go.uber.org/zap" ) +func (handler *CacheJobHandler) CacheJobStatus(c echo.Context) error { + id, err := strconv.ParseInt(c.QueryParam("id"), 10, 64) + if err != nil || id <= 0 { + return echo.NewHTTPError(400, "invalid job id") + } + status, err := handler.cacheJobService.CacheJobStatus(id) + if err != nil { + return repositoryHTTPError(err) + } + return c.JSON(200, status) +} + type CacheJobHandler struct { cacheJobService *service.CacheJobService } @@ -33,15 +48,15 @@ func (handler *CacheJobHandler) CreateCacheJobHandler(c echo.Context) error { zap.S().Errorf("MetaProxyCommon repoType:%s is not exist RepoTypesMapping", createCacheJobReq.Datatype) return util.ErrorPageNotFound(c) } - if createCacheJobReq.Org == "" && createCacheJobReq.Repo == "" { + if createCacheJobReq.Namespace == "" && createCacheJobReq.Repo == "" { zap.S().Errorf("MetaProxyCommon org and repo is null") return util.ErrorRepoNotFound(c) } - jobId, err := handler.cacheJobService.CreateCacheJob(c, createCacheJobReq) + result, err := handler.cacheJobService.CreateCacheJobResult(c, createCacheJobReq) if err != nil { return util.ResponseError(c, err) } - return util.ResponseData(c, util.Body{Msg: "success", Data: jobId}) + return c.JSON(http.StatusOK, result) } func (handler *CacheJobHandler) StopCacheJobHandler(c echo.Context) error { @@ -52,10 +67,7 @@ func (handler *CacheJobHandler) StopCacheJobHandler(c echo.Context) error { }) } err := handler.cacheJobService.StopCacheJob(jobStatusReq) - if err != nil { - return util.ResponseError(c, err) - } - return util.ResponseData(c, nil) + return handler.cacheActionResponse(c, jobStatusReq.Id, err) } func (handler *CacheJobHandler) ResumeCacheJobHandler(c echo.Context) error { @@ -66,10 +78,7 @@ func (handler *CacheJobHandler) ResumeCacheJobHandler(c echo.Context) error { }) } err := handler.cacheJobService.ResumeCacheJob(c, resumeJobReq) - if err != nil { - return util.ResponseError(c, err) - } - return util.ResponseData(c, nil) + return handler.cacheActionResponse(c, resumeJobReq.Id, err) } func (handler *CacheJobHandler) RealtimeCacheJobHandler(c echo.Context) error { @@ -82,3 +91,25 @@ func (handler *CacheJobHandler) RealtimeCacheJobHandler(c echo.Context) error { resp := handler.cacheJobService.RealtimeCacheJob(realtimeReq) return util.ResponseData(c, resp) } + +func (handler *CacheJobHandler) cacheActionResponse(c echo.Context, id int64, err error) error { + if err != nil { + var conflict *service.CacheJobConflict + if errors.As(err, &conflict) { + return c.JSON(409, map[string]any{"error": conflict.Code, "code": conflict.Code, "activeJobId": conflict.ActiveJobID}) + } + return repositoryHTTPError(err) + } + status, err := handler.cacheJobService.CacheJobStatus(id) + if errors.Is(err, os.ErrNotExist) { + return util.ResponseData(c, nil) + } // Legacy mount jobs have no preheat snapshot. + if err != nil { + return repositoryHTTPError(err) + } + code := http.StatusOK + if status.State == "pausing" || status.State == "canceling" || status.State == "resuming" { + code = http.StatusAccepted + } + return c.JSON(code, status) +} diff --git a/internal/handler/cache_query.go b/internal/handler/cache_query.go new file mode 100644 index 0000000..c2b5b06 --- /dev/null +++ b/internal/handler/cache_query.go @@ -0,0 +1,46 @@ +package handler + +import ( + "dingospeed/pkg/repository" + "github.com/labstack/echo/v4" + "net/url" + "strings" +) + +func CacheQueryGuard(next echo.HandlerFunc) echo.HandlerFunc { + return func(c echo.Context) error { + if !strings.HasPrefix(c.Request().URL.Path, "/api/cache/") { + return next(c) + } + q, err := url.ParseQuery(c.Request().URL.RawQuery) + if err != nil { + return echo.NewHTTPError(400, "invalid query") + } + for name, values := range q { + if name == "orgRepo" || name == "org" { + return echo.NewHTTPError(400, "use namespace and repo") + } + if len(values) != 1 { + return echo.NewHTTPError(400, "duplicate query: "+name) + } + } + namespace, repo, typ := q.Get("namespace"), q.Get("repo"), q.Get("repoType") + if typ != "" && typ != "models" && typ != "datasets" && typ != "spaces" { + return echo.NewHTTPError(400, "invalid repoType") + } + if namespace != "" { + if err := repository.Segment(namespace); err != nil { + return echo.NewHTTPError(400, err.Error()) + } + } + if repo != "" { + if namespace == "" { + return echo.NewHTTPError(400, "namespace is required with repo") + } + if err := repository.Relative(repo); err != nil { + return echo.NewHTTPError(400, err.Error()) + } + } + return next(c) + } +} diff --git a/internal/handler/file_handler.go b/internal/handler/file_handler.go index 9ffc78f..a7f6166 100644 --- a/internal/handler/file_handler.go +++ b/internal/handler/file_handler.go @@ -15,8 +15,8 @@ package handler import ( + "dingospeed/pkg/repository" "fmt" - "net/url" "strconv" "dingospeed/internal/service" @@ -48,6 +48,9 @@ func (handler *FileHandler) HeadFileHandler1(c echo.Context) error { repoType, orgRepo, commit, filePath, err := paramProcess(c, 1) if err != nil { zap.S().Error("解码出错:%v", err) + if he, ok := err.(*echo.HTTPError); ok { + return he + } return util.ErrorRequestParam(c) } return handler.fileService.FileHeadCommon(c, repoType, orgRepo, commit, filePath) @@ -57,6 +60,9 @@ func (handler *FileHandler) HeadFileHandler2(c echo.Context) error { repoType, orgRepo, commit, filePath, err := paramProcess(c, 2) if err != nil { zap.S().Error("解码出错:%v", err) + if he, ok := err.(*echo.HTTPError); ok { + return he + } return util.ErrorRequestParam(c) } return handler.fileService.FileHeadCommon(c, repoType, orgRepo, commit, filePath) @@ -66,6 +72,9 @@ func (handler *FileHandler) HeadFileHandler3(c echo.Context) error { repoType, orgRepo, commit, filePath, err := paramProcess(c, 3) if err != nil { zap.S().Error("解码出错:%v", err) + if he, ok := err.(*echo.HTTPError); ok { + return he + } return util.ErrorRequestParam(c) } return handler.fileService.FileHeadCommon(c, repoType, orgRepo, commit, filePath) @@ -75,6 +84,9 @@ func (handler *FileHandler) GetFileHandler1(c echo.Context) error { repoType, orgRepo, commit, filePath, err := paramProcess(c, 1) if err != nil { zap.S().Error("解码出错:%v", err) + if he, ok := err.(*echo.HTTPError); ok { + return he + } return util.ErrorRequestParam(c) } return handler.fileGetCommon(c, repoType, orgRepo, commit, filePath) @@ -84,6 +96,9 @@ func (handler *FileHandler) GetFileHandler2(c echo.Context) error { repoType, orgRepo, commit, filePath, err := paramProcess(c, 2) if err != nil { zap.S().Error("解码出错:%v", err) + if he, ok := err.(*echo.HTTPError); ok { + return he + } return util.ErrorRequestParam(c) } return handler.fileGetCommon(c, repoType, orgRepo, commit, filePath) @@ -93,6 +108,9 @@ func (handler *FileHandler) GetFileHandler3(c echo.Context) error { repoType, orgRepo, commit, filePath, err := paramProcess(c, 3) if err != nil { zap.S().Error("解码出错:%v", err) + if he, ok := err.(*echo.HTTPError); ok { + return he + } return util.ErrorRequestParam(c) } return handler.fileGetCommon(c, repoType, orgRepo, commit, filePath) @@ -111,12 +129,18 @@ func paramProcess(c echo.Context, processMode int) (string, string, string, stri org = c.Param("org") repo = c.Param("repo") commit = c.Param("commit") - filePath = c.Param("filePath") + filePath = c.Param("*") + if filePath == "" { + filePath = c.Param("filePath") + } } else if processMode == 2 { orgOrRepoType := c.Param("orgOrRepoType") repo = c.Param("repo") commit = c.Param("commit") - filePath = c.Param("filePath") + filePath = c.Param("*") + if filePath == "" { + filePath = c.Param("filePath") + } if _, ok := consts.RepoTypesMapping[orgOrRepoType]; ok { repoType = orgOrRepoType org = "" @@ -127,12 +151,19 @@ func paramProcess(c echo.Context, processMode int) (string, string, string, stri } else if processMode == 3 { repo = c.Param("repo") commit = c.Param("commit") - filePath = c.Param("filePath") + filePath = c.Param("*") + if filePath == "" { + filePath = c.Param("filePath") + } repoType = "models" } else { panic("param process error.") } - orgRepo := util.GetOrgRepo(org, repo) + key, err := HFProtocolKey(c, repoType, org, repo) + if err != nil { + return repoType, "", commit, filePath, err + } + orgRepo := key.ID() c.Set(consts.PromOrgRepo, orgRepo) if _, ok := consts.RepoTypesMapping[repoType]; !ok { @@ -143,8 +174,10 @@ func paramProcess(c echo.Context, processMode int) (string, string, string, stri zap.S().Errorf("FileGetCommon or and repo is null") return repoType, orgRepo, commit, filePath, fmt.Errorf("FileGetCommon or and repo is null") } - filePath, err := url.QueryUnescape(filePath) - return repoType, orgRepo, commit, filePath, err + if err := repository.Segment(commit); err != nil { + return repoType, orgRepo, commit, filePath, err + } + return repoType, orgRepo, commit, filePath, repository.Relative(filePath) } func (handler *FileHandler) fileGetCommon(c echo.Context, repoType, orgRepo, commit, filePath string) error { diff --git a/internal/handler/hf_projection_handler.go b/internal/handler/hf_projection_handler.go new file mode 100644 index 0000000..50d2046 --- /dev/null +++ b/internal/handler/hf_projection_handler.go @@ -0,0 +1,82 @@ +package handler + +import ( + "errors" + "net/http" + "os" + "time" + + "dingospeed/pkg/config" + "dingospeed/pkg/hfprojection" + "dingospeed/pkg/repository" + "github.com/labstack/echo/v4" + "go.uber.org/zap" +) + +func HFLocalCatalog(c echo.Context) error { + if _, err := parseUniqueQuery(c); err != nil { + return err + } + if c.Param("namespace") != repository.HuggingFace && c.Param("namespace") != repository.ModelScope { + return echo.NewHTTPError(400, "local projection requires a remote repository") + } + k := repository.RepoKey{Namespace: c.Param("namespace"), RepoType: c.Param("repoType"), Repo: "placeholder"} + if err := k.Validate(); err != nil { + return echo.NewHTTPError(400, err.Error()) + } + var out hfprojection.Catalog + var err error + if k.Namespace == repository.ModelScope { + out, err = (hfprojection.Reader{Root: config.SysConfig.Repos(), Namespace: k.Namespace}).Catalog(k.RepoType) + } else { + out, err = hfprojection.DefaultIndex.Refresh(config.SysConfig.Repos(), k.RepoType) + } + if err != nil { + return hfProjectionError(err) + } + return c.JSON(200, out) +} + +func HFLocalManifest(c echo.Context) error { + k := requestRepoKey(c) + if k.Namespace != repository.HuggingFace && k.Namespace != repository.ModelScope { + return echo.NewHTTPError(400, "local projection requires a remote repository") + } + var out hfprojection.Manifest + var err error + if k.Namespace == repository.ModelScope { + out, err = (hfprojection.Reader{Root: config.SysConfig.Repos(), Namespace: k.Namespace}).Manifest(k.RepoType, k.Repo, c.QueryParam("revision")) + } else { + out, err = hfprojection.DefaultIndex.Manifest(config.SysConfig.Repos(), k.RepoType, k.Repo, c.QueryParam("revision")) + } + if err != nil { + return hfProjectionError(err) + } + return c.JSON(200, out) +} + +func HFLocalFile(c echo.Context) error { + k := requestRepoKey(c) + if k.Namespace != repository.HuggingFace && k.Namespace != repository.ModelScope { + return echo.NewHTTPError(400, "local projection requires a remote repository") + } + f, content, err := (hfprojection.Reader{Root: config.SysConfig.Repos(), Namespace: k.Namespace}).OpenFile(k.RepoType, k.Repo, c.QueryParam("revision"), c.QueryParam("path")) + if errors.Is(err, hfprojection.ErrIncomplete) { + return echo.NewHTTPError(409, err.Error()) + } + if err != nil { + return hfProjectionError(err) + } + defer f.Close() + c.Response().Header().Set("Cache-Control", "no-store") + http.ServeContent(c.Response(), c.Request(), c.QueryParam("path"), time.Time{}, content) + return nil +} + +func hfProjectionError(err error) error { + if errors.Is(err, os.ErrNotExist) { + return echo.NewHTTPError(404, "local cache entry not found") + } + zap.S().Warnw("Could not read HF cache projection", "error", err) + return echo.NewHTTPError(500, "local cache is unreadable or metadata is invalid") +} diff --git a/internal/handler/meta_handler.go b/internal/handler/meta_handler.go index 6e786f6..e8d4c9f 100644 --- a/internal/handler/meta_handler.go +++ b/internal/handler/meta_handler.go @@ -45,7 +45,11 @@ func (handler *MetaHandler) GetMetadataHandler(c echo.Context) error { repo := c.Param("repo") revision := c.Param("revision") method := strings.ToLower(c.Request().Method) - orgRepo := util.GetOrgRepo(org, repo) + key, keyErr := HFProtocolKey(c, repoType, org, repo) + if keyErr != nil { + return keyErr + } + orgRepo := key.ID() c.Set(consts.PromOrgRepo, orgRepo) if _, ok := consts.RepoTypesMapping[repoType]; !ok { zap.S().Errorf("repoType:%s is not exist RepoTypesMapping", repoType) @@ -84,7 +88,11 @@ func (handler *MetaHandler) GetRepoTreeHandler(c echo.Context) error { repo := c.Param("repo") revision := c.Param("revision") pathInRepo := c.Param("*") - orgRepo := util.GetOrgRepo(org, repo) + key, keyErr := HFProtocolKey(c, repoType, org, repo) + if keyErr != nil { + return keyErr + } + orgRepo := key.ID() c.Set(consts.PromOrgRepo, orgRepo) if _, ok := consts.RepoTypesMapping[repoType]; !ok { zap.S().Errorf("repoType:%s is not exist RepoTypesMapping", repoType) @@ -173,6 +181,21 @@ func (handler *MetaHandler) RepoRefsHandler(c echo.Context) error { } func (handler *MetaHandler) ForwardToNewSiteHandler(c echo.Context) error { + if c.Get("forcedLocal") == true { + return echo.NewHTTPError(http.StatusNotFound, "hosted operation is not available via proxy") + } + if ownAPIPath(c.Request().URL.Path) { + return echo.NewHTTPError(404, "endpoint removed") + } + segments := strings.Split(strings.TrimPrefix(c.Request().URL.Path, "/"), "/") + if len(segments) > 2 && segments[0] == "api" && segments[1] != "v1" { + segments = segments[2:] + } + if len(segments) >= 2 { + if c.Get("forcedProvider") != "huggingface" && segments[0] == "dingo-local" { + return echo.NewHTTPError(http.StatusNotFound, "hosted operation is not available via proxy") + } + } return handler.metaService.ForwardToNewSite(c) } @@ -198,3 +221,12 @@ func (handler *MetaHandler) RepositoryFilesHandler(c echo.Context) error { } return util.ResponseData(c, files) } + +func ownAPIPath(path string) bool { + for _, prefix := range []string{"/api/local-", "/api/repositories/", "/api/uploads/", "/api/upload-chunks/", "/api/upload-progress/", "/api/publish/", "/api/publish-tree/", "/api/cache/", "/api/fileOffset/"} { + if strings.HasPrefix(path, prefix) { + return true + } + } + return false +} diff --git a/internal/handler/meta_local_errors_test.go b/internal/handler/meta_local_errors_test.go new file mode 100644 index 0000000..c2ad463 --- /dev/null +++ b/internal/handler/meta_local_errors_test.go @@ -0,0 +1,119 @@ +package handler + +import ( + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strings" + "sync/atomic" + "testing" + "time" + + "dingospeed/internal/dao" + "dingospeed/internal/data" + "dingospeed/internal/service" + "dingospeed/pkg/config" + "dingospeed/pkg/dependency" + "github.com/labstack/echo/v4" +) + +func TestLocalMetadataHTTPErrorContract(t *testing.T) { + const body = `{"sha":"abc123","id":"org/repo"}` + var calls atomic.Int32 + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { calls.Add(1); w.WriteHeader(502) })) + defer upstream.Close() + old := config.SysConfig + defer func() { config.SysConfig = old }() + for _, method := range []string{"GET", "HEAD"} { + for _, tc := range []struct { + name string + status int + cached bool + }{ + {"success", 200, false}, {"missing", 404, false}, {"stat", 503, false}, {"read", 503, false}, {"corrupt", 500, false}, + {"success", 200, true}, {"missing", 404, true}, {"stat", 503, true}, {"read", 503, true}, {"corrupt", 500, true}, + } { + name := tc.name + if tc.cached { + name += "-cached-sha" + } + t.Run(method+"/"+name, func(t *testing.T) { + root := t.TempDir() + oldMonitor := dependency.Default + dependency.Default = &dependency.Monitor{} + t.Cleanup(func() { dependency.Default = oldMonitor }) + if tc.name == "stat" { + root = "invalid\x00repos" + } + config.SysConfig = &config.Config{Server: config.ServerConfig{Repos: root, HfScheme: "http", HfNetLoc: strings.TrimPrefix(upstream.URL, "http://")}, Retry: config.Retry{Attempts: 1}} + base := data.NewBaseData() + lock := dao.NewLockDao(base) + file := dao.NewFileDao(nil, base, lock) + if tc.cached { + base.Cache.Set(dao.GetMetaShaRepoKey("models/huggingface/org/repo", "main", ""), "abc123", time.Minute) + } + writeCache := func(revision, verb string) { + p := filepath.Join(root, "api/models/org/repo/revision", revision, "meta_"+verb+".json") + if err := os.MkdirAll(filepath.Dir(p), 0700); err != nil { + t.Fatal(err) + } + if err := file.WriteCacheRequest(p, 200, map[string]string{"content-type": "application/json"}, []byte(body)); err != nil { + t.Fatal(err) + } + } + if tc.name == "success" { + writeCache("main", "get") + writeCache("abc123", strings.ToLower(method)) + } + if tc.name == "read" || tc.name == "corrupt" { + revision, verb := "main", "get" + if tc.cached { + revision = "abc123" + verb = strings.ToLower(method) + } + p := filepath.Join(root, "api/models/org/repo/revision", revision, "meta_"+verb+".json") + if err := os.MkdirAll(filepath.Dir(p), 0700); err != nil { + t.Fatal(err) + } + if tc.name == "read" { + if err := os.Mkdir(p, 0700); err != nil { + t.Fatal(err) + } + } else if err := os.WriteFile(p, []byte("bad-json"), 0600); err != nil { + t.Fatal(err) + } + } + h := NewMetaHandler(service.NewMetaService(file, dao.NewMetaDao(file, lock, base))) + e := echo.New() + // The native HF route identifies the provider explicitly; avoid + // unrelated namespace registry I/O masking metadata failures. + e.Add(method, "/api/:repoType/:org/:repo/revision/:revision", func(c echo.Context) error { c.Set("forcedProvider", "huggingface"); return h.GetMetadataHandler(c) }) + r := httptest.NewRecorder() + e.ServeHTTP(r, httptest.NewRequest(method, "/api/models/org/repo/revision/main", nil)) + if r.Code != tc.status { + t.Fatalf("status=%d want=%d body=%s", r.Code, tc.status, r.Body.String()) + } + if calls.Load() != 0 { + t.Fatal("local metadata failure triggered upstream") + } + observation := dependency.Default.Snapshot(dependency.MetadataRead, time.Now()) + if tc.status == 503 && !observation.Unresolved { + t.Fatal("storage failure missing from heartbeat observation") + } + if tc.status == 200 && observation.LastObservation.IsZero() { + t.Fatal("successful cached-SHA read missing from recovery observation") + } + if (tc.status == 404 || tc.name == "corrupt") && observation.Unresolved { + t.Fatal("non-storage error incorrectly marked as storage failure") + } + if method == "GET" && tc.status == 200 && r.Body.String() != body { + t.Fatalf("normal payload changed: %s", r.Body.String()) + } + if tc.status >= 500 && (strings.Contains(r.Body.String(), root) || strings.Contains(r.Body.String(), "meta_get.json")) { + t.Fatal("internal path leaked") + } + }) + } + } +} diff --git a/internal/handler/meta_storage_contract_test.go b/internal/handler/meta_storage_contract_test.go new file mode 100644 index 0000000..b187452 --- /dev/null +++ b/internal/handler/meta_storage_contract_test.go @@ -0,0 +1,135 @@ +package handler + +import ( + "fmt" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strings" + "sync/atomic" + "testing" + + "dingospeed/internal/dao" + "dingospeed/internal/data" + "dingospeed/internal/service" + "dingospeed/pkg/config" + "github.com/labstack/echo/v4" +) + +// Online metadata cache failures do not discard successful HF responses. +// Local 404/503 and namespace validation retain their distinct behavior. +func TestMetadataStorageLegacyHTTPContract(t *testing.T) { + const body = `{"sha":"remote-commit","id":"org/repo"}` + var calls atomic.Int64 + var upstreamCode atomic.Int64 + upstreamCode.Store(200) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + calls.Add(1) + if r.Header.Get("Authorization") != "Bearer test-only" { + t.Error("authorization lost") + } + w.Header().Set("Content-Type", "application/json") + w.Header().Set("ETag", `"fixture"`) + w.WriteHeader(int(upstreamCode.Load())) + if r.Method != http.MethodHead { + fmt.Fprint(w, body) + } + })) + defer server.Close() + old := config.SysConfig + defer func() { config.SysConfig = old }() + for _, tc := range []struct { + name string + online bool + failure string + upstream, status int + count int64 + }{ + {"online-success", true, "", 200, 200, 2}, + {"online-mkdir-failure", true, "stat", 200, 200, 2}, + {"online-write-failure", true, "write", 200, 200, 2}, + {"online-descriptor-write-failure", true, "registration", 200, 200, 2}, + {"offline-missing", false, "", 200, 404, 0}, + {"offline-stat-failure", false, "stat", 200, 503, 0}, + {"offline-read-failure", false, "read", 200, 503, 0}, + {"upstream-denied", true, "", 403, 403, 1}, + {"upstream-missing", true, "", 404, 404, 1}, + {"upstream-unavailable", true, "", 503, 503, 1}, + } { + for _, method := range []string{http.MethodGet, http.MethodHead} { + t.Run(tc.name+"/"+method, func(t *testing.T) { + repos := t.TempDir() + if tc.failure == "registration" { + if err := os.MkdirAll(filepath.Join(repos, "api/.repository-registration.lock"), 0700); err != nil { + t.Fatal(err) + } + } + if tc.failure == "stat" { + repos = "invalid\x00repos" + } + if tc.failure == "read" || tc.failure == "write" { + cacheMethod := strings.ToLower(method) + if tc.failure == "read" { + cacheMethod = "get" + } + path := filepath.Join(repos, "api/models/org/repo/revision/main/meta_"+cacheMethod+".json") + if err := os.MkdirAll(path, 0700); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(path, "keep"), []byte("unchanged"), 0600); err != nil { + t.Fatal(err) + } + } + config.SysConfig = &config.Config{Server: config.ServerConfig{Online: tc.online, Repos: repos, HfScheme: "http", HfNetLoc: strings.TrimPrefix(server.URL, "http://")}, Retry: config.Retry{Attempts: 1}} + base := data.NewBaseData() + lock := dao.NewLockDao(base) + file := dao.NewFileDao(nil, base, lock) + h := NewMetaHandler(service.NewMetaService(file, dao.NewMetaDao(file, lock, base))) + e := echo.New() + e.Add(method, "/api/:repoType/:org/:repo/revision/:revision", h.GetMetadataHandler) + req := httptest.NewRequest(method, "/api/models/org/repo/revision/main", nil) + req.Header.Set("Authorization", "Bearer test-only") + rec := httptest.NewRecorder() + before := calls.Load() + upstreamCode.Store(int64(tc.upstream)) + e.ServeHTTP(rec, req) + if rec.Code != tc.status || calls.Load()-before != tc.count { + t.Fatalf("status=%d calls=%d; want %d/%d body=%s", rec.Code, calls.Load()-before, tc.status, tc.count, rec.Body.String()) + } + if tc.status == 200 && method == http.MethodGet && rec.Body.String() != body { + t.Fatalf("upstream payload changed: %s", rec.Body.String()) + } + if tc.status == 200 && rec.Header().Get("ETag") != `"fixture"` { + t.Fatal("upstream ETag lost") + } + if tc.status == 200 && method == http.MethodHead && strings.TrimSpace(rec.Body.String()) != "null" { + t.Fatal("legacy HEAD handler serialization changed") + } + if tc.status == 200 && tc.failure == "" { + if method == http.MethodGet && rec.Body.String() != body { + t.Fatalf("body changed: %s", rec.Body.String()) + } + for _, revision := range []string{"main", "remote-commit"} { + cache, err := file.ReadCacheRequest(filepath.Join(repos, "api/models/org/repo/revision", revision, "meta_"+strings.ToLower(method)+".json")) + if err != nil { + t.Fatal(err) + } + if method == http.MethodGet && string(cache.OriginContent) != body { + t.Fatal("cache payload changed") + } + } + } else if tc.failure == "write" || tc.failure == "read" { + cacheMethod := strings.ToLower(method) + if tc.failure == "read" { + cacheMethod = "get" + } + content, err := os.ReadFile(filepath.Join(repos, "api/models/org/repo/revision/main/meta_"+cacheMethod+".json/keep")) + if err != nil || string(content) != "unchanged" { + t.Fatal("existing cache modified") + } + } + }) + } + } +} diff --git a/internal/handler/modelscope_handler.go b/internal/handler/modelscope_handler.go index 3f3facb..9b16d34 100644 --- a/internal/handler/modelscope_handler.go +++ b/internal/handler/modelscope_handler.go @@ -5,6 +5,8 @@ import ( "strings" "dingospeed/internal/service" + "dingospeed/pkg/config" + "dingospeed/pkg/repository" "dingospeed/pkg/util" "github.com/labstack/echo/v4" @@ -31,7 +33,14 @@ func (m *ModelscopeHandler) ModelInfoHandler(c echo.Context) error { return util.ResponseError(c, err) } - org, repo, repoType := parts[3], parts[4], parts[2] + org, repo, repoType := c.Param("org"), c.Param("repo"), c.Param("repoType") + key := repository.RepoKey{Namespace: repository.ModelScope, RepoType: repoType, Repo: org + "/" + repo} + if err := key.Validate(); err != nil { + return echo.NewHTTPError(400, err.Error()) + } + if err := repository.Register(config.SysConfig.Repos(), repository.Remote(key)); err != nil { + return err + } if err := m.modelscopeService.ForwardModelInfo(c, org, repo, repoType); err != nil { return util.ResponseError(c, err) } @@ -47,7 +56,14 @@ func (m *ModelscopeHandler) RevisionsHandler(c echo.Context) error { return util.ResponseError(c, err) } - org, repo, repoType := parts[3], parts[4], parts[2] + org, repo, repoType := c.Param("org"), c.Param("repo"), c.Param("repoType") + key := repository.RepoKey{Namespace: repository.ModelScope, RepoType: repoType, Repo: org + "/" + repo} + if err := key.Validate(); err != nil { + return echo.NewHTTPError(400, err.Error()) + } + if err := repository.Register(config.SysConfig.Repos(), repository.Remote(key)); err != nil { + return err + } if err := m.modelscopeService.ForwardRevisions(c, org, repo, repoType); err != nil { return util.ResponseError(c, err) } @@ -63,7 +79,14 @@ func (m *ModelscopeHandler) FileListHandler(c echo.Context) error { return util.ResponseError(c, err) } - org, repo, repoType := parts[3], parts[4], parts[2] + org, repo, repoType := c.Param("org"), c.Param("repo"), c.Param("repoType") + key := repository.RepoKey{Namespace: repository.ModelScope, RepoType: repoType, Repo: org + "/" + repo} + if err := key.Validate(); err != nil { + return echo.NewHTTPError(400, err.Error()) + } + if err := repository.Register(config.SysConfig.Repos(), repository.Remote(key)); err != nil { + return err + } if err := m.modelscopeService.ForwardFileList(c, org, repo, repoType); err != nil { return util.ResponseError(c, err) } @@ -79,11 +102,15 @@ func (m *ModelscopeHandler) FileDownloadHandler(c echo.Context) error { return util.ResponseError(c, err) } - org, repo, repoType := parts[3], parts[4], parts[2] - if err := m.modelscopeService.HandleFileDownload(c, org, repo, repoType); err != nil { - return util.ResponseError(c, err) + org, repo, repoType := c.Param("org"), c.Param("repo"), c.Param("repoType") + key := repository.RepoKey{Namespace: repository.ModelScope, RepoType: repoType, Repo: org + "/" + repo} + if err := key.Validate(); err != nil { + return echo.NewHTTPError(400, err.Error()) } - return nil + if err := repository.Register(config.SysConfig.Repos(), repository.Remote(key)); err != nil { + return err + } + return m.modelscopeService.HandleFileDownload(c, org, repo, repoType) } // FileTreeHandler 处理数据集文件列表请求 @@ -95,7 +122,14 @@ func (m *ModelscopeHandler) FileTreeHandler(c echo.Context) error { return util.ResponseError(c, err) } - org, repo, repoType := parts[3], parts[4], parts[2] + org, repo, repoType := c.Param("org"), c.Param("repo"), c.Param("repoType") + key := repository.RepoKey{Namespace: repository.ModelScope, RepoType: repoType, Repo: org + "/" + repo} + if err := key.Validate(); err != nil { + return echo.NewHTTPError(400, err.Error()) + } + if err := repository.Register(config.SysConfig.Repos(), repository.Remote(key)); err != nil { + return err + } if err := m.modelscopeService.ForwardRepoTree(c, org, repo, repoType); err != nil { return util.ResponseError(c, err) } diff --git a/internal/handler/repository_handler.go b/internal/handler/repository_handler.go new file mode 100644 index 0000000..e2a2345 --- /dev/null +++ b/internal/handler/repository_handler.go @@ -0,0 +1,395 @@ +package handler + +import ( + "dingospeed/internal/dao" + "dingospeed/internal/service" + "dingospeed/pkg/common" + "dingospeed/pkg/config" + "dingospeed/pkg/repository" + "fmt" + "github.com/labstack/echo/v4" + "net/http" + "net/url" + "os" + "path/filepath" + "strconv" + "strings" +) + +var repositoryMS = service.NewModelscopeService() + +// RepositoryLocator validates the independent locator once. Values from +// net/url.Query are already decoded and must never be decoded again. +func RepositoryLocator(fileRequired, revisionRequired bool) echo.MiddlewareFunc { + return func(next echo.HandlerFunc) echo.HandlerFunc { + return func(c echo.Context) error { + q, err := parseUniqueQuery(c) + if err != nil { + return err + } + k := repository.RepoKey{Namespace: c.Param("namespace"), RepoType: c.Param("repoType"), Repo: q["repo"]} + if err = k.Validate(); err != nil { + return echo.NewHTTPError(400, err.Error()) + } + revision := q["revision"] + if revisionRequired || revision != "" { + if err = repository.Segment(revision); err != nil { + return echo.NewHTTPError(400, err.Error()) + } + } + file := q["path"] + if fileRequired || file != "" { + if err = repository.Relative(file); err != nil { + return echo.NewHTTPError(400, err.Error()) + } + } + c.Set("repositoryKey", k) + return next(c) + } + } +} + +func parseUniqueQuery(c echo.Context) (map[string]string, error) { + var values url.Values + var err error + // ParseQuery additionally reports malformed escapes and unescaped semicolons. + values, err = url.ParseQuery(c.Request().URL.RawQuery) + if err != nil { + return nil, echo.NewHTTPError(400, "invalid query encoding") + } + q := map[string]string{} + for key, vs := range values { + if len(vs) != 1 { + return nil, echo.NewHTTPError(400, "duplicate query parameter: "+key) + } + q[key] = vs[0] + } + for _, name := range []string{"namespace", "org", "orgRepo", "repoType"} { + if _, ok := q[name]; ok { + return nil, echo.NewHTTPError(400, "locator field must not override route: "+name) + } + } + return q, nil +} + +func requestRepoKey(c echo.Context) repository.RepoKey { + return c.Get("repositoryKey").(repository.RepoKey) +} + +// HFProtocolKey adapts only a standard upstream repo ID. Registered hosted +// namespaces require the same trusted download gateway as the own API. +func HFProtocolKey(c echo.Context, repoType, org, repo string) (repository.RepoKey, error) { + k := repository.RepoKey{Namespace: repository.HuggingFace, RepoType: repoType, Repo: repo} + if org != "" { + k.Repo = org + "/" + repo + } + forceHF := c.Get("forcedProvider") == repository.HuggingFace + hosted := !forceHF && (c.Get("forcedLocal") == true || org == repository.Local) + if hosted { + k.Namespace = org + k.Repo = repo + } + if err := k.Validate(); err != nil { + return k, echo.NewHTTPError(400, err.Error()) + } + if hosted { + d, err := repository.Read(config.SysConfig.Repos(), k) + if org == repository.Local && os.IsNotExist(err) { + // Pre-namespace uploads have manifests but no platform descriptor. + return k, nil + } + if err != nil { + return k, repositoryHTTPError(err) + } + if d.Source != "hosted" { + return k, echo.NewHTTPError(404, "hosted repository not found") + } + } + + return k, nil +} + +// The reserved HF source supports cold reads. Other namespaces need a hosted +// descriptor; corrupt descriptors and unavailable storage are never cache misses. +func readableRepository(k repository.RepoKey) error { + _, err := repository.Read(config.SysConfig.Repos(), k) + if k.Namespace == repository.HuggingFace && os.IsNotExist(err) { + return nil + } + return err +} + +func (h *FileHandler) RepositoryFile(c echo.Context) error { + if id := c.Request().Header.Get(common.LocalTransferHeader); id != "" { + finish, ok := common.ClaimLocalTransfer(id) + if !ok { + return echo.NewHTTPError(409, "local transfer no longer active") + } + defer finish() + if trace := common.LocalCacheTrace(id); trace != nil { + c.SetRequest(c.Request().WithContext(common.WithCacheTrace(c.Request().Context(), trace))) + } + } + + k := requestRepoKey(c) + if k.Namespace == repository.ModelScope { + return repositoryMS.RepositoryFile(c, k, c.QueryParam("revision"), c.QueryParam("path")) + } + if err := readableRepository(k); err != nil { + return repositoryHTTPError(err) + } + if c.Request().Method == http.MethodHead { + return h.fileService.FileHeadCommon(c, k.RepoType, k.ID(), c.QueryParam("revision"), c.QueryParam("path")) + } + return h.fileGetCommon(c, k.RepoType, k.ID(), c.QueryParam("revision"), c.QueryParam("path")) +} +func (h *MetaHandler) RepositorySnapshot(c echo.Context) error { + k := requestRepoKey(c) + d, err := repository.Read(config.SysConfig.Repos(), k) + if err != nil { + return repositoryHTTPError(err) + } + if d.Source != "hosted" { + return echo.NewHTTPError(400, "snapshot requires a hosted repository") + } + snap, err := h.metaService.GetLocalSnapshot(k.RepoType, k.ID(), c.QueryParam("revision")) + if err != nil { + return repositoryHTTPError(err) + } + verify := c.QueryParam("verify") + if verify == "true" || verify == "sha256" { + files := make([]dao.LocalManifestFile, len(snap.Files)) + for i, f := range snap.Files { + files[i] = dao.LocalManifestFile{Path: f.Path, Size: f.Size, Sha256: f.Sha256} + } + if err := dao.VerifyPublishedFiles(k.RepoType, k.ID(), snap.Commit, files); err != nil { + return repositoryHTTPError(err) + } + if verify == "sha256" { + if err := dao.VerifyPublishedSHA256(k.RepoType, k.ID(), files); err != nil { + return repositoryHTTPError(err) + } + } + } + return c.JSON(200, map[string]interface{}{"namespace": k.Namespace, "repoType": k.RepoType, "repo": k.Repo, "revision": c.QueryParam("revision"), "commit": snap.Commit, "files": snap.Files, "complete": true, "fileCount": len(snap.Files), "verified": verify == "true" || verify == "sha256", "contentVerified": verify == "sha256"}) +} +func (h *MetaHandler) RepositoryArchive(c echo.Context) error { + k := requestRepoKey(c) + d, err := repository.Read(config.SysConfig.Repos(), k) + if err != nil { + return repositoryHTTPError(err) + } + if d.Source != "hosted" { + return echo.NewHTTPError(400, "archive requires a hosted repository") + } + revision := c.QueryParam("revision") + if _, err = h.metaService.GetLocalSnapshot(k.RepoType, k.ID(), revision); err != nil { + return repositoryHTTPError(err) + } + setLocalArchiveHeaders(c, filepath.Base(k.Repo)+"-"+revision+".zip") + if c.Request().Method == http.MethodHead { + return c.NoContent(200) + } + return h.metaService.StreamLocalArchive(k.RepoType, k.ID(), revision, filepath.Base(k.Repo), c.Response()) +} +func (h *MetaHandler) RepositoryTree(c echo.Context) error { + k := requestRepoKey(c) + if k.Namespace == repository.ModelScope { + items, err := repositoryMS.RepositoryTree(c, k, c.QueryParam("revision"), c.QueryParam("path"), strings.EqualFold(c.QueryParam("recursive"), "true")) + if err != nil { + return err + } + return c.JSON(200, items) + } + if err := readableRepository(k); err != nil { + return repositoryHTTPError(err) + } + items, err := h.metaService.GetRepoTree(k.RepoType, k.ID(), c.QueryParam("revision"), c.QueryParam("path"), strings.EqualFold(c.QueryParam("recursive"), "true"), c.Request().Header.Get("Authorization")) + if err != nil { + return repositoryHTTPError(err) + } + if k.Namespace == repository.HuggingFace { + if err := repository.Register(config.SysConfig.Repos(), repository.Remote(k)); err != nil { + return repositoryHTTPError(err) + } + } + return c.JSON(200, items) +} +func (h *MetaHandler) RepositoryMetadata(c echo.Context) error { + if id := c.Request().Header.Get(common.LocalTransferHeader); id != "" { + finish, ok := common.ClaimLocalTransfer(id) + if !ok { + return echo.NewHTTPError(409, "local transfer no longer active") + } + defer finish() + if trace := common.LocalCacheTrace(id); trace != nil { + c.SetRequest(c.Request().WithContext(common.WithCacheTrace(c.Request().Context(), trace))) + } + } + + k := requestRepoKey(c) + if k.Namespace == repository.ModelScope { + meta, err := repositoryMS.RepositoryMetadata(c, k, c.QueryParam("revision")) + if err != nil { + return err + } + c.Response().Header().Set("X-Repo-Commit", fmt.Sprint(meta["sha"])) + if c.Request().Method == http.MethodHead { + return c.NoContent(200) + } + return c.JSON(200, meta) + } + if k.Namespace != repository.HuggingFace { + if _, err := repository.Read(config.SysConfig.Repos(), k); err != nil { + return repositoryHTTPError(err) + } + } + meta, err := h.metaService.GetMetadata(k.RepoType, k.ID(), c.QueryParam("revision"), strings.ToLower(c.Request().Method), c.Request().Header.Get("Authorization")) + if err != nil { + return repositoryHTTPError(err) + } + for key, value := range meta.Headers { + c.Response().Header().Set(key, value) + } + if c.Request().Method == http.MethodHead { + return c.NoContent(meta.StatusCode) + } + return c.Blob(meta.StatusCode, "application/json", meta.OriginContent) +} +func (h *MetaHandler) RepositoryDirectories(c echo.Context) error { + if _, err := parseUniqueQuery(c); err != nil { + return err + } + namespace, typ := c.Param("namespace"), c.Param("repoType") + if err := (repository.RepoKey{Namespace: namespace, RepoType: typ, Repo: "placeholder"}).Validate(); err != nil { + return echo.NewHTTPError(400, err.Error()) + } + all, err := repository.List(config.SysConfig.Repos()) + if err != nil { + return repositoryHTTPError(err) + } + result := []repository.RepoKey{} + for _, d := range all { + if d.Namespace == namespace && d.RepoType == typ { + result = append(result, d.RepoKey) + } + } + return c.JSON(200, result) +} +func (h *MetaHandler) RepositoryRevisions(c echo.Context) error { + k := requestRepoKey(c) + if k.Namespace == repository.ModelScope { + items, err := repositoryMS.RepositoryRevisions(c, k) + if err != nil { + return err + } + return c.JSON(200, items) + } + if _, err := repository.Read(config.SysConfig.Repos(), k); err != nil { + return repositoryHTTPError(err) + } + entries, err := os.ReadDir(filepath.Join(k.APIRoot(config.SysConfig.Repos()), "revision")) + if os.IsNotExist(err) { + entries = nil + } else if err != nil { + return repositoryHTTPError(err) + } + result := []map[string]string{} + for _, entry := range entries { + if !entry.IsDir() { + continue + } + commit, err := h.metaService.RevisionCommit(k.RepoType, k.ID(), entry.Name()) + if err != nil { + return repositoryHTTPError(err) + } + if entry.Name() != commit { + result = append(result, map[string]string{"name": entry.Name(), "commit": commit}) + } + } + return c.JSON(200, result) +} +func repositoryHTTPError(err error) error { + if e, ok := err.(*echo.HTTPError); ok { + return e + } + if os.IsNotExist(err) { + return echo.NewHTTPError(404, "repository entry not found") + } + if e, ok := err.(interface{ StatusCode() int }); ok { + return echo.NewHTTPError(e.StatusCode(), err.Error()) + } + return echo.NewHTTPError(500, fmt.Sprintf("repository access failed: %v", err)) +} + +func (h *MetaHandler) RepositoryFiles(c echo.Context) error { + k := requestRepoKey(c) + if k.Namespace == repository.ModelScope { + items, err := repositoryMS.RepositoryTree(c, k, c.QueryParam("revision"), c.QueryParam("path"), false) + if err != nil { + return err + } + files := []map[string]interface{}{} + for _, item := range items { + q := url.Values{"repo": {k.Repo}, "revision": {c.QueryParam("revision")}, "path": {item.Path}} + link := "/api/repositories/" + url.PathEscape(k.RepoType) + "/" + url.PathEscape(k.Namespace) + "/file?" + q.Encode() + files = append(files, map[string]interface{}{"name": filepath.Base(item.Path), "size": item.Size, "isDir": item.Type == "directory", "link": link}) + } + return c.JSON(200, files) + } + if _, err := repository.Read(config.SysConfig.Repos(), k); err != nil { + return repositoryHTTPError(err) + } + commit, err := h.metaService.RevisionCommit(k.RepoType, k.ID(), c.QueryParam("revision")) + if err != nil { + return repositoryHTTPError(err) + } + files, err := h.metaService.RepositoryFiles(k.RepoType, k.ID(), commit, c.QueryParam("path")) + if err != nil { + return repositoryHTTPError(err) + } + return c.JSON(200, files) +} +func (h *FileHandler) RepositoryOffset(c echo.Context) error { + k := requestRepoKey(c) + if _, err := repository.Read(config.SysConfig.Repos(), k); err != nil { + return repositoryHTTPError(err) + } + etag := c.QueryParam("etag") + if err := repository.Segment(etag); err != nil { + return echo.NewHTTPError(400, err.Error()) + } + size, err := strconv.ParseInt(c.QueryParam("size"), 10, 64) + if err != nil || size < 0 { + return echo.NewHTTPError(400, "invalid size") + } + return c.JSON(200, h.fileService.GetFileOffset(k.RepoType, k.Namespace, k.Repo, etag, size)) +} + +func (h *MetaHandler) ProtocolRefs(c echo.Context) error { + k, err := HFProtocolKey(c, c.Param("repoType"), c.Param("org"), c.Param("repo")) + if err != nil { + return err + } + if k.Namespace == repository.HuggingFace { + return h.metaService.ForwardToNewSite(c) + } + entries, err := os.ReadDir(filepath.Join(k.APIRoot(config.SysConfig.Repos()), "revision")) + if err != nil { + return repositoryHTTPError(err) + } + branches := []map[string]string{} + for _, entry := range entries { + if !entry.IsDir() { + continue + } + commit, err := h.metaService.RevisionCommit(k.RepoType, k.ID(), entry.Name()) + if err != nil { + return repositoryHTTPError(err) + } + if entry.Name() != commit { + branches = append(branches, map[string]string{"name": entry.Name(), "ref": "refs/heads/" + entry.Name(), "targetCommit": commit}) + } + } + return c.JSON(200, map[string]interface{}{"branches": branches, "tags": []interface{}{}, "converts": []interface{}{}}) +} diff --git a/internal/handler/scheduler_registration.go b/internal/handler/scheduler_registration.go new file mode 100644 index 0000000..ef5b11d --- /dev/null +++ b/internal/handler/scheduler_registration.go @@ -0,0 +1,29 @@ +package handler + +import ( + "dingospeed/internal/service" + "dingospeed/pkg/config" + "encoding/json" + "github.com/labstack/echo/v4" + "io" + "net/http" +) + +// Configuration is intentionally unauthenticated during connectivity rollout. +func SchedulerRegistration(c echo.Context) error { + if c.Request().Method == http.MethodPut { + var r config.Registration + dec := json.NewDecoder(http.MaxBytesReader(c.Response(), c.Request().Body, 4096)) + dec.DisallowUnknownFields() + if err := dec.Decode(&r); err != nil { + return c.JSON(400, map[string]string{"error": err.Error()}) + } + if dec.Decode(&struct{}{}) != io.EOF { + return c.JSON(400, map[string]string{"error": "expected one JSON object"}) + } + if err := service.SaveSchedulerRegistration(r); err != nil { + return c.JSON(409, map[string]string{"error": err.Error()}) + } + } + return c.JSON(200, map[string]any{"config": config.SysConfig.Registration(), "status": config.SysConfig.RegistrationState()}) +} diff --git a/internal/handler/scheduler_registration_test.go b/internal/handler/scheduler_registration_test.go new file mode 100644 index 0000000..250f585 --- /dev/null +++ b/internal/handler/scheduler_registration_test.go @@ -0,0 +1,32 @@ +package handler + +import ( + "dingospeed/pkg/config" + "encoding/json" + "github.com/labstack/echo/v4" + "net/http/httptest" + "strings" + "testing" +) + +func TestSchedulerRegistrationWithoutCredentials(t *testing.T) { + old := config.SysConfig + defer func() { config.SysConfig = old }() + cfg := &config.Config{} + // Repos uses the configured server repo root. + cfg.Server.Repos = t.TempDir() + config.SysConfig = cfg + t.Setenv("DINGO_NODE_MANAGEMENT_TOKEN", "ignored-old-token") + body := `{"enabled":true,"nodeId":"external-node","address":"scheduler:19091","host":"speed","port":8090,"managementUrl":"http://speed:8091","downloadUrl":"http://speed:8090"}` + for _, method := range []string{"PUT", "GET"} { + req := httptest.NewRequest(method, "/api/scheduler-registration", strings.NewReader(body)) + rec := httptest.NewRecorder() + if err := SchedulerRegistration(echo.New().NewContext(req, rec)); err != nil || rec.Code != 200 { + t.Fatalf("%s: %d %s %v", method, rec.Code, rec.Body.String(), err) + } + var out struct{ Config config.Registration } + if err := json.Unmarshal(rec.Body.Bytes(), &out); err != nil || out.Config.ManagementURL != "http://speed:8091" { + t.Fatalf("readback: %s %v", rec.Body.String(), err) + } + } +} diff --git a/internal/handler/transfer_settings.go b/internal/handler/transfer_settings.go new file mode 100644 index 0000000..63a6227 --- /dev/null +++ b/internal/handler/transfer_settings.go @@ -0,0 +1,59 @@ +package handler + +import ( + "dingospeed/pkg/transfersettings" + "encoding/json" + "github.com/labstack/echo/v4" + "io" + "net/http" + "strings" +) + +func TransferSettings(c echo.Context) error { + if c.Request().Method == http.MethodPut { + var s transfersettings.Settings + dec := json.NewDecoder(io.LimitReader(c.Request().Body, 4096)) + dec.DisallowUnknownFields() + if err := dec.Decode(&s); err != nil { + return c.JSON(400, map[string]string{"error": err.Error()}) + } + var extra any + if err := dec.Decode(&extra); err != io.EOF { + return c.JSON(400, map[string]string{"error": "expected one JSON object"}) + } + if err := transfersettings.Validate(s); err != nil { + return c.JSON(400, map[string]string{"error": err.Error()}) + } + if err := transfersettings.Save(s); err != nil { + return c.JSON(500, map[string]string{"error": err.Error()}) + } + } + return c.JSON(200, transfersettings.Current()) +} + +func LimitUpload(next echo.HandlerFunc) echo.HandlerFunc { + return func(c echo.Context) error { + release, err := transfersettings.AcquireUpload(c.Request().Context()) + if err != nil { + return err + } + defer release() + return next(c) + } +} + +func LimitDownload(next echo.HandlerFunc) echo.HandlerFunc { + return func(c echo.Context) error { + path := c.Path() + data := strings.Contains(path, "/resolve/") || strings.HasSuffix(path, "/local-file") || strings.HasSuffix(path, "/file") || strings.HasSuffix(path, "/archive") || path == "/api/v1/:repoType/:org/:repo/repo" + if !data { + return next(c) + } + release, err := transfersettings.AcquireDownload(c.Request().Context()) + if err != nil { + return err + } + defer release() + return next(c) + } +} diff --git a/internal/handler/upload_chunk_handler_test.go b/internal/handler/upload_chunk_handler_test.go index 3262d93..767f7e7 100644 --- a/internal/handler/upload_chunk_handler_test.go +++ b/internal/handler/upload_chunk_handler_test.go @@ -5,6 +5,7 @@ import ( "fmt" "net/http" "net/http/httptest" + "net/url" "testing" "github.com/bytedance/sonic" @@ -21,14 +22,14 @@ func doChunk(t *testing.T, h *UploadHandler, e *echo.Echo, filePath string, cont if chunkSha == "" { chunkSha = shaOf(body) } - target := fmt.Sprintf("/api/local-upload-chunk/models/dingo-local/demo/main/%s?size=%d&sha256=%s&offset=%d&chunkSha256=%s", - filePath, len(content), shaOf(content), offset, chunkSha) + target := fmt.Sprintf("/api/upload-chunks/models/dingo-local?repo=demo&revision=main&path=%s&size=%d&sha256=%s&offset=%d&chunkSha256=%s", + url.QueryEscape(filePath), len(content), shaOf(content), offset, chunkSha) req := httptest.NewRequest(http.MethodPut, target, bytes.NewReader(body)) // httptest.NewRequest 已按 body 长度填好 ContentLength，分块接口依赖它。 rec := httptest.NewRecorder() c := e.NewContext(req, rec) - c.SetParamNames("repoType", "org", "repo", "revision", "*") - c.SetParamValues("models", "dingo-local", "demo", "main", filePath) + c.SetParamNames("repoType", "namespace") + c.SetParamValues("models", "dingo-local") if err := h.UploadChunk(c); err != nil { t.Fatalf("chunk handler returned error: %v", err) } @@ -113,12 +114,12 @@ func TestChunkProgressReportsMissingRanges(t *testing.T) { content := bytes.Repeat([]byte("q"), handlerBlockSize*4) doChunk(t, h, e, "weights.bin", content, 0, handlerBlockSize, "") - target := fmt.Sprintf("/api/local-upload-progress/models/dingo-local/demo/main/weights.bin?sha256=%s", shaOf(content)) + target := fmt.Sprintf("/api/upload-progress/models/dingo-local?repo=demo&revision=main&path=weights.bin&sha256=%s", shaOf(content)) req := httptest.NewRequest(http.MethodGet, target, nil) rec := httptest.NewRecorder() c := e.NewContext(req, rec) - c.SetParamNames("repoType", "org", "repo", "revision", "*") - c.SetParamValues("models", "dingo-local", "demo", "main", "weights.bin") + c.SetParamNames("repoType", "namespace") + c.SetParamValues("models", "dingo-local", "weights.bin") if err := h.QueryProgress(c); err != nil { t.Fatalf("progress handler returned error: %v", err) } diff --git a/internal/handler/upload_handler.go b/internal/handler/upload_handler.go index 83de607..5c9435b 100644 --- a/internal/handler/upload_handler.go +++ b/internal/handler/upload_handler.go @@ -22,13 +22,21 @@ func NewUploadHandler(uploadService *service.UploadService) *UploadHandler { return &UploadHandler{uploadService: uploadService} } +func (h *UploadHandler) DeleteRepository(c echo.Context) error { + result, err := h.uploadService.DeleteRepository(requestRepoKey(c)) + if err != nil { + return writeUploadError(c, "repository deletion failed", err) + } + return c.JSON(http.StatusOK, result) +} + func (h *UploadHandler) UploadWholeFile(c echo.Context) error { param := dao.LocalUploadParam{ RepoType: c.Param("repoType"), - Org: c.Param("org"), - Repo: c.Param("repo"), - Revision: c.Param("revision"), - FilePath: c.Param("*"), + Namespace: c.Param("namespace"), + Repo: c.QueryParam("repo"), + Revision: c.QueryParam("revision"), + FilePath: c.QueryParam("path"), Sha256: c.QueryParam("sha256"), Overwrite: strings.EqualFold(c.QueryParam("overwrite"), "true"), Deferred: strings.EqualFold(c.QueryParam("defer"), "true"), @@ -47,12 +55,12 @@ func (h *UploadHandler) UploadWholeFile(c echo.Context) error { // 分块上传因此永远是暂缓生效的，清单才是可见性闸门。 func (h *UploadHandler) UploadChunk(c echo.Context) error { param := dao.LocalChunkUploadParam{ - RepoType: c.Param("repoType"), - Org: c.Param("org"), - Repo: c.Param("repo"), - Revision: c.Param("revision"), - FilePath: c.Param("*"), - Sha256: c.QueryParam("sha256"), + RepoType: c.Param("repoType"), + Namespace: c.Param("namespace"), + Repo: c.QueryParam("repo"), + Revision: c.QueryParam("revision"), + FilePath: c.QueryParam("path"), + Sha256: c.QueryParam("sha256"), } result, err := h.uploadService.UploadChunk( param, @@ -102,12 +110,13 @@ func (h *UploadHandler) PublishFiles(c echo.Context) error { }) } param := dao.LocalPublishParam{ - RepoType: c.Param("repoType"), - Org: c.Param("org"), - Repo: c.Param("repo"), - Revision: c.Param("revision"), - Overwrite: strings.EqualFold(c.QueryParam("overwrite"), "true"), - Files: make([]dao.LocalManifestFile, 0, len(req.Files)), + CreateOnly: strings.EqualFold(c.QueryParam("createOnly"), "true"), + RepoType: c.Param("repoType"), + Namespace: c.Param("namespace"), + Repo: c.QueryParam("repo"), + Revision: c.QueryParam("revision"), + Overwrite: strings.EqualFold(c.QueryParam("overwrite"), "true"), + Files: make([]dao.LocalManifestFile, 0, len(req.Files)), } for _, item := range req.Files { param.Files = append(param.Files, dao.LocalManifestFile{ @@ -149,9 +158,9 @@ func (h *UploadHandler) PublishTree(c echo.Context) error { } param := dao.LocalPublishTreeParam{ RepoType: c.Param("repoType"), - Org: c.Param("org"), - Repo: c.Param("repo"), - Revision: c.Param("revision"), + Namespace: c.Param("namespace"), + Repo: c.QueryParam("repo"), + Revision: c.QueryParam("revision"), BaseCommit: req.BaseCommit, Files: make([]dao.LocalManifestFile, 0, len(req.Files)), } @@ -171,12 +180,12 @@ func (h *UploadHandler) PublishTree(c echo.Context) error { func (h *UploadHandler) QueryProgress(c echo.Context) error { param := dao.LocalUploadParam{ - RepoType: c.Param("repoType"), - Org: c.Param("org"), - Repo: c.Param("repo"), - Revision: c.Param("revision"), - FilePath: c.Param("*"), - Sha256: c.QueryParam("sha256"), + RepoType: c.Param("repoType"), + Namespace: c.Param("namespace"), + Repo: c.QueryParam("repo"), + Revision: c.QueryParam("revision"), + FilePath: c.QueryParam("path"), + Sha256: c.QueryParam("sha256"), } result, err := h.uploadService.QueryProgress(param) if err != nil { diff --git a/internal/handler/upload_handler_test.go b/internal/handler/upload_handler_test.go index 712143f..26de3d2 100644 --- a/internal/handler/upload_handler_test.go +++ b/internal/handler/upload_handler_test.go @@ -7,6 +7,7 @@ import ( "fmt" "net/http" "net/http/httptest" + "net/url" "strings" "testing" @@ -43,13 +44,13 @@ func shaOf(b []byte) string { // doUpload 走一次 POST /api/local-upload/... 的完整 handler 路径。 func doUpload(t *testing.T, h *UploadHandler, e *echo.Echo, filePath string, content []byte, query string) *httptest.ResponseRecorder { t.Helper() - target := fmt.Sprintf("/api/local-upload/models/dingo-local/demo/main/%s?size=%d&sha256=%s&%s", - filePath, len(content), shaOf(content), query) + target := fmt.Sprintf("/api/uploads/models/dingo-local?repo=demo&revision=main&path=%s&size=%d&sha256=%s&%s", + url.QueryEscape(filePath), len(content), shaOf(content), query) req := httptest.NewRequest(http.MethodPost, target, bytes.NewReader(content)) rec := httptest.NewRecorder() c := e.NewContext(req, rec) - c.SetParamNames("repoType", "org", "repo", "revision", "*") - c.SetParamValues("models", "dingo-local", "demo", "main", filePath) + c.SetParamNames("repoType", "namespace") + c.SetParamValues("models", "dingo-local") if err := h.UploadWholeFile(c); err != nil { t.Fatalf("upload handler returned error: %v", err) } @@ -58,12 +59,12 @@ func doUpload(t *testing.T, h *UploadHandler, e *echo.Echo, filePath string, con func doPublish(t *testing.T, h *UploadHandler, e *echo.Echo, body string, query string) *httptest.ResponseRecorder { t.Helper() - target := "/api/local-publish/models/dingo-local/demo/main?" + query + target := "/api/publish/models/dingo-local?repo=demo&revision=main&" + query req := httptest.NewRequest(http.MethodPost, target, strings.NewReader(body)) rec := httptest.NewRecorder() c := e.NewContext(req, rec) - c.SetParamNames("repoType", "org", "repo", "revision") - c.SetParamValues("models", "dingo-local", "demo", "main") + c.SetParamNames("repoType", "namespace") + c.SetParamValues("models", "dingo-local") if err := h.PublishFiles(c); err != nil { t.Fatalf("publish handler returned error: %v", err) } diff --git a/internal/handler/upload_inventory_handler.go b/internal/handler/upload_inventory_handler.go new file mode 100644 index 0000000..99b1e1d --- /dev/null +++ b/internal/handler/upload_inventory_handler.go @@ -0,0 +1,24 @@ +package handler + +import ( + "crypto/subtle" + "dingospeed/internal/service" + "github.com/labstack/echo/v4" + "net/http" +) + +// UploadInventory serves only the durable snapshot already produced by the +// reconciler. Scheduler never interprets a half-finished live filesystem scan. +func UploadInventory(c echo.Context) error { + snap, err := service.LoadUploadInventorySnapshot() + if err != nil { + return echo.NewHTTPError(http.StatusServiceUnavailable, "upload inventory is not ready") + } + provided := c.QueryParam("token") + if snap.ReportToken == "" || subtle.ConstantTimeCompare([]byte(provided), []byte(snap.ReportToken)) != 1 { + return echo.NewHTTPError(http.StatusForbidden, "invalid inventory report token") + } + response := *snap + response.ReportToken = "" + return c.JSON(http.StatusOK, &response) +} diff --git a/internal/handler/upload_reconcile.go b/internal/handler/upload_reconcile.go new file mode 100644 index 0000000..4777bed --- /dev/null +++ b/internal/handler/upload_reconcile.go @@ -0,0 +1,28 @@ +package handler + +import ( + "dingospeed/pkg/config" + "dingospeed/pkg/inventory" + "encoding/json" + "github.com/labstack/echo/v4" + "io" + "net/http" +) + +func UploadReconcile(c echo.Context) error { + var p struct { + Epoch string `json:"epoch"` + } + d := json.NewDecoder(http.MaxBytesReader(c.Response(), c.Request().Body, 4096)) + d.DisallowUnknownFields() + if err := d.Decode(&p); err != nil { + return echo.NewHTTPError(400, err.Error()) + } + if d.Decode(&struct{}{}) != io.EOF || len(p.Epoch) != 36 { + return echo.NewHTTPError(400, "invalid reconciliation epoch") + } + if err := inventory.RequestReconcile(config.SysConfig.Repos(), p.Epoch); err != nil { + return echo.NewHTTPError(409, err.Error()) + } + return c.JSON(202, map[string]string{"status": "pending", "epoch": p.Epoch}) +} diff --git a/internal/model/query/query.go b/internal/model/query/query.go index e433cc4..09a5080 100644 --- a/internal/model/query/query.go +++ b/internal/model/query/query.go @@ -1,26 +1,30 @@ package query type CreateCacheJobReq struct { + Org string `json:"org,omitempty"` + OrgRepo string `json:"orgRepo,omitempty"` Type int32 `json:"type"` InstanceId string `json:"instanceId"` Datatype string `json:"datatype"` - Org string `json:"org"` + Namespace string `json:"namespace"` Repo string `json:"repo"` RepositoryId int64 `json:"repositoryId"` } type ResumeCacheJobReq struct { + Org string `json:"org,omitempty"` Id int64 `json:"id"` Type int32 `json:"type"` InstanceId string `json:"instanceId"` Datatype string `json:"datatype"` - Org string `json:"org"` + Namespace string `json:"namespace"` Repo string `json:"repo"` Token string `json:"token"` UsedStorage int64 `json:"usedStorage"` } type JobStatusReq struct { + Intent string `json:"intent,omitempty"` Id int64 `json:"id"` InstanceId string `json:"instanceId"` } diff --git a/internal/router/cache_pause_modelscope_test.go b/internal/router/cache_pause_modelscope_test.go new file mode 100644 index 0000000..178dcc4 --- /dev/null +++ b/internal/router/cache_pause_modelscope_test.go @@ -0,0 +1,90 @@ +package router + +import ( + "context" + "dingospeed/internal/dao" + "dingospeed/internal/data" + "dingospeed/internal/model/query" + "dingospeed/internal/service" + "dingospeed/pkg/app" + "dingospeed/pkg/config" + "dingospeed/pkg/consts" + "fmt" + "github.com/labstack/echo/v4" + "net" + "net/http" + "net/http/httptest" + "strconv" + "strings" + "testing" + "time" +) + +func TestModelScopePauseJoinsLoopbackCacheHandler(t *testing.T) { + entered, canceled := make(chan struct{}), make(chan struct{}) + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if strings.HasSuffix(r.URL.Path, "/repo/files") { + fmt.Fprint(w, `{"Code":200,"Data":{"LatestCommitter":{"Id":"fixed"},"Files":[{"Name":"model.bin","Path":"model.bin","Type":"blob","Size":8192,"Sha256":"file-oid","Revision":"fixed"}]}}`) + return + } + if strings.HasSuffix(r.URL.Path, "/repo") { + w.Header().Set("Content-Length", "8192") + w.WriteHeader(200) + w.(http.Flusher).Flush() + close(entered) + <-r.Context().Done() + close(canceled) + return + } + w.WriteHeader(404) + })) + defer upstream.Close() + read, _ := namespaceEngines(t) + config.SysConfig.Server.Online = true + config.SysConfig.Download.RemoteFileRangeSize = 16384 + config.SysConfig.Modelscope = config.Modelscope{OfficialBaseURL: upstream.URL, MaxRetry: 1} + loopback := httptest.NewServer(read) + defer loopback.Close() + host, port, _ := net.SplitHostPort(strings.TrimPrefix(loopback.URL, "http://")) + config.SysConfig.Server.Host = host + config.SysConfig.Server.Port, _ = strconv.Atoi(port) + base := data.NewBaseData() + locks := dao.NewLockDao(base) + downloads := dao.NewDownloaderDao(nil) + files := dao.NewFileDao(downloads, base, locks) + svc := service.NewCacheJobService(files, dao.NewMetaDao(files, locks, base), downloads, nil) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + request := httptest.NewRequest("POST", "/", nil).WithContext(app.NewContext(ctx, app.New(app.Context(ctx)))) + id, err := svc.CreateCacheJob(echo.New().NewContext(request, httptest.NewRecorder()), &query.CreateCacheJobReq{Type: consts.CacheTypePreheat, Namespace: "modelscope", Datatype: "models", Repo: "owner/model"}) + if err != nil { + t.Fatal(err) + } + select { + case <-entered: + case <-time.After(5 * time.Second): + t.Fatal("download did not start") + } + if err := svc.StopCacheJob(&query.JobStatusReq{Id: id}); err != nil { + t.Fatal(err) + } + deadline := time.Now().Add(5 * time.Second) + for { + snapshot, err := svc.CacheJobStatus(id) + if err != nil { + t.Fatal(err) + } + if snapshot.State == "paused" { + break + } + if time.Now().After(deadline) { + t.Fatalf("did not join cache handler: %+v", snapshot) + } + time.Sleep(time.Millisecond) + } + select { + case <-canceled: + case <-time.After(time.Second): + t.Fatal("source still active after paused confirmation") + } +} diff --git a/internal/router/hf_projection_test.go b/internal/router/hf_projection_test.go new file mode 100644 index 0000000..07fb372 --- /dev/null +++ b/internal/router/hf_projection_test.go @@ -0,0 +1,51 @@ +package router + +import ( + "encoding/hex" + "encoding/json" + "net/http" + "os" + "path/filepath" + "testing" + + "dingospeed/pkg/config" +) + +func TestHFProjectionHTTPDoesNotRegisterOrDownload(t *testing.T) { + e, _ := namespaceEngines(t) + // The server is online, but a projection request must never use its remote + // read-through path (which would register a descriptor or access upstream). + config.SysConfig.Server.Online = true + root := config.SysConfig.Repos() + api := filepath.Join(root, "api", "models", "team", "demo") + meta := filepath.Join(api, "revision", "main", "meta_get.json") + os.MkdirAll(filepath.Dir(meta), 0755) + body := []byte(`{"sha":"commit1","siblings":[{"rfilename":"README.md","size":5,"blobId":"readme"},{"rfilename":"missing.bin","size":9,"blobId":"missing"}]}`) + wrapped, _ := json.Marshal(map[string]any{"status_code": 200, "content": hex.EncodeToString(body)}) + if err := os.WriteFile(meta, wrapped, 0444); err != nil { + t.Fatal(err) + } + blob := filepath.Join(root, "files", "models", "team", "demo", "blobs", "readme") + os.MkdirAll(filepath.Dir(blob), 0755) + if err := os.WriteFile(blob, []byte("hello"), 0444); err != nil { + t.Fatal(err) + } + for _, route := range []string{"local-catalog", "local-manifest?repo=team%2Fdemo&revision=main", "local-file?repo=team%2Fdemo&revision=main&path=README.md"} { + r := namespaceRequest(e, http.MethodGet, "/api/repositories/models/huggingface/"+route, nil) + requireStatus(t, r, 200) + } + r := namespaceRequest(e, http.MethodGet, "/api/repositories/models/huggingface/local-file?repo=team%2Fdemo&revision=main&path=missing.bin", nil) + requireStatus(t, r, 404) + r = namespaceRequest(e, http.MethodGet, "/api/repositories/models/huggingface/local-manifest?repo=team%2Fdemo&revision=absent", nil) + requireStatus(t, r, 404) + r = namespaceRequest(e, http.MethodGet, "/api/repositories/models/datacanvas/local-catalog", nil) + requireStatus(t, r, 400) + r = namespaceRequest(e, http.MethodGet, "/api/repositories/models/huggingface/local-manifest?repo=team%2Fdemo&repo=other&revision=main", nil) + requireStatus(t, r, 400) + if _, err := os.Stat(filepath.Join(api, "repository.json")); !os.IsNotExist(err) { + t.Fatal("read registered HF repository") + } + if _, err := os.Stat(filepath.Join(filepath.Dir(blob), "missing")); !os.IsNotExist(err) { + t.Fatal("missing file was fetched") + } +} diff --git a/internal/router/http_router.go b/internal/router/http_router.go index dfed90f..2da3eda 100644 --- a/internal/router/http_router.go +++ b/internal/router/http_router.go @@ -17,6 +17,7 @@ package router import ( "dingospeed/internal/handler" "dingospeed/pkg/config" + "dingospeed/pkg/middleware" "github.com/labstack/echo/v4" "github.com/prometheus/client_golang/prometheus/promhttp" @@ -46,6 +47,9 @@ func NewHttpRouter(echo *echo.Echo, fileHandler *handler.FileHandler, metaHandle } func (r *HttpRouter) initRouter() { + r.echo.Pre(providerPrefix) + r.echo.Use(middleware.DecodeRouteParams) + r.echo.Use(handler.LimitDownload) // 系统信息 r.echo.GET("/info", r.sysHandler.Info) if config.SysConfig.EnableMetric() { @@ -60,21 +64,30 @@ func (r *HttpRouter) initRouter() { } func (r *HttpRouter) routerForSpeed() { // alayanew - // The local repository contract is opt-in so existing deployments keep the - // exact same route table until their control plane is ready to consume it. - if config.SysConfig.Server.LocalRepositoryAPI { - r.echo.GET("/api/local-repositories/:repoType/:org/:repo/revisions/:revision", r.metaHandler.GetLocalSnapshotHandler) - r.echo.GET("/api/local-repositories/:repoType/:org/:repo/revisions/:revision/archive", r.metaHandler.GetLocalArchiveHandler) - r.echo.HEAD("/api/local-repositories/:repoType/:org/:repo/revisions/:revision/archive", r.metaHandler.GetLocalArchiveHandler) - } + p := "/api/repositories/:repoType/:namespace" + r.echo.GET(p+"/local-catalog", handler.HFLocalCatalog) + r.echo.GET(p+"/local-manifest", handler.HFLocalManifest, handler.RepositoryLocator(false, true)) + r.echo.GET(p+"/local-file", handler.HFLocalFile, handler.RepositoryLocator(true, true)) + r.echo.GET(p+"/directories", r.metaHandler.RepositoryDirectories) + r.echo.GET(p+"/revisions", r.metaHandler.RepositoryRevisions, handler.RepositoryLocator(false, false)) + r.echo.GET(p+"/snapshot", r.metaHandler.RepositorySnapshot, handler.RepositoryLocator(false, true)) + r.echo.GET(p+"/files", r.metaHandler.RepositoryFiles, handler.RepositoryLocator(false, true)) + r.echo.GET(p+"/offset", r.fileHandler.RepositoryOffset, handler.RepositoryLocator(false, false)) + r.echo.GET(p+"/tree", r.metaHandler.RepositoryTree, handler.RepositoryLocator(false, true)) + r.echo.GET(p+"/metadata", r.metaHandler.RepositoryMetadata, handler.RepositoryLocator(false, true)) + r.echo.HEAD(p+"/metadata", r.metaHandler.RepositoryMetadata, handler.RepositoryLocator(false, true)) + r.echo.GET(p+"/archive", r.metaHandler.RepositoryArchive, handler.RepositoryLocator(false, true)) + r.echo.HEAD(p+"/archive", r.metaHandler.RepositoryArchive, handler.RepositoryLocator(false, true)) + r.echo.GET(p+"/file", r.fileHandler.RepositoryFile, handler.RepositoryLocator(true, true)) + r.echo.HEAD(p+"/file", r.fileHandler.RepositoryFile, handler.RepositoryLocator(true, true)) // 单个文件下载 - r.echo.HEAD("/:repoType/:org/:repo/resolve/:commit/:filePath", r.fileHandler.HeadFileHandler1) - r.echo.HEAD("/:orgOrRepoType/:repo/resolve/:commit/:filePath", r.fileHandler.HeadFileHandler2) - r.echo.HEAD("/:repo/resolve/:commit/:filePath", r.fileHandler.HeadFileHandler3) - r.echo.GET("/:repoType/:org/:repo/resolve/:commit/:filePath", r.fileHandler.GetFileHandler1) - r.echo.GET("/:orgOrRepoType/:repo/resolve/:commit/:filePath", r.fileHandler.GetFileHandler2) - r.echo.GET("/:repo/resolve/:commit/:filePath", r.fileHandler.GetFileHandler3) + r.echo.HEAD("/:repoType/:org/:repo/resolve/:commit/*", r.fileHandler.HeadFileHandler1) + r.echo.HEAD("/:orgOrRepoType/:repo/resolve/:commit/*", r.fileHandler.HeadFileHandler2) + r.echo.HEAD("/:repo/resolve/:commit/*", r.fileHandler.HeadFileHandler3) + r.echo.GET("/:repoType/:org/:repo/resolve/:commit/*", r.fileHandler.GetFileHandler1) + r.echo.GET("/:orgOrRepoType/:repo/resolve/:commit/*", r.fileHandler.GetFileHandler2) + r.echo.GET("/:repo/resolve/:commit/*", r.fileHandler.GetFileHandler3) // 模型&数据集元数据 r.echo.HEAD("/api/:repoType/:org/:repo/revision/:revision", r.metaHandler.GetMetadataHandler) @@ -83,22 +96,20 @@ func (r *HttpRouter) routerForSpeed() { // alayanew r.echo.GET("/api/:repoType/:org/:repo/tree/:revision/*", r.metaHandler.GetRepoTreeHandler) // refs + r.echo.GET("/api/:repoType/:org/:repo/refs", r.metaHandler.ProtocolRefs) // r.echo.GET("/api/:repoType/:org/:repo/refs", r.metaHandler.RepoRefsHandler) 修复转发响应码，走统一转发。 r.echo.GET("/api/whoami-v2", r.metaHandler.WhoamiV2Handler) r.echo.GET("/repos", r.metaHandler.ReposHandler) r.echo.Any("/*", r.metaHandler.ForwardToNewSiteHandler) } -func (r *HttpRouter) routerForScheduler() { // alayanew - r.echo.GET("/api/:repoType/:org/:repo/files/:commit/", r.metaHandler.RepositoryFilesHandler) - r.echo.GET("/api/:repoType/:org/:repo/files/:commit/:filePath", r.metaHandler.RepositoryFilesHandler) - - r.echo.GET("/api/fileOffset/:dataType/:org/:repo/:etag/:fileSize", r.fileHandler.GetFileOffset) +func (r *HttpRouter) routerForScheduler() { r.echo.GET("/api/fileProcessSync", r.fileHandler.FileProcessSync) - + r.echo.GET("/api/upload-inventory", handler.UploadInventory) } func (r *HttpRouter) routerForCacheJob() { // alayanew + r.echo.GET("/api/cacheJob/status", r.cacheJobHandler.CacheJobStatus) r.echo.POST("/api/cacheJob/create", r.cacheJobHandler.CreateCacheJobHandler) r.echo.POST("/api/cacheJob/stop", r.cacheJobHandler.StopCacheJobHandler) r.echo.POST("/api/cacheJob/resume", r.cacheJobHandler.ResumeCacheJobHandler) @@ -110,6 +121,7 @@ func (r *HttpRouter) routerForModelscope() { // modelscope r.echo.GET("/api/v1/:repoType/:org/:repo/revisions", r.modelscopeHandler.RevisionsHandler) r.echo.GET("/api/v1/:repoType/:org/:repo/repo/files", r.modelscopeHandler.FileListHandler) r.echo.GET("/api/v1/:repoType/:org/:repo/repo", r.modelscopeHandler.FileDownloadHandler) + r.echo.HEAD("/api/v1/:repoType/:org/:repo/repo", r.modelscopeHandler.FileDownloadHandler) r.echo.GET("/api/v1/:repoType/:org/:repo/repo/tree", r.modelscopeHandler.FileTreeHandler) r.echo.GET("/api/v1/datasets/:datasetId/repo/tree", r.modelscopeHandler.DatasetFileTreeHandler) } diff --git a/internal/router/http_router_test.go b/internal/router/http_router_test.go index 1cba27a..38da911 100644 --- a/internal/router/http_router_test.go +++ b/internal/router/http_router_test.go @@ -1,37 +1,79 @@ package router import ( + "net/http/httptest" + "os" + "strings" "testing" "dingospeed/internal/handler" "dingospeed/pkg/config" + "dingospeed/pkg/prom" + "dingospeed/pkg/util" "github.com/labstack/echo/v4" ) -const localSnapshotRoute = "/api/local-repositories/:repoType/:org/:repo/revisions/:revision" -const localArchiveRoute = localSnapshotRoute + "/archive" - -func TestLocalRepositoryRoutesAreOptIn(t *testing.T) { - oldConfig := config.SysConfig - t.Cleanup(func() { config.SysConfig = oldConfig }) - - routes := func(enabled bool) map[string]bool { - config.SysConfig = &config.Config{Server: config.ServerConfig{LocalRepositoryAPI: enabled}} +func TestStorageMetricsUseExistingRouteAndFlag(t *testing.T) { + old := config.SysConfig + t.Cleanup(func() { config.SysConfig = old }) + // Instantiate an existing metric so its lazy label series is exposed too. + prom.RequestTotalCnt.WithLabelValues("storage-metrics-test").Inc() + t.Cleanup(func() { prom.RequestTotalCnt.DeleteLabelValues("storage-metrics-test") }) + util.ObserveFileAccessFailure("read", &os.PathError{Op: "read", Path: "private-path", Err: os.ErrPermission}) + for _, enabled := range []bool{false, true, true, false} { + config.SysConfig = &config.Config{Server: config.ServerConfig{Metrics: enabled, Repos: "invalid\x00repos"}} e := echo.New() NewHttpRouter(e, handler.NewFileHandler(nil, nil, nil), handler.NewMetaHandler(nil), handler.NewSysHandler(nil), handler.NewCacheJobHandler(nil), handler.NewModelscopeHandler(nil)) - got := map[string]bool{} + found := false for _, route := range e.Routes() { - got[route.Method+" "+route.Path] = true + if route.Method == "GET" && route.Path == "/metrics" { + found = true + } + } + if found != enabled { + t.Fatalf("metrics route enabled=%v want %v", found, enabled) + } + if !enabled { + continue + } // Existing catch-all routing is not changed. + r := httptest.NewRecorder() + e.ServeHTTP(r, httptest.NewRequest("GET", "/metrics", nil)) + if r.Code != 200 { + t.Fatalf("scrape failed: %d %s", r.Code, r.Body.String()) + } + body := r.Body.String() + for _, want := range []string{ + "# TYPE dingospeed_storage_access_errors_total counter", + `dingospeed_storage_access_errors_total{kind="permission_denied",operation="read"} 1`, + `request_total_cnt{source="storage-metrics-test"} 1`, + } { + if !strings.Contains(body, want) { + t.Fatalf("missing %s", want) + } + } + if strings.Contains(body, "private-path") { + t.Fatal("path leaked to metrics") } - return got } +} - if got := routes(false); got["GET "+localSnapshotRoute] || got["GET "+localArchiveRoute] || got["HEAD "+localArchiveRoute] { - t.Fatal("local repository API route is enabled by default") +func TestRepositoryRoutesReplaceLocalAliases(t *testing.T) { + old := config.SysConfig + defer func() { config.SysConfig = old }() + config.SysConfig = &config.Config{} + e := echo.New() + NewHttpRouter(e, handler.NewFileHandler(nil, nil, nil), handler.NewMetaHandler(nil), handler.NewSysHandler(nil), handler.NewCacheJobHandler(nil), handler.NewModelscopeHandler(nil)) + got := map[string]bool{} + for _, r := range e.Routes() { + if strings.Contains(r.Path, "/api/local-") { + t.Fatalf("legacy alias remains: %s", r.Path) + } + got[r.Method+" "+r.Path] = true } - got := routes(true) - if !got["GET "+localSnapshotRoute] || !got["GET "+localArchiveRoute] || !got["HEAD "+localArchiveRoute] { - t.Fatal("local repository API route is missing when enabled") + for _, op := range []string{"snapshot", "archive", "file", "metadata", "tree", "directories", "revisions"} { + if !got["GET /api/repositories/:repoType/:namespace/"+op] { + t.Fatalf("missing %s", op) + } } } diff --git a/internal/router/namespace_integration_test.go b/internal/router/namespace_integration_test.go new file mode 100644 index 0000000..f2b592b --- /dev/null +++ b/internal/router/namespace_integration_test.go @@ -0,0 +1,382 @@ +package router + +import ( + "archive/zip" + "bytes" + "crypto/sha256" + "encoding/hex" + "encoding/json" + "fmt" + "io" + "net/http" + "net/http/httptest" + "net/url" + "os" + "strings" + "sync/atomic" + "testing" + + "dingospeed/internal/dao" + "dingospeed/internal/data" + "dingospeed/internal/downloader" + "dingospeed/internal/handler" + "dingospeed/internal/service" + "dingospeed/pkg/config" + "dingospeed/pkg/repository" + "github.com/labstack/echo/v4" +) + +func namespaceEngines(t *testing.T) (*echo.Echo, *echo.Echo) { + t.Helper() + old := config.SysConfig + t.Cleanup(func() { config.SysConfig = old }) + config.SysConfig = &config.Config{Server: config.ServerConfig{Repos: t.TempDir()}, Upload: config.Upload{Namespace: "dingo-local", ConcurrentLimit: 4}, Download: config.Download{GoroutineMaxNumPerFile: 4, BlockSize: 1024, RespChanSize: 32, RespChunkSize: 1024, RemoteFileRangeSize: 4096, RemoteFileBufferSize: 4096, ReqTimeout: 5}, Retry: config.Retry{Attempts: 1}} + base := data.NewBaseData() + lock := dao.NewLockDao(base) + file := dao.NewFileDao(dao.NewDownloaderDao(nil), base, lock) + meta := handler.NewMetaHandler(service.NewMetaService(file, dao.NewMetaDao(file, lock, base))) + read := echo.New() + NewHttpRouter(read, handler.NewFileHandler(service.NewFileService(file), nil, nil), meta, handler.NewSysHandler(nil), handler.NewCacheJobHandler(nil), handler.NewModelscopeHandler(service.NewModelscopeService())) + upload := echo.New() + NewUploadRouter(UploadEcho{upload}, handler.NewUploadHandler(service.NewUploadService(dao.NewUploadDao(file, lock))), handler.NewCacheAdminHandler(service.NewCacheAdminService(dao.NewCacheAdminDao(file)))) + return read, upload +} +func namespaceRequest(e *echo.Echo, method, target string, body []byte) *httptest.ResponseRecorder { + req := httptest.NewRequest(method, target, bytes.NewReader(body)) + rec := httptest.NewRecorder() + e.ServeHTTP(rec, req) + return rec +} +func requireStatus(t *testing.T, r *httptest.ResponseRecorder, status int) { + t.Helper() + if r.Code != status { + t.Fatalf("status %d want %d: %s", r.Code, status, r.Body.String()) + } +} +func digest(content []byte) string { s := sha256.Sum256(content); return hex.EncodeToString(s[:]) } +func locator(repo, revision, path string) string { + q := url.Values{"repo": {repo}, "revision": {revision}} + if path != "" { + q.Set("path", path) + } + return q.Encode() +} + +func TestNamespacesUploadPublishReadAndArchive(t *testing.T) { + read, upload := namespaceEngines(t) + const repo = "team/resolve/model-a" + files := []string{"weights/nested/model.bin", "notes/50%.txt", "notes/literal%2F.txt"} + for _, ns := range []string{"alice", "bob", "datacanvas", "用户#50%"} { + routeNS := url.PathEscape(ns) + items := []map[string]interface{}{} + for _, file := range files { + content := []byte(ns + ":" + file) + q := locator(repo, "main", file) + "&size=" + fmt.Sprint(len(content)) + "&sha256=" + digest(content) + "&defer=true" + r := namespaceRequest(upload, "POST", "/api/uploads/models/"+routeNS+"?"+q, content) + requireStatus(t, r, 201) + var got map[string]interface{} + json.Unmarshal(r.Body.Bytes(), &got) + if got["namespace"] != ns || got["repo"] != repo || got["commit"] != "" { + t.Fatalf("staging identity wrong: %s", r.Body.String()) + } + r = namespaceRequest(upload, "GET", "/api/upload-progress/models/"+routeNS+"?"+locator(repo, "main", file)+"&sha256="+digest(content), nil) + requireStatus(t, r, 200) + json.Unmarshal(r.Body.Bytes(), &got) + if got["blobComplete"] != true || got["effective"] != false { + t.Fatalf("staging progress wrong: %s", r.Body.String()) + } + items = append(items, map[string]interface{}{"path": file, "size": len(content), "sha256": digest(content)}) + } + body, _ := json.Marshal(map[string]interface{}{"files": items}) + r := namespaceRequest(upload, "POST", "/api/publish/models/"+routeNS+"?"+locator(repo, "main", ""), body) + requireStatus(t, r, 201) + r = namespaceRequest(read, "GET", "/api/repositories/models/"+routeNS+"/snapshot?"+locator(repo, "main", ""), nil) + requireStatus(t, r, 200) + var snap map[string]interface{} + json.Unmarshal(r.Body.Bytes(), &snap) + if snap["namespace"] != ns || snap["repo"] != repo || len(snap["files"].([]interface{})) != len(files) { + t.Fatalf("snapshot identity wrong: %s", r.Body.String()) + } + for _, file := range files { + r = namespaceRequest(read, "GET", "/api/repositories/models/"+routeNS+"/file?"+locator(repo, "main", file), nil) + requireStatus(t, r, 200) + if r.Body.String() != ns+":"+file { + t.Fatalf("cross-namespace file: %s", r.Body.String()) + } + requireStatus(t, namespaceRequest(read, "HEAD", "/api/repositories/models/"+routeNS+"/file?"+locator(repo, "main", file), nil), 200) + } + r = namespaceRequest(read, "GET", "/api/repositories/models/"+routeNS+"/archive?"+locator(repo, "main", ""), nil) + requireStatus(t, r, 200) + archive, err := zip.NewReader(bytes.NewReader(r.Body.Bytes()), int64(r.Body.Len())) + if err != nil { + t.Fatal(err) + } + if len(archive.File) != len(files) { + t.Fatalf("archive file count %d", len(archive.File)) + } + for _, entry := range archive.File { + f, err := entry.Open() + if err != nil { + t.Fatal(err) + } + b, _ := io.ReadAll(f) + f.Close() + if !strings.HasPrefix(string(b), ns+":") { + t.Fatal("archive crossed namespaces") + } + } + r = namespaceRequest(read, "GET", "/api/repositories/models/"+routeNS+"/directories", nil) + requireStatus(t, r, 200) + if !strings.Contains(r.Body.String(), `"repo":"`+repo+`"`) { + t.Fatal(r.Body.String()) + } + r = namespaceRequest(read, "GET", "/api/repositories/models/"+routeNS+"/revisions?repo="+url.QueryEscape(repo), nil) + requireStatus(t, r, 200) + if !strings.Contains(r.Body.String(), `"name":"main"`) { + t.Fatal(r.Body.String()) + } + } + // Publish-tree edits only alice. The other namespace and the immutable base + // commit remain independently readable while the new branch head changes. + base := namespaceRequest(read, "GET", "/api/repositories/models/alice/snapshot?"+locator(repo, "main", ""), nil) + requireStatus(t, base, 200) + var snap struct { + Commit string `json:"commit"` + } + if err := json.Unmarshal(base.Body.Bytes(), &snap); err != nil { + t.Fatal(err) + } + changed := []byte("alice edited content") + newPath := "edits/repository.json" + staged := namespaceRequest(upload, "POST", "/api/uploads/models/alice?"+locator(repo, "main", newPath)+"&size="+fmt.Sprint(len(changed))+"&sha256="+digest(changed)+"&defer=true", changed) + requireStatus(t, staged, 201) + body, _ := json.Marshal(map[string]interface{}{"baseCommit": snap.Commit, "files": []map[string]interface{}{{"path": newPath, "size": len(changed), "sha256": digest(changed)}}}) + updated := namespaceRequest(upload, "POST", "/api/publish-tree/models/alice?"+locator(repo, "main", ""), body) + requireStatus(t, updated, 201) + if !strings.Contains(updated.Body.String(), `"removed":3`) { + t.Fatal(updated.Body.String()) + } + requireStatus(t, namespaceRequest(read, "GET", "/api/repositories/models/alice/file?"+locator(repo, "main", files[0]), nil), 404) + requireStatus(t, namespaceRequest(read, "GET", "/api/repositories/models/alice/file?"+locator(repo, snap.Commit, files[0]), nil), 200) + untouched := namespaceRequest(read, "GET", "/api/repositories/models/bob/file?"+locator(repo, "main", files[0]), nil) + requireStatus(t, untouched, 200) + if untouched.Body.String() != "bob:"+files[0] { + t.Fatal("tree publish changed bob") + } + // A different edit based on the superseded head must fail its optimistic lock. + body, _ = json.Marshal(map[string]interface{}{"baseCommit": snap.Commit, "files": []interface{}{}}) + requireStatus(t, namespaceRequest(upload, "POST", "/api/publish-tree/models/alice?"+locator(repo, "main", ""), body), 409) + +} + +func TestRepositoryAPIRejectsBadLocatorsWithoutServiceToken(t *testing.T) { + read, upload := namespaceEngines(t) + for _, q := range []string{"repo=a&repo=b&revision=main&path=x", "repo=a&revision=main&revision=v1&path=x", "repo=a&revision=main&path=x&path=y", "repo=a//b&revision=main&path=x", "repo=a/../b&revision=main&path=x", "repo=a&revision=../main&path=x", "repo=a&revision=main&path=../secret", "repo=a&revision=main&path=x&namespace=bob", "repo=a&revision=main&path=x%ZZ"} { + requireStatus(t, namespaceRequest(upload, "POST", "/api/uploads/models/alice?"+q, []byte("x")), 400) + } + for _, ns := range []string{"huggingface", "modelscope"} { + requireStatus(t, namespaceRequest(upload, "POST", "/api/uploads/models/"+ns+"?repo=Qwen/Qwen3&revision=main&path=a&size=1&sha256="+digest([]byte("x")), []byte("x")), 400) + } + requireStatus(t, namespaceRequest(read, "GET", "/api/repositories/models/unknown/snapshot?repo=team/model&revision=main", nil), 404) + requireStatus(t, namespaceRequest(read, "GET", "/api/local-repositories/models/dingo-local/demo/revisions/main", nil), 404) + requireStatus(t, namespaceRequest(upload, "GET", "/api/cache/summary", nil), 200) +} + +func TestHFAndModelScopeProtocolIsolationWithMockUpstreams(t *testing.T) { + read, _ := namespaceEngines(t) + hfContent := []byte("huggingface file content") + msContent := []byte("modelscope original bytes") + var denied atomic.Bool + var hfRequests, msRequests atomic.Int64 + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Header.Get("X-Dingo-Service-Token") != "" { + t.Error("machine token leaked upstream") + } + if strings.Contains(r.URL.Path, "huggingface/") || strings.Contains(r.URL.Path, "modelscope/") { + t.Error("platform namespace leaked upstream") + } + if denied.Load() { + w.WriteHeader(403) + return + } + path := strings.TrimPrefix(r.URL.Path, "/prefix") + w.Header().Set("Content-Type", "application/json") + if strings.HasPrefix(path, "/api/v1/") { + msRequests.Add(1) + if path == "/api/v1/models/Qwen/Qwen3/repo" { + if r.URL.Query().Get("FilePath") != "nested/weights.bin" { + t.Error("MS file path lost") + } + w.Header().Set("Content-Type", "application/octet-stream") + if r.Header.Get("Range") == "bytes=0-0" { + w.Header().Set("Content-Range", fmt.Sprintf("bytes 0-0/%d", len(msContent))) + w.WriteHeader(206) + w.Write(msContent[:1]) + return + } + w.Header().Set("Content-Length", fmt.Sprint(len(msContent))) + w.Write(msContent) + return + } + if path == "/api/v1/models/Qwen/Qwen3/revisions" { + fmt.Fprint(w, `{"Code":200,"Data":{"RevisionMap":{"Branches":[{"Revision":"master","CommitId":"ms-commit-one"}],"Tags":[]}}}`) + return + } + if path == "/api/v1/models/Qwen/Qwen3/repo/files" { + fmt.Fprintf(w, `{"Code":200,"Data":{"Files":[{"Type":"blob","Path":"nested/weights.bin","Name":"weights.bin","Size":%d,"Sha256":%q}]}}`, len(msContent), digest(msContent)) + return + } + fmt.Fprint(w, `{"Code":200,"Data":{"Id":"Qwen/Qwen3"}}`) + return + } + hfRequests.Add(1) + switch { + case strings.HasPrefix(path, "/api/models/Qwen/Qwen3/revision/"): + fmt.Fprintf(w, `{"id":"Qwen/Qwen3","sha":"commit-one","siblings":[{"rfilename":"nested/weights.bin"}],"usedStorage":%d}`, len(hfContent)) + case path == "/api/models/Qwen/Qwen3/paths-info/commit-one": + body, _ := io.ReadAll(r.Body) + if !strings.Contains(string(body), "nested/weights.bin") { + t.Error("HF nested file path lost") + } + fmt.Fprintf(w, `[{"type":"file","oid":%q,"size":%d,"path":"nested/weights.bin","lfs":{"oid":%q,"size":%d}}]`, digest(hfContent), len(hfContent), digest(hfContent), len(hfContent)) + case path == "/Qwen/Qwen3/resolve/commit-one/nested/weights.bin": + w.Header().Set("Content-Type", "application/octet-stream") + w.Header().Set("Content-Length", fmt.Sprint(len(hfContent))) + w.Write(hfContent) + case path == "/api/models/Qwen/Qwen3/tree/main": + if r.URL.Query().Get("cursor") == "next" { + fmt.Fprint(w, `[{"type":"file","path":"page-two.txt","size":1}]`) + return + } + w.Header().Set("Link", `; rel="next"`) + fmt.Fprint(w, `[{"type":"file","path":"nested/weights.bin","size":23}]`) + case path == "/api/models/Qwen/Qwen3/refs": + fmt.Fprint(w, `{"branches":[{"name":"main","targetCommit":"commit-one"}]}`) + default: + t.Errorf("unexpected upstream path: %s", r.URL.Path) + http.NotFound(w, r) + } + })) + defer upstream.Close() + if err := dao.RegisterHosted("models", "Qwen", "Qwen3"); err != nil { + t.Fatal(err) + } + config.SysConfig.Server.Online = true + config.SysConfig.Server.HfScheme = "http" + config.SysConfig.Server.HfNetLoc = strings.TrimPrefix(upstream.URL, "http://") + "/prefix" + config.SysConfig.Modelscope.OfficialBaseURL = upstream.URL + "/prefix" + config.SysConfig.Modelscope.MaxRetry = 1 + config.SysConfig.Modelscope.ChunkSize = 1024 + originalRoot := config.SysConfig.Server.Repos + for _, cold := range []struct{ method, operation string }{{"HEAD", "file"}, {"GET", "file"}, {"GET", "tree"}} { + config.SysConfig.Server.Repos = t.TempDir() + k := repository.RepoKey{Namespace: repository.HuggingFace, RepoType: "models", Repo: "Qwen/Qwen3"} + if _, err := repository.Read(config.SysConfig.Repos(), k); !os.IsNotExist(err) { + t.Fatal("cold test already registered") + } + target := "/api/repositories/models/huggingface/" + cold.operation + "?" + locator("Qwen/Qwen3", "main", "nested/weights.bin") + if cold.operation == "tree" { + target = "/api/repositories/models/huggingface/tree?" + locator("Qwen/Qwen3", "main", "") + } + response := namespaceRequest(read, cold.method, target, nil) + requireStatus(t, response, 200) + if cold.operation == "tree" && !strings.Contains(response.Body.String(), "page-two.txt") { + t.Fatal("canonical tree lost pagination") + } + if cold.method == "GET" && cold.operation == "file" && !bytes.Equal(response.Body.Bytes(), hfContent) { + t.Fatal("cold canonical HF file mismatch") + } + if _, err := repository.Read(config.SysConfig.Repos(), k); err != nil { + t.Fatal("cold HF read did not register source", err) + } + } + config.SysConfig.Server.Repos = originalRoot + for _, path := range []string{"/api/models/Qwen/Qwen3/revision/main", "/api/models/Qwen/Qwen3/tree/main?recursive=true", "/api/models/Qwen/Qwen3/refs"} { + requireStatus(t, namespaceRequest(read, "GET", "/huggingface"+path, nil), 200) + requireStatus(t, namespaceRequest(read, "GET", path, nil), 200) + } + paged := namespaceRequest(read, "GET", "/huggingface/api/models/Qwen/Qwen3/tree/main?recursive=true", nil) + requireStatus(t, paged, 200) + if link := paged.Header().Get("Link"); !strings.Contains(link, "/huggingface/api/models/Qwen/Qwen3/tree/main?cursor=next") || strings.Contains(link, "/prefix/") { + t.Fatalf("HF pagination lost explicit endpoint: %s", link) + } + r := namespaceRequest(read, "GET", "/huggingface/Qwen/Qwen3/resolve/main/nested/weights.bin", nil) + requireStatus(t, r, 200) + if !bytes.Equal(r.Body.Bytes(), hfContent) { + t.Fatalf("HF payload mismatch: %q", r.Body.String()) + } + plain := namespaceRequest(read, "GET", "/Qwen/Qwen3/resolve/main/nested/weights.bin", nil) + requireStatus(t, plain, 200) + if !bytes.Equal(plain.Body.Bytes(), hfContent) { + t.Fatal("hosted namespace intercepted original upstream download") + } + for _, path := range []string{"/api/v1/models/Qwen/Qwen3", "/api/v1/models/Qwen/Qwen3/revisions", "/api/v1/models/Qwen/Qwen3/repo/files"} { + requireStatus(t, namespaceRequest(read, "GET", "/modelscope"+path, nil), 200) + } + msURL := "/modelscope/api/v1/models/Qwen/Qwen3/repo?Revision=master&FilePath=nested%2Fweights.bin" + r = namespaceRequest(read, "GET", msURL, nil) + if r.Code != 200 && r.Code != 206 { + t.Fatal(r.Body.String()) + } + if !bytes.Equal(r.Body.Bytes(), msContent) { + t.Fatalf("MS payload mismatch: %q", r.Body.String()) + } + hf := repository.RepoKey{Namespace: repository.HuggingFace, RepoType: "models", Repo: "Qwen/Qwen3"} + ms := repository.RepoKey{Namespace: repository.ModelScope, RepoType: "models", Repo: "Qwen/Qwen3"} + for _, key := range []repository.RepoKey{hf, ms} { + if _, err := repository.Read(config.SysConfig.Repos(), key); err != nil { + t.Fatal(err) + } + } + if _, err := os.Stat(hf.Blob(config.SysConfig.Repos(), digest(hfContent))); err != nil { + t.Fatal(err) + } + fh, err := os.Open(ms.Blob(config.SysConfig.Repos(), digest(msContent))) + if err != nil { + t.Fatal(err) + } + header := &downloader.DingCacheHeader{} + if err = header.Read(fh); err != nil { + t.Fatal(err) + } + raw, err := io.ReadAll(fh) + fh.Close() + if err != nil || !bytes.Equal(raw, msContent) { + t.Fatalf("MS dingcache payload wrong: %v %q", err, raw) + } + if _, err = os.Stat(ms.Resolve(config.SysConfig.Repos(), "ms-commit-one", "nested/weights.bin")); err != nil { + t.Fatal(err) + } + + for _, op := range []string{"metadata", "tree", "files", "revisions"} { + target := "/api/repositories/models/modelscope/" + op + "?repo=Qwen%2FQwen3&revision=master" + canonical := namespaceRequest(read, "GET", target, nil) + requireStatus(t, canonical, 200) + if op == "metadata" && (!strings.Contains(canonical.Body.String(), `"rfilename":"nested/weights.bin"`) || !strings.Contains(canonical.Body.String(), `"sha":"ms-commit-one"`)) { + t.Fatalf("ModelScope canonical metadata: %s", canonical.Body.String()) + } + if op == "revisions" && !strings.Contains(canonical.Body.String(), `"name":"master"`) { + t.Fatalf("ModelScope canonical revisions: %s", canonical.Body.String()) + } + } + canonicalFile := "/api/repositories/models/modelscope/file?" + locator("Qwen/Qwen3", "master", "nested/weights.bin") + canonical := namespaceRequest(read, "GET", canonicalFile, nil) + if canonical.Code != 200 && canonical.Code != 206 { + t.Fatal(canonical.Body.String()) + } + if !bytes.Equal(canonical.Body.Bytes(), msContent) { + t.Fatalf("canonical ModelScope payload wrong: %q", canonical.Body.String()) + } + requireStatus(t, namespaceRequest(read, "HEAD", canonicalFile, nil), 200) + if hfRequests.Load() == 0 || msRequests.Load() == 0 { + t.Fatal("missing upstream protocol requests") + } + denied.Store(true) + for _, op := range []string{"metadata", "tree", "revisions"} { + requireStatus(t, namespaceRequest(read, "GET", "/api/repositories/models/modelscope/"+op+"?repo=Qwen%2FQwen3&revision=master", nil), 403) + } + requireStatus(t, namespaceRequest(read, "GET", "/api/repositories/models/huggingface/tree?repo=Qwen%2FQwen3&revision=main", nil), 403) + requireStatus(t, namespaceRequest(read, "GET", canonicalFile, nil), 403) + requireStatus(t, namespaceRequest(read, "GET", "/huggingface/Qwen/Qwen3/resolve/main/nested/weights.bin", nil), 403) + requireStatus(t, namespaceRequest(read, "GET", msURL, nil), 403) +} diff --git a/internal/router/provider_prefix.go b/internal/router/provider_prefix.go new file mode 100644 index 0000000..7dd7e94 --- /dev/null +++ b/internal/router/provider_prefix.go @@ -0,0 +1,47 @@ +package router + +import ( + "github.com/labstack/echo/v4" + "strings" +) + +// Explicit provider endpoints stay unambiguous when an upstream owner has the +// same spelling as a registered hosted namespace. +func providerPrefix(next echo.HandlerFunc) echo.HandlerFunc { + return func(c echo.Context) error { + req := c.Request() + for _, provider := range []string{"huggingface", "modelscope"} { + prefix := "/" + provider + parts := strings.Split(strings.TrimPrefix(req.URL.Path, "/"), "/") + file := len(parts) >= 6 && parts[0] == provider && parts[3] == "resolve" + typedFile := len(parts) >= 7 && parts[0] == provider && (parts[1] == "datasets" || parts[1] == "spaces") && parts[4] == "resolve" + explicit := provider == "huggingface" && (strings.HasPrefix(req.URL.Path, prefix+"/api/") || file || typedFile) || provider == "modelscope" && strings.HasPrefix(req.URL.Path, prefix+"/api/v1/") + if explicit { + c.Set("forcedProvider", provider) + c.Set("providerPrefix", prefix) + req.URL.Path = strings.TrimPrefix(req.URL.Path, prefix) + if req.URL.RawPath != "" { + req.URL.RawPath = strings.TrimPrefix(req.URL.RawPath, prefix) + } + if provider == "modelscope" && !strings.HasPrefix(req.URL.Path, "/api/v1/") { + return echo.NewHTTPError(404, "unsupported ModelScope protocol route") + } + break + } + } + // Hosted HF-compatible downloads use an explicit endpoint. A personal + // namespace must never intercept the same upstream HF organization. + parts := strings.Split(strings.TrimPrefix(req.URL.Path, "/"), "/") + localFile := len(parts) >= 6 && parts[0] == "dingo-local" && parts[3] == "resolve" + localTypedFile := len(parts) >= 7 && parts[0] == "dingo-local" && (parts[1] == "datasets" || parts[1] == "spaces") && parts[4] == "resolve" + if strings.HasPrefix(req.URL.Path, "/dingo-local/api/") || localFile || localTypedFile { + c.Set("forcedLocal", true) + c.Set("providerPrefix", "/dingo-local") + req.URL.Path = strings.TrimPrefix(req.URL.Path, "/dingo-local") + if req.URL.RawPath != "" { + req.URL.RawPath = strings.TrimPrefix(req.URL.RawPath, "/dingo-local") + } + } + return next(c) + } +} diff --git a/internal/router/provider_prefix_test.go b/internal/router/provider_prefix_test.go new file mode 100644 index 0000000..ce00f08 --- /dev/null +++ b/internal/router/provider_prefix_test.go @@ -0,0 +1,34 @@ +package router + +import ( + "net/http/httptest" + "testing" + + "github.com/labstack/echo/v4" +) + +func TestProviderPrefixPreservesOriginalOwnerNames(t *testing.T) { + for _, tc := range []struct { + input, path string + local bool + }{ + {"/huggingface/demo/resolve/main/weights.bin", "/huggingface/demo/resolve/main/weights.bin", false}, + {"/modelscope/demo/resolve/main/nested/weights.bin", "/modelscope/demo/resolve/main/nested/weights.bin", false}, + {"/huggingface/team/demo/resolve/main/weights.bin", "/team/demo/resolve/main/weights.bin", false}, + {"/dingo-local/demo/resolve/main/nested/weights.bin", "/dingo-local/demo/resolve/main/nested/weights.bin", false}, + {"/dingo-local/alice/demo/resolve/main/weights.bin", "/alice/demo/resolve/main/weights.bin", true}, + {"/dingo-local/datasets/alice/demo/resolve/main/weights.bin", "/datasets/alice/demo/resolve/main/weights.bin", true}, + {"/dingo-local/api/models/alice/demo/revision/main", "/api/models/alice/demo/revision/main", true}, + } { + c := echo.New().NewContext(httptest.NewRequest("GET", tc.input, nil), httptest.NewRecorder()) + err := providerPrefix(func(c echo.Context) error { + if c.Request().URL.Path != tc.path || (c.Get("forcedLocal") == true) != tc.local { + t.Fatalf("%s resolved to %s local=%v", tc.input, c.Request().URL.Path, c.Get("forcedLocal")) + } + return nil + })(c) + if err != nil { + t.Fatal(err) + } + } +} diff --git a/internal/router/upload_router.go b/internal/router/upload_router.go index 79fef8b..5e8400d 100644 --- a/internal/router/upload_router.go +++ b/internal/router/upload_router.go @@ -2,6 +2,7 @@ package router import ( "dingospeed/internal/handler" + "dingospeed/pkg/middleware" "github.com/labstack/echo/v4" ) @@ -18,18 +19,24 @@ type UploadEcho struct { func NewUploadRouter(uploadEcho UploadEcho, uploadHandler *handler.UploadHandler, cacheAdminHandler *handler.CacheAdminHandler) *UploadRouter { r := &UploadRouter{echo: uploadEcho.Echo, uploadHandler: uploadHandler, cacheAdminHandler: cacheAdminHandler} + r.echo.Use(middleware.DecodeRouteParams, handler.CacheQueryGuard) r.initRouter() r.initCacheAdminRouter() return r } // initCacheAdminRouter 挂缓存管理。放在上传服务而不是下载服务上： -// 这些接口会删数据，必须待在只监听回环地址、且要求上传 token 的那个端口后面。 +// 这些接口经控制台转发到上传端口，保留既有跨站写入检查。 // // 只有 JSON 接口，没有页面：管理界面在 spinfield 控制台里（web/console 的 // 「缓存管理」页），经 ingest agent 转发到这里。dingospeed 自己不再提供页面， // 否则同一套界面会有两份实现，各改各的。 func (r *UploadRouter) initCacheAdminRouter() { + r.echo.POST("/api/upload-inventory/reconcile", handler.UploadReconcile) + r.echo.GET("/api/scheduler-registration", handler.SchedulerRegistration) + r.echo.PUT("/api/scheduler-registration", handler.SchedulerRegistration) + r.echo.GET("/api/transfer-settings", handler.TransferSettings) + r.echo.PUT("/api/transfer-settings", handler.TransferSettings) r.echo.GET("/api/cache/summary", r.cacheAdminHandler.Summary) r.echo.GET("/api/cache/repos", r.cacheAdminHandler.ListRepos) r.echo.GET("/api/cache/files", r.cacheAdminHandler.ListFiles) @@ -39,14 +46,12 @@ func (r *UploadRouter) initCacheAdminRouter() { } func (r *UploadRouter) initRouter() { - r.echo.GET("/api/local-upload-progress/:repoType/:org/:repo/:revision/*", r.uploadHandler.QueryProgress) - r.echo.POST("/api/local-upload/:repoType/:org/:repo/:revision/*", r.uploadHandler.UploadWholeFile) - // 分块是幂等的（重复写同一块直接返回 already_present），用 PUT 而不是 POST。 - r.echo.PUT("/api/local-upload-chunk/:repoType/:org/:repo/:revision/*", r.uploadHandler.UploadChunk) - r.echo.POST("/api/local-publish/:repoType/:org/:repo/:revision", r.uploadHandler.PublishFiles) - // 整树发布：清单是目标全量而非增量，因此能表达删除。仓库编辑走这条路径， - // 新建上传继续走 local-publish（那条做并集合并，仓库/revision 可以不存在）。 - r.echo.POST("/api/local-publish-tree/:repoType/:org/:repo/:revision", r.uploadHandler.PublishTree) + r.echo.DELETE("/api/repositories/:repoType/:namespace", r.uploadHandler.DeleteRepository, handler.RepositoryLocator(false, false)) + r.echo.GET("/api/upload-progress/:repoType/:namespace", r.uploadHandler.QueryProgress, handler.RepositoryLocator(true, true)) + r.echo.POST("/api/uploads/:repoType/:namespace", r.uploadHandler.UploadWholeFile, handler.RepositoryLocator(true, true), handler.LimitUpload) + r.echo.PUT("/api/upload-chunks/:repoType/:namespace", r.uploadHandler.UploadChunk, handler.RepositoryLocator(true, true), handler.LimitUpload) + r.echo.POST("/api/publish/:repoType/:namespace", r.uploadHandler.PublishFiles, handler.RepositoryLocator(false, true)) + r.echo.POST("/api/publish-tree/:repoType/:namespace", r.uploadHandler.PublishTree, handler.RepositoryLocator(false, true)) } func (r *UploadRouter) Echo() *echo.Echo { diff --git a/internal/server/scheduler.go b/internal/server/scheduler.go index 42aeb7b..ea00451 100644 --- a/internal/server/scheduler.go +++ b/internal/server/scheduler.go @@ -4,7 +4,9 @@ import ( "context" "crypto/tls" "crypto/x509" + "errors" "io/ioutil" + "sync" "dingospeed/internal/service" "dingospeed/pkg/config" @@ -16,7 +18,7 @@ import ( ) type SchedulerServer struct { - conn *grpc.ClientConn + conn *registrationConnection schedulerService *service.SchedulerService localOperationService *service.LocalOperationService sysService *service.SysService @@ -31,28 +33,88 @@ func NewSchedulerServer(schedulerService *service.SchedulerService, sysService * } func (s *SchedulerServer) Start(ctx context.Context) error { - if !config.SysConfig.IsCluster() { - return nil - } - ssl := config.SysConfig.Server.Ssl - creds := credential(ssl.CrtFile, ssl.KeyFile, ssl.CaFile, "zetyun.com") - conn, err := grpc.NewClient(config.SysConfig.Scheduler.Addr, grpc.WithTransportCredentials(creds)) - if err != nil { - zap.S().Errorf("连接失败: %v", err) - return err - } + config.SysConfig.Registration() + config.SysConfig.SetSchedulerModel("standalone") + conn := ®istrationConnection{} s.conn = conn client := manager.NewManagerClient(conn) - s.schedulerService.Client = client + s.schedulerService.BindClient(client) s.sysService.Client = client s.schedulerService.Ctx = ctx - s.schedulerService.Register() + go s.schedulerService.Register() s.localOperationService.Ctx = ctx s.localOperationService.Initialize() return nil } +// One stable client is shared by all services. Configuration changes replace only +// its transport, not service pointers or background consumers. +type registrationConnection struct { + mu sync.Mutex + conn *grpc.ClientConn + address string + closed bool +} + +func (c *registrationConnection) current() (*grpc.ClientConn, error) { + c.mu.Lock() + defer c.mu.Unlock() + r := config.SysConfig.Registration() + if c.closed { + return nil, errors.New("scheduler client stopped") + } + if !r.Enabled { + if c.conn != nil { + _ = c.conn.Close() + c.conn = nil + } + return nil, errors.New("scheduler registration disabled") + } + if c.conn != nil && c.address == r.Address { + return c.conn, nil + } + if c.conn != nil { + _ = c.conn.Close() + c.conn = nil + } + ssl := config.SysConfig.Server.Ssl + creds := credential(ssl.CrtFile, ssl.KeyFile, ssl.CaFile, "zetyun.com") + if creds == nil { + return nil, errors.New("scheduler TLS certificate/key/CA unavailable") + } + conn, err := grpc.NewClient(r.Address, grpc.WithTransportCredentials(creds)) + if err != nil { + return nil, err + } + c.conn = conn + c.address = r.Address + return conn, nil +} +func (c *registrationConnection) Invoke(ctx context.Context, method string, args, reply any, opts ...grpc.CallOption) error { + conn, err := c.current() + if err != nil { + return err + } + return conn.Invoke(ctx, method, args, reply, opts...) +} +func (c *registrationConnection) NewStream(ctx context.Context, desc *grpc.StreamDesc, method string, opts ...grpc.CallOption) (grpc.ClientStream, error) { + conn, err := c.current() + if err != nil { + return nil, err + } + return conn.NewStream(ctx, desc, method, opts...) +} +func (c *registrationConnection) Close() error { + c.mu.Lock() + defer c.mu.Unlock() + c.closed = true + if c.conn != nil { + return c.conn.Close() + } + return nil +} + func credential(crtFile, keyFile, caFile, svcName string) credentials.TransportCredentials { cert, err := tls.LoadX509KeyPair(crtFile, keyFile) if err != nil { diff --git a/internal/server/storage_probe.go b/internal/server/storage_probe.go new file mode 100644 index 0000000..34356f2 --- /dev/null +++ b/internal/server/storage_probe.go @@ -0,0 +1,78 @@ +package server + +import ( + "context" + "errors" + "path/filepath" + "sync" + "time" + + "dingospeed/pkg/config" + "dingospeed/pkg/dependency" + "dingospeed/pkg/storageprobe" + "github.com/prometheus/client_golang/prometheus" +) + +type StorageProbeServer struct { + config *config.Config + registry prometheus.Registerer + mu sync.Mutex + probe *storageprobe.Probe + stopped bool + started bool +} + +func NewStorageProbeServer(c *config.Config) *StorageProbeServer { + return &StorageProbeServer{config: c, registry: prometheus.DefaultRegisterer} +} + +func (s *StorageProbeServer) Start(ctx context.Context) error { + s.mu.Lock() + defer s.mu.Unlock() + if s.stopped || s.started { + return nil + } + if !s.config.StorageProbe.Enabled { + if err := s.registry.Register(dependency.ImpactCollector(dependency.Default, nil)); err != nil { + return err + } + s.started = true + return nil + } + if s.config.Server.Repos == "" { + return errors.New("storage probe requires server.repos") + } + // Only lexical resolution here: even startup must not await a broken mount. + root, err := filepath.Abs(s.config.Server.Repos) + if err != nil { + return err + } + read, write := storageprobe.Filesystem(root) + p, err := storageprobe.New(read, write, 10*time.Second, 3*time.Second) + if err != nil { + return err + } + probeCollector := storageprobe.Collector(p) + if err = s.registry.Register(probeCollector); err != nil { + return err + } + if err = s.registry.Register(dependency.ImpactCollector(dependency.Default, p.Evidence)); err != nil { + s.registry.Unregister(probeCollector) + return err + } + s.probe = p + s.started = true + p.Start(ctx) + return nil +} + +func (s *StorageProbeServer) Stop(context.Context) error { + s.mu.Lock() + s.stopped = true + p := s.probe + s.mu.Unlock() + if p != nil { + p.Stop() + } + return nil +} diff --git a/internal/server/storage_probe_test.go b/internal/server/storage_probe_test.go new file mode 100644 index 0000000..ba9946f --- /dev/null +++ b/internal/server/storage_probe_test.go @@ -0,0 +1,60 @@ +package server + +import ( + "context" + "testing" + + "dingospeed/pkg/config" + "github.com/prometheus/client_golang/prometheus" +) + +func TestStorageProbeDisabledHasNoSideEffects(t *testing.T) { + c := &config.Config{} + c.Server.Repos = "invalid/unused/root" + s := NewStorageProbeServer(c) + registry := prometheus.NewPedanticRegistry() + s.registry = registry + if err := s.Start(context.Background()); err != nil { + t.Fatal(err) + } + if s.probe != nil { + t.Fatal("disabled probe started") + } + metrics, err := registry.Gather() + if err != nil || len(metrics) != 2 { + t.Fatal("disabled probe should only register memory-only impact metrics") + } + for _, f := range metrics { + if f.GetName() == "dingospeed_storage_impact_probe_enabled" && f.Metric[0].Gauge.GetValue() != 0 { + t.Fatal("disabled probe reported enabled") + } + } + if err := s.Stop(context.Background()); err != nil { + t.Fatal(err) + } +} + +func TestStorageProbeLifecycle(t *testing.T) { + c := &config.Config{} + c.StorageProbe.Enabled = true + c.Server.Repos = t.TempDir() // No sentinel: Start must not perform filesystem I/O. + s := NewStorageProbeServer(c) + registry := prometheus.NewPedanticRegistry() + s.registry = registry + defer s.Stop(context.Background()) + for i := 0; i < 2; i++ { + if err := s.Start(context.Background()); err != nil { + t.Fatal(err) + } + } + metrics, err := registry.Gather() + if err != nil || len(metrics) != 6 { + t.Fatalf("unexpected collector registration: %d %v", len(metrics), err) + } + if err := s.Stop(context.Background()); err != nil { + t.Fatal(err) + } + if err := s.Start(context.Background()); err != nil { + t.Fatal(err) + } +} diff --git a/internal/service/baseline_retry_test.go b/internal/service/baseline_retry_test.go new file mode 100644 index 0000000..277ecf0 --- /dev/null +++ b/internal/service/baseline_retry_test.go @@ -0,0 +1,241 @@ +package service + +import ( + "bytes" + "context" + "crypto/sha256" + "dingospeed/internal/dao" + "dingospeed/internal/data" + "dingospeed/internal/downloader" + "dingospeed/pkg/config" + "dingospeed/pkg/inventory" + "dingospeed/pkg/repository" + "encoding/hex" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "testing" + "time" +) + +type baselineFixture struct { + svc *SchedulerService + root, epoch string + fail bool + reports [][]byte + statuses []string + upload func() +} + +func newBaselineFixture(t *testing.T) *baselineFixture { + t.Helper() + f := &baselineFixture{root: t.TempDir(), epoch: "11111111-1111-1111-1111-111111111111"} + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/api/v1/upload-inventory/nodes/node-a/session": + json.NewEncoder(w).Encode(inventory.Session{InstanceID: "node-a", PendingEpoch: f.epoch}) + case "/api/v1/upload-inventory/nodes/node-a/reconcile-progress": + var p map[string]string + json.NewDecoder(r.Body).Decode(&p) + f.statuses = append(f.statuses, p["status"]) + json.NewEncoder(w).Encode(map[string]bool{"ok": true}) + case "/api/v1/upload-inventory/reports": + b, _ := io.ReadAll(r.Body) + f.reports = append(f.reports, b) + if f.fail { + w.WriteHeader(503) + return + } + var p inventory.Report + json.Unmarshal(b, &p) + json.NewEncoder(w).Encode(inventory.Ack{Epoch: p.Epoch, Sequence: p.Sequence, Digest: p.Digest(), Status: "accepted"}) + default: + http.NotFound(w, r) + } + })) + t.Cleanup(server.Close) + old := config.SysConfig + t.Cleanup(func() { config.SysConfig = old }) + config.SysConfig = &config.Config{Server: config.ServerConfig{Repos: f.root}, Download: config.Download{BlockSize: 16}, Upload: config.Upload{Namespace: "dingo-local"}, Scheduler: config.Scheduler{HTTPURL: server.URL, Discovery: config.Discovery{InstanceId: "node-a"}}} + base := data.NewBaseData() + locks := dao.NewLockDao(base) + files := dao.NewFileDao(nil, base, locks) + uploads := dao.NewUploadDao(files, locks) + f.svc = &SchedulerService{metaService: NewMetaService(files, nil)} + f.upload = func() { + t.Helper() + body := []byte("new commit") + hash := sha256.Sum256(body) + _, err := uploads.UploadWholeFile(dao.LocalUploadParam{Namespace: "team", RepoType: "models", Repo: "changed", Revision: "main", FilePath: "a", Size: int64(len(body)), Sha256: hex.EncodeToString(hash[:])}, bytes.NewReader(body)) + if err != nil { + t.Fatal(err) + } + } + if err := f.svc.openInventorySession(context.Background()); err != nil { + t.Fatal(err) + } + return f +} +func retryNow(t *testing.T, root string) { + t.Helper() + if err := inventory.Update(root, func(s *inventory.State) error { s.RetryAt = time.Time{}; return nil }); err != nil { + t.Fatal(err) + } +} + +func TestBaselineTrustsPublishedMetadataWithoutHashingPayload(t *testing.T) { + f := newBaselineFixture(t) + f.upload() + hash := sha256.Sum256([]byte("new commit")) + digest := hex.EncodeToString(hash[:]) + blob := dao.RepositoryKey("models", "team/changed").Blob(f.root, digest) + cache, err := downloader.NewDingCache(blob, 16) + if err != nil { + t.Fatal(err) + } + // Change the payload while preserving size and complete block metadata. + // Inventory reconciliation trusts the published SHA256, not these bytes. + err = cache.WriteBlock(0, bytes.Repeat([]byte("x"), 16)) + cache.Close() + if err != nil { + t.Fatal(err) + } + if err = f.svc.SyncPublications(context.Background()); err != nil { + t.Fatal(err) + } + if len(f.reports) != 1 { + t.Fatalf("reports=%d, want one baseline", len(f.reports)) + } + var p inventory.Report + if err = json.Unmarshal(f.reports[0], &p); err != nil { + t.Fatal(err) + } + if !p.Baseline || len(p.Files) != 1 || p.Files[0].SHA256 != digest || p.Files[0].Size != 10 { + t.Fatalf("published metadata was not preserved: %+v", p) + } + q, err := inventory.Read(f.root) + if err != nil || q.ReconcileEpoch != "" || len(q.Entries) != 0 { + t.Fatalf("baseline not confirmed: %+v %v", q, err) + } +} + +func TestBaselineRejectsMissingPublishedFile(t *testing.T) { + f := newBaselineFixture(t) + f.upload() + hash := sha256.Sum256([]byte("new commit")) + blob := dao.RepositoryKey("models", "team/changed").Blob(f.root, hex.EncodeToString(hash[:])) + if err := os.Remove(blob); err != nil { + t.Fatal(err) + } + if err := f.svc.SyncPublications(context.Background()); err == nil { + t.Fatal("missing published file accepted") + } + if len(f.reports) != 0 { + t.Fatal("incomplete baseline sent to Scheduler") + } +} +func TestBaselineBuildFailureStopsScanningAndExplicitResetResumes(t *testing.T) { + f := newBaselineFixture(t) + ctx := context.Background() + bad := filepath.Join(f.root, "api", "models", repository.Local, "broken", repository.Marker) + if err := os.MkdirAll(filepath.Dir(bad), 0755); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(bad, []byte("invalid"), 0644); err != nil { + t.Fatal(err) + } + for i := 0; i < maxBaselineBuildAttempts; i++ { + retryNow(t, f.root) + if err := f.svc.SyncPublications(ctx); err == nil { + t.Fatal("expected build failure") + } + } + q, err := inventory.Read(f.root) + if err != nil || !q.BuildBlocked || q.BuildAttempts != 3 { + t.Fatalf("budget not persisted: %+v %v", q, err) + } + if len(f.reports) != 0 || f.statuses[len(f.statuses)-1] != "needs_attention" { + t.Fatal("partial report sent or attention not reported") + } + if err := os.Remove(bad); err != nil { + t.Fatal(err) + } + f.upload() + // Reconnection and duplicate delivery must not reset the budget. + if err := f.svc.openInventorySession(ctx); err != nil { + t.Fatal(err) + } + if err := inventory.RequestReconcile(f.root, f.epoch); err != nil { + t.Fatal(err) + } + if err := f.svc.SyncPublications(ctx); err != nil { + t.Fatal(err) + } + q, _ = inventory.Read(f.root) + if !q.BuildBlocked || len(q.Entries) != 1 || len(f.reports) != 0 { + t.Fatal("blocked baseline consumed new commit or scanned again") + } + if nextInventoryWake() < time.Hour { + t.Fatal("blocked baseline busy polls") + } + f.epoch = "22222222-2222-2222-2222-222222222222" + if err := inventory.RequestReconcile(f.root, f.epoch); err != nil { + t.Fatal(err) + } + if err := f.svc.SyncPublications(ctx); err != nil { + t.Fatal(err) + } + q, _ = inventory.Read(f.root) + if q.BuildBlocked || q.BuildAttempts != 0 || len(q.Entries) != 0 || len(f.reports) != 1 { + t.Fatalf("reset failed: %+v", q) + } + var p inventory.Report + json.Unmarshal(f.reports[0], &p) + if len(p.Files) != 1 { + t.Fatal("new commit missing from rebuilt baseline") + } +} +func TestBaselineNetworkRetryKeepsBytesAndLaterCommit(t *testing.T) { + f := newBaselineFixture(t) + ctx := context.Background() + f.fail = true + if err := f.svc.SyncPublications(ctx); err == nil { + t.Fatal("expected HTTP failure") + } + fixed := append([]byte(nil), f.reports[0]...) + f.upload() + for i := 0; i < 4; i++ { + retryNow(t, f.root) + if err := f.svc.SyncPublications(ctx); err == nil { + t.Fatal("expected HTTP failure") + } + } + q, _ := inventory.Read(f.root) + if q.BuildBlocked || q.BuildAttempts != 0 { + t.Fatal("network error consumed build budget") + } + f.fail = false + retryNow(t, f.root) + if err := f.svc.SyncPublications(ctx); err != nil { + t.Fatal(err) + } + for _, b := range f.reports { + if !bytes.Equal(b, fixed) { + t.Fatal("baseline retry changed") + } + } + q, _ = inventory.Read(f.root) + if len(q.Entries) != 1 { + t.Fatal("baseline ACK removed later commit") + } + if err := f.svc.SyncPublications(ctx); err != nil { + t.Fatal(err) + } + q, _ = inventory.Read(f.root) + if len(q.Entries) != 0 { + t.Fatal("later commit not delivered") + } +} diff --git a/internal/service/cache_admin_service.go b/internal/service/cache_admin_service.go index f3388e4..919bb95 100644 --- a/internal/service/cache_admin_service.go +++ b/internal/service/cache_admin_service.go @@ -34,14 +34,15 @@ func NewCacheAdminService(cacheAdminDao *dao.CacheAdminDao) *CacheAdminService { // CacheQuery 是两个列表接口共用的查询条件。 type CacheQuery struct { - RepoType string - OrgRepo string - Source string - Keyword string - Sort string - Order string - Page int - PageSize int + Namespace string + RepoType string + OrgRepo string + Source string + Keyword string + Sort string + Order string + Page int + PageSize int } type CacheSummary struct { @@ -103,6 +104,9 @@ func (s *CacheAdminService) ListFiles(query CacheQuery) (*CacheFilePage, error) rows := s.cacheAdminDao.ListFiles(query.RepoType, query.OrgRepo) filtered := make([]*dao.CacheFileRow, 0, len(rows)) for _, row := range rows { + if query.Namespace != "" && row.Namespace != query.Namespace { + continue + } if query.Source != "" && row.Source != query.Source { continue } @@ -127,6 +131,9 @@ func (s *CacheAdminService) ListOrphans(query CacheQuery) (*RecyclePage, error) filtered := make([]*dao.RecycleRow, 0, len(rows)) var totalSize int64 for _, row := range rows { + if query.Namespace != "" && row.Namespace != query.Namespace { + continue + } if query.Source != "" && row.Source != query.Source { continue } @@ -150,14 +157,22 @@ func (s *CacheAdminService) SoftDelete(items []dao.DeleteItem) ([]*dao.DeleteRes if len(items) == 0 { return nil, uploadError{status: 400, code: "CACHE_INVALID_ARGUMENT", msg: "items is empty"} } - return s.cacheAdminDao.SoftDelete(items) + result, err := s.cacheAdminDao.SoftDelete(items) + if err == nil { + dao.NotifyPublished() + } + return result, err } func (s *CacheAdminService) PurgeOrphans(items []dao.DeleteItem) ([]*dao.DeleteResult, error) { if len(items) == 0 { return nil, uploadError{status: 400, code: "CACHE_INVALID_ARGUMENT", msg: "items is empty"} } - return s.cacheAdminDao.PurgeOrphans(items) + result, err := s.cacheAdminDao.PurgeOrphans(items) + if err == nil { + dao.NotifyPublished() + } + return result, err } // sumDistinctBlobs 按 (repoType, orgRepo, sha) 去重后求和，与 ListRepos/Summary diff --git a/internal/service/cache_job_dedup.go b/internal/service/cache_job_dedup.go new file mode 100644 index 0000000..8dfcf4b --- /dev/null +++ b/internal/service/cache_job_dedup.go @@ -0,0 +1,72 @@ +package service + +import ( + "errors" + "os" + "path/filepath" + "strconv" + "strings" + + "dingospeed/internal/dao" + "dingospeed/pkg/config" + "dingospeed/pkg/repository" +) + +// Data stays numeric for existing cacheJob/create clients. +type CacheJobCreateResult struct { + ID int64 `json:"data"` + Msg string `json:"msg"` + Disposition string `json:"disposition"` + Commit string `json:"commit,omitempty"` +} + +// Called with createMu held through submission. Each Speed owns its node's +// worker pool; all control-plane replicas reach this same admission boundary. +func (p *CacheJobService) reusableCacheJob(key repository.RepoKey, snapshot *dao.CommitHfSha, local bool, instanceID string) (*CacheJobStatus, error) { + entries, err := os.ReadDir(filepath.Join(config.SysConfig.Repos(), "cache-jobs")) + if errors.Is(err, os.ErrNotExist) { + return nil, nil + } + if err != nil { + return nil, err + } + var completed *CacheJobStatus + for _, entry := range entries { + if entry.IsDir() || !strings.HasSuffix(entry.Name(), ".json") { + continue + } + id, err := strconv.ParseInt(strings.TrimSuffix(entry.Name(), ".json"), 10, 64) + if err != nil || id <= 0 { + continue + } + status, err := p.CacheJobStatus(id) + if errors.Is(err, os.ErrNotExist) { + continue + } + if err != nil { + return nil, err + } + if status.RepoKey != key || status.Commit != snapshot.Sha { + continue + } + switch status.State { + case "running", "pausing", "resuming", "canceling": + return &status, nil + case "waiting": + if _, queued := p.cachePool.GetTask(int(id)); queued { + return &status, nil + } + case "complete": + if status.Local != local || (!local && status.InstanceID != instanceID) { + continue + } + if completed == nil || status.ID > completed.ID { + completed = &status + } + } + } + if completed != nil && dao.PreheatCacheComplete(key, snapshot) { + return completed, nil + } + return nil, nil +} diff --git a/internal/service/cache_job_dedup_test.go b/internal/service/cache_job_dedup_test.go new file mode 100644 index 0000000..1949946 --- /dev/null +++ b/internal/service/cache_job_dedup_test.go @@ -0,0 +1,257 @@ +package service + +import ( + "context" + "crypto/sha256" + "fmt" + "net/http" + "net/http/httptest" + "net/url" + "os" + "strconv" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "dingospeed/internal/dao" + "dingospeed/internal/data" + "dingospeed/internal/model/query" + "dingospeed/pkg/app" + "dingospeed/pkg/common" + "dingospeed/pkg/config" + "dingospeed/pkg/consts" + "dingospeed/pkg/repository" + "github.com/labstack/echo/v4" +) + +func waitCacheJob(t *testing.T, svc *CacheJobService, id int64, want string) { + t.Helper() + deadline := time.Now().Add(5 * time.Second) + for time.Now().Before(deadline) { + status, err := svc.CacheJobStatus(id) + if err != nil { + t.Fatal(err) + } + if status.State == want { + return + } + time.Sleep(10 * time.Millisecond) + } + status, err := svc.CacheJobStatus(id) + t.Fatalf("want %s, got %+v %v", want, status, err) +} + +func TestPreheatConcurrentRequestsReuseAndRevalidate(t *testing.T) { + var version atomic.Int32 + version.Store(1) + gate := make(chan struct{}) + var releaseOnce sync.Once + release := func() { releaseOnce.Do(func() { close(gate) }) } + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch { + case strings.Contains(r.URL.Path, "/revision/"): + fmt.Fprintf(w, `{"sha":"commit-%d","siblings":[{"rfilename":"model.bin"}]}`, version.Load()) + case strings.Contains(r.URL.Path, "/paths-info/"): + fmt.Fprint(w, `[{"type":"file","path":"model.bin","oid":"file-oid","size":8}]`) + case strings.Contains(r.URL.Path, "/resolve/"): + select { + case <-gate: + case <-r.Context().Done(): + return + } + w.Header().Set("Content-Length", "8") + fmt.Fprint(w, "abcdefgh") + default: + w.WriteHeader(404) + } + })) + defer upstream.Close() + msConfig(t, upstream.URL) + config.SysConfig.Server.HfScheme = "http" + config.SysConfig.Server.HfNetLoc = strings.TrimPrefix(upstream.URL, "http://") + config.SysConfig.Scheduler.Mode = "standalone" + base := data.NewBaseData() + locks := dao.NewLockDao(base) + downloads := dao.NewDownloaderDao(nil) + files := dao.NewFileDao(downloads, base, locks) + svc := NewCacheJobService(files, dao.NewMetaDao(files, locks, base), downloads, nil) + defer svc.cachePool.Close() + defer release() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + create := func() (CacheJobCreateResult, error) { + req := httptest.NewRequest("POST", "/api/cacheJob/create", nil).WithContext(app.NewContext(ctx, app.New(app.Context(ctx)))) + return svc.CreateCacheJobResult(echo.New().NewContext(req, httptest.NewRecorder()), &query.CreateCacheJobReq{Type: consts.CacheTypePreheat, Namespace: "huggingface", Datatype: "models", Repo: "owner/model"}) + } + // Start all callers together, including the first one: admission must be atomic. + const callers = 12 + results := make([]CacheJobCreateResult, callers) + errs := make([]error, callers) + var wg sync.WaitGroup + for i := range results { + wg.Add(1) + go func(i int) { defer wg.Done(); results[i], errs[i] = create() }(i) + } + wg.Wait() + created := 0 + for i, result := range results { + if errs[i] != nil { + t.Fatal(errs[i]) + } + if result.ID != results[0].ID { + t.Fatalf("duplicate job IDs: %+v", results) + } + if result.Disposition == "created" { + created++ + } else if result.Disposition != "running" { + t.Fatalf("unexpected result: %+v", result) + } + } + if created != 1 { + t.Fatalf("created %d jobs", created) + } + release() + id := results[0].ID + waitCacheJob(t, svc, id, "complete") + result, err := create() + if err != nil || result.ID != id || result.Disposition != "cached" { + t.Fatalf("completed reuse: %+v %v", result, err) + } + + // Removing a reference after completion must allow the cache to be filled again. + if err := os.Remove(dao.ResolvePath("models", "huggingface/owner/model", "commit-1", "model.bin")); err != nil { + t.Fatal(err) + } + result, err = create() + if err != nil || result.ID == id || result.Disposition != "created" { + t.Fatalf("deleted cache reuse: %+v %v", result, err) + } + waitCacheJob(t, svc, result.ID, "complete") + version.Store(2) + next, err := create() + if err != nil || next.ID == result.ID || next.Disposition != "created" || next.Commit != "commit-2" { + t.Fatalf("new commit: %+v %v", next, err) + } + waitCacheJob(t, svc, next.ID, "complete") + // A restarted worker pool can reuse a complete persisted job, but not an orphaned runner. + restarted := &CacheJobService{cachePool: common.NewPool(1, true)} + defer restarted.cachePool.Close() + metadata, err := svc.preheatMetadata(ctx, repository.RepoKey{Namespace: "huggingface", RepoType: "models", Repo: "owner/model"}, "main", "") + if err != nil { + t.Fatal(err) + } + reused, err := restarted.reusableCacheJob(repository.RepoKey{Namespace: "huggingface", RepoType: "models", Repo: "owner/model"}, metadata, true, "") + if err != nil || reused == nil || reused.ID != next.ID { + t.Fatalf("restart reuse: %+v %v", reused, err) + } +} + +func TestReusablePreheatIdentityAndTerminalStates(t *testing.T) { + for _, provider := range []string{"huggingface", "modelscope"} { + t.Run(provider, func(t *testing.T) { + msConfig(t, "http://unused.invalid") + svc := &CacheJobService{cachePool: common.NewPool(1, true)} + defer svc.cachePool.Close() + key := repository.RepoKey{Namespace: provider, RepoType: "models", Repo: "owner/model"} + snapshot := &dao.CommitHfSha{Sha: "commit-1"} + for _, state := range []string{"failed", "interrupted", "stopped", "running", "waiting"} { + if err := writeCacheJobStatus(CacheJobStatus{ID: 42, RepoKey: key, Commit: snapshot.Sha, State: state, Local: true}); err != nil { + t.Fatal(err) + } + got, err := svc.reusableCacheJob(key, snapshot, true, "") + if err != nil || got != nil { + t.Fatalf("reused %s: %+v %v", state, got, err) + } + } + if err := writeCacheJobStatus(CacheJobStatus{ID: 42, RepoKey: key, Commit: snapshot.Sha, State: "complete", InstanceID: "node-a"}); err != nil { + t.Fatal(err) + } + got, err := svc.reusableCacheJob(key, snapshot, false, "node-a") + if err != nil || got == nil { + t.Fatalf("not reused: %+v %v", got, err) + } + for _, change := range []string{"node", "provider", "repo", "commit", "mode"} { + otherKey, otherSnapshot, node, local := key, *snapshot, "node-a", false + switch change { + case "node": + node = "node-b" + case "provider": + otherKey.Namespace = "other-provider" + case "repo": + otherKey.Repo = "another/model" + case "commit": + otherSnapshot.Sha = "commit-2" + case "mode": + local = true + } + got, err := svc.reusableCacheJob(otherKey, &otherSnapshot, local, node) + if err != nil || got != nil { + t.Fatalf("cross-%s reuse: %+v %v", change, got, err) + } + } + }) + } +} + +func TestModelScopePreheatDedupThroughRepositoryAPI(t *testing.T) { + const payload = "modelscope payload" + oid := fmt.Sprintf("%x", sha256.Sum256([]byte(payload))) + var bodies atomic.Int32 + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if strings.HasSuffix(r.URL.Path, "/repo/files") { + fmt.Fprintf(w, `{"Code":200,"Data":{"LatestCommitter":{"Id":"ms-commit"},"Files":[{"Type":"blob","Path":"model.bin","Size":%d,"Sha256":%q,"Revision":"ms-commit"}]}}`, len(payload), oid) + return + } + bodies.Add(1) + w.Header().Set("Content-Length", fmt.Sprint(len(payload))) + fmt.Fprint(w, payload) + })) + defer upstream.Close() + msConfig(t, upstream.URL) + config.SysConfig.Scheduler.Mode = "standalone" + key := repository.RepoKey{Namespace: repository.ModelScope, RepoType: "models", Repo: "owner/model"} + provider := NewModelscopeService() + e := echo.New() + e.GET("/api/repositories/models/modelscope/metadata", func(c echo.Context) error { + meta, err := provider.RepositoryMetadata(c, key, c.QueryParam("revision")) + if err != nil { + return err + } + return c.JSON(200, meta) + }) + e.GET("/api/repositories/models/modelscope/file", func(c echo.Context) error { + return provider.RepositoryFile(c, key, c.QueryParam("revision"), c.QueryParam("path")) + }) + server := httptest.NewServer(e) + defer server.Close() + address, err := url.Parse(server.URL) + if err != nil { + t.Fatal(err) + } + config.SysConfig.Server.Host = address.Hostname() + config.SysConfig.Server.Port, err = strconv.Atoi(address.Port()) + if err != nil { + t.Fatal(err) + } + svc := NewCacheJobService(nil, nil, nil, nil) + defer svc.cachePool.Close() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + create := func() CacheJobCreateResult { + req := httptest.NewRequest("POST", "/", nil).WithContext(app.NewContext(ctx, app.New(app.Context(ctx)))) + result, err := svc.CreateCacheJobResult(echo.New().NewContext(req, httptest.NewRecorder()), &query.CreateCacheJobReq{Type: consts.CacheTypePreheat, Namespace: key.Namespace, Datatype: key.RepoType, Repo: key.Repo}) + if err != nil { + t.Fatal(err) + } + return result + } + first := create() + waitCacheJob(t, svc, first.ID, "complete") + before := bodies.Load() + second := create() + if second.ID != first.ID || second.Disposition != "cached" || bodies.Load() != before { + t.Fatalf("ModelScope repeated download: first=%+v second=%+v bodies=%d/%d", first, second, before, bodies.Load()) + } +} diff --git a/internal/service/cache_job_failed_retry_test.go b/internal/service/cache_job_failed_retry_test.go new file mode 100644 index 0000000..c268307 --- /dev/null +++ b/internal/service/cache_job_failed_retry_test.go @@ -0,0 +1,186 @@ +package service + +import ( + "bytes" + "context" + "crypto/sha256" + "errors" + "fmt" + "net/http" + "net/http/httptest" + "net/url" + "strconv" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "dingospeed/internal/dao" + "dingospeed/internal/model/query" + "dingospeed/internal/service/task" + "dingospeed/pkg/app" + "dingospeed/pkg/config" + "dingospeed/pkg/consts" + "dingospeed/pkg/repository" + "dingospeed/pkg/util" + "github.com/labstack/echo/v4" +) + +func TestModelScopeFailedTaskRetriesPinnedCache(t *testing.T) { + payload := bytes.Repeat([]byte("retry-data-123456"), 1024) + oid := fmt.Sprintf("%x", sha256.Sum256(payload)) + var bodies, resumedOffset atomic.Int64 + var recoverSource atomic.Bool + gate := make(chan struct{}) + var once sync.Once + release := func() { once.Do(func() { close(gate) }) } + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if strings.HasSuffix(r.URL.Path, "/repo/files") { + if recoverSource.Load() && r.URL.Query().Get("Revision") != "fixed" { + t.Errorf("retry changed pinned revision: %s", r.URL) + } + fmt.Fprintf(w, `{"Code":200,"Data":{"LatestCommitter":{"Id":"fixed"},"Files":[{"Type":"blob","Path":"model.bin","Size":%d,"Sha256":%q,"Revision":"fixed"}]}}`, len(payload), oid) + return + } + bodies.Add(1) + if r.URL.Query().Get("Revision") != "fixed" { + t.Error("file request changed commit") + } + start, end, err := util.FileRange(r.Header.Get("Range"), int64(len(payload))) + if err != nil { + t.Error(err) + w.WriteHeader(400) + return + } + w.Header().Set("Content-Length", fmt.Sprint(end-start)) + if start > 0 { + w.Header().Set("Content-Range", fmt.Sprintf("bytes %d-%d/%d", start, end-1, len(payload))) + w.WriteHeader(206) + } + if !recoverSource.Load() { + w.Write(payload[:4096]) + w.(http.Flusher).Flush() + select { + case <-gate: + case <-r.Context().Done(): + } + return // Deliberate short body after the first blocks have reached disk. + } + resumedOffset.Store(start) + w.Write(payload[start:end]) + })) + defer upstream.Close() + defer release() + msConfig(t, upstream.URL) + config.SysConfig.Download.RemoteFileRangeSize = 0 + key := repository.RepoKey{Namespace: repository.ModelScope, RepoType: "models", Repo: "owner/model"} + provider := NewModelscopeService() + e := echo.New() + e.GET("/api/repositories/models/modelscope/metadata", func(c echo.Context) error { + meta, err := provider.RepositoryMetadata(c, key, c.QueryParam("revision")) + if err != nil { + return err + } + return c.JSON(200, meta) + }) + e.GET("/api/repositories/models/modelscope/file", func(c echo.Context) error { + return provider.RepositoryFile(c, key, c.QueryParam("revision"), c.QueryParam("path")) + }) + server := httptest.NewServer(e) + defer server.Close() + address, _ := url.Parse(server.URL) + config.SysConfig.Server.Host = address.Hostname() + config.SysConfig.Server.Port, _ = strconv.Atoi(address.Port()) + svc := NewCacheJobService(nil, nil, nil, nil) + defer svc.cachePool.Close() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + caller := func() echo.Context { + req := httptest.NewRequest("POST", "/", nil).WithContext(app.NewContext(ctx, app.New(app.Context(ctx)))) + return echo.New().NewContext(req, httptest.NewRecorder()) + } + first, err := svc.CreateCacheJobResult(caller(), &query.CreateCacheJobReq{Type: consts.CacheTypePreheat, Namespace: key.Namespace, Datatype: key.RepoType, Repo: key.Repo}) + if err != nil { + t.Fatal(err) + } + deadline := time.Now().Add(5 * time.Second) + for { + status, err := svc.CacheJobStatus(first.ID) + if err != nil { + t.Fatal(err) + } + if status.CachedBytes >= 4096 { + break + } + if time.Now().After(deadline) { + t.Fatalf("no cached blocks: %+v", status) + } + time.Sleep(time.Millisecond) + } + release() + waitCacheJob(t, svc, first.ID, "failed") + failed, _ := svc.CacheJobStatus(first.ID) + if !strings.Contains(failed.Error, "model.bin") || !strings.Contains(failed.Error, "upstream response ended") || failed.CachedBytes < 4096 { + t.Fatalf("failure lost cause or cache: %+v", failed) + } + for { + if _, active := svc.cachePool.GetTask(int(first.ID)); !active { + break + } + if time.Now().After(deadline) { + t.Fatal("failed worker still owns task") + } + time.Sleep(time.Millisecond) + } + recoverSource.Store(true) + if err := svc.ResumeCacheJob(caller(), &query.ResumeCacheJobReq{Id: first.ID}); err != nil { + t.Fatal(err) + } + waitCacheJob(t, svc, first.ID, "complete") + final, _ := svc.CacheJobStatus(first.ID) + if final.Commit != "fixed" || !final.Local || final.InstanceID != "" || final.Generation != failed.Generation+1 || final.Error != "" || final.CachedBytes != uint64(len(payload)) || resumedOffset.Load() != 4096 || bodies.Load() != 2 { + t.Fatalf("retry changed identity or lost cached blocks: %+v offset=%d requests=%d", final, resumedOffset.Load(), bodies.Load()) + } + actual, _ := msPayload(t, key, oid) + if !bytes.Equal(actual, payload) { + t.Fatal("retried cached file differs from source") + } + t.Logf("failed -> complete; task=%d commit=%s generations=%d->%d resumedOffset=%d sha256=%s", first.ID, final.Commit, failed.Generation, final.Generation, resumedOffset.Load(), oid) +} + +func TestFailedTaskResumeKeepsAdmissionGuards(t *testing.T) { + msConfig(t, "http://unused.invalid") + svc := NewCacheJobService(nil, nil, nil, nil) + defer svc.cachePool.Close() + key := repository.RepoKey{Namespace: repository.ModelScope, RepoType: "models", Repo: "owner/model"} + for _, tc := range []struct{ state, commit, want string }{ + {"failed", "", "missing_pinned_commit"}, + {"complete", "fixed", "invalid_task_state"}, + {"canceled", "fixed", "invalid_task_state"}, + } { + if err := writeCacheJobStatus(CacheJobStatus{ID: 991, RepoKey: key, State: tc.state, Commit: tc.commit, Local: true}); err != nil { + t.Fatal(err) + } + err := svc.ResumeCacheJob(nil, &query.ResumeCacheJobReq{Id: 991}) + var conflict *CacheJobConflict + if !errors.As(err, &conflict) || conflict.Code != tc.want { + t.Fatalf("state=%s error=%v", tc.state, err) + } + } + if err := writeCacheJobStatus(CacheJobStatus{ID: 991, RepoKey: key, State: "failed", Commit: "fixed", Local: true}); err != nil { + t.Fatal(err) + } + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + active := &task.PreheatCacheTask{CacheTask: task.CacheTask{TaskNo: 992, Ctx: ctx, CancelFunc: cancel, LocalJob: true}, Sha: &dao.CommitHfSha{Sha: "fixed"}} + if err := trackCacheJob(active, key, context.Background()); err != nil { + t.Fatal(err) + } + defer active.OnFinish(context.Canceled) + err := svc.ResumeCacheJob(nil, &query.ResumeCacheJobReq{Id: 991}) + var conflict *CacheJobConflict + if !errors.As(err, &conflict) || conflict.Code != "active_job_conflict" || conflict.ActiveJobID != 992 { + t.Fatalf("failed task bypassed single-writer guard: %v", err) + } +} diff --git a/internal/service/cache_job_metrics.go b/internal/service/cache_job_metrics.go new file mode 100644 index 0000000..ef15475 --- /dev/null +++ b/internal/service/cache_job_metrics.go @@ -0,0 +1,41 @@ +package service + +import ( + "github.com/prometheus/client_golang/prometheus" + "github.com/prometheus/client_golang/prometheus/promauto" + "time" +) + +var cacheStopDuration = promauto.NewHistogramVec(prometheus.HistogramOpts{ + Name: "dingospeed_cache_stop_seconds", Help: "Time from accepted stop intent to joined download completion", Buckets: []float64{.01, .05, .1, .5, 1, 2, 5, 10, 30}, +}, []string{"intent"}) +var cacheTransitions = promauto.NewCounterVec(prometheus.CounterOpts{ + Name: "dingospeed_cache_transitions_total", Help: "Durably committed cache lifecycle transitions", +}, []string{"state"}) +var cachePersistenceFailures = promauto.NewCounter(prometheus.CounterOpts{ + Name: "dingospeed_cache_snapshot_write_failures_total", Help: "Failed durable cache snapshot writes", +}) + +func init() { + for _, state := range []string{"pausing", "resuming", "canceling"} { + promauto.NewGaugeFunc(prometheus.GaugeOpts{Name: "dingospeed_cache_transition_oldest_seconds", Help: "Oldest active cache transition", ConstLabels: prometheus.Labels{"state": state}}, func() float64 { + cacheStateMu.Lock() + defer cacheStateMu.Unlock() + oldest := 0.0 + for _, execution := range cacheExecutions { + if execution.state != state { + continue + } + since := execution.started + if !execution.stopStarted.IsZero() { + since = execution.stopStarted + } + age := time.Since(since).Seconds() + if age > oldest { + oldest = age + } + } + return oldest + }) + } +} diff --git a/internal/service/cache_job_pause_integration_test.go b/internal/service/cache_job_pause_integration_test.go new file mode 100644 index 0000000..b28b369 --- /dev/null +++ b/internal/service/cache_job_pause_integration_test.go @@ -0,0 +1,160 @@ +package service + +import ( + "context" + "crypto/sha256" + "dingospeed/internal/dao" + "dingospeed/internal/data" + "dingospeed/internal/model/query" + "dingospeed/pkg/app" + "dingospeed/pkg/config" + "dingospeed/pkg/consts" + "dingospeed/pkg/hfprojection" + "fmt" + "github.com/labstack/echo/v4" + "io" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" + "time" +) + +func TestPreheatPauseResumeClosesHTTPAndReusesPinnedBlocks(t *testing.T) { + for i := 0; i < 3; i++ { + t.Run(fmt.Sprint(i), testPreheatPauseResume) + } +} +func testPreheatPauseResume(t *testing.T) { + testPreheatPauseResumeRegistration(t, false) +} + +func TestPreheatResumeAfterJoiningScheduler(t *testing.T) { + testPreheatPauseResumeRegistration(t, true) +} + +func testPreheatPauseResumeRegistration(t *testing.T, joinScheduler bool) { + payload := strings.Repeat("abcdefgh", 1024) + var requests, resumedOffset atomic.Int64 + sourceCanceled := make(chan struct{}) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch { + case strings.Contains(r.URL.Path, "/revision/"): + fmt.Fprint(w, `{"sha":"fixed","siblings":[{"rfilename":"model.bin"}]}`) + case strings.Contains(r.URL.Path, "/paths-info/"): + fmt.Fprintf(w, `[{"type":"file","path":"model.bin","oid":"file-oid","size":%d}]`, len(payload)) + case strings.Contains(r.URL.Path, "/resolve/"): + number := requests.Add(1) + if !strings.Contains(r.URL.Path, "/fixed/") { + t.Error("download changed commit") + } + start, end := 0, len(payload)-1 + if value := r.Header.Get("Range"); value != "" { + fmt.Sscanf(value, "bytes=%d-%d", &start, &end) + w.Header().Set("Content-Range", fmt.Sprintf("bytes %d-%d/%d", start, end, len(payload))) + } + w.Header().Set("Content-Length", fmt.Sprint(end-start+1)) + if start > 0 { + w.WriteHeader(206) + } + if number == 1 { + io.WriteString(w, payload[:4096]) + w.(http.Flusher).Flush() + <-r.Context().Done() + close(sourceCanceled) + return + } + resumedOffset.Store(int64(start)) + io.WriteString(w, payload[start:end+1]) + default: + w.WriteHeader(404) + } + })) + defer server.Close() + msConfig(t, server.URL) + config.SysConfig.Server.HfScheme = "http" + config.SysConfig.Server.HfNetLoc = strings.TrimPrefix(server.URL, "http://") + config.SysConfig.Download.RemoteFileRangeSize = 16384 + base := data.NewBaseData() + locks := dao.NewLockDao(base) + downloads := dao.NewDownloaderDao(nil) + files := dao.NewFileDao(downloads, base, locks) + svc := NewCacheJobService(files, dao.NewMetaDao(files, locks, base), downloads, nil) + defer svc.cachePool.Close() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + caller := func() echo.Context { + return echo.New().NewContext(httptest.NewRequest("POST", "/", nil).WithContext(app.NewContext(ctx, app.New(app.Context(ctx)))), httptest.NewRecorder()) + } + result, err := svc.CreateCacheJobResult(caller(), &query.CreateCacheJobReq{Type: consts.CacheTypePreheat, Namespace: "huggingface", Datatype: "models", Repo: "owner/model"}) + if err != nil { + t.Fatal(err) + } + deadline := time.Now().Add(5 * time.Second) + for { + status, err := svc.CacheJobStatus(result.ID) + if err != nil { + t.Fatal(err) + } + if status.CachedBytes >= 4096 { + break + } + if time.Now().After(deadline) { + t.Fatalf("no progress: %+v", status) + } + time.Sleep(time.Millisecond) + } + if err := svc.StopCacheJob(&query.JobStatusReq{Id: result.ID}); err != nil { + t.Fatal(err) + } + waitCacheJob(t, svc, result.ID, "paused") + select { + case <-sourceCanceled: + case <-time.After(time.Second): + t.Fatal("paused before source canceled") + } + paused, _ := svc.CacheJobStatus(result.ID) + if paused.Error != "" || requests.Load() != 1 { + t.Fatalf("paused=%+v requests=%d", paused, requests.Load()) + } + // Pool cleanup follows completion; wait for ownership release before resuming. + for { + if _, ok := svc.cachePool.GetTask(int(result.ID)); !ok { + break + } + if time.Now().After(deadline) { + t.Fatal("pool retained stopped worker") + } + time.Sleep(time.Millisecond) + } + if joinScheduler { + if err := config.SysConfig.SaveRegistration(config.Registration{Enabled: true, NodeID: "node-a", Address: "scheduler:19091", Host: "speed", Port: 8090}); err != nil { + t.Fatal(err) + } + } + if err := svc.ResumeCacheJob(caller(), &query.ResumeCacheJobReq{Id: result.ID}); err != nil { + t.Fatal(err) + } + waitCacheJob(t, svc, result.ID, "complete") + final, _ := svc.CacheJobStatus(result.ID) + if !final.Local || final.InstanceID != "" { + t.Fatalf("resume changed task authority: %+v", final) + } + if final.Generation != paused.Generation+1 || final.SnapshotVersion <= paused.SnapshotVersion || resumedOffset.Load() <= 0 { + t.Fatalf("resume did not reuse blocks: %+v offset=%d", final, resumedOffset.Load()) + } + reader := hfprojection.Reader{Root: config.SysConfig.Repos()} + f, contents, err := reader.OpenFile("models", "owner/model", "fixed", "model.bin") + if err != nil { + t.Fatal(err) + } + defer f.Close() + body, err := io.ReadAll(contents) + if err != nil { + t.Fatal(err) + } + if sha256.Sum256(body) != sha256.Sum256([]byte(payload)) { + t.Fatal("resumed file differs from source") + } +} diff --git a/internal/service/cache_job_pause_test.go b/internal/service/cache_job_pause_test.go new file mode 100644 index 0000000..a4f1bbe --- /dev/null +++ b/internal/service/cache_job_pause_test.go @@ -0,0 +1,112 @@ +package service + +import ( + "context" + "dingospeed/internal/dao" + "dingospeed/internal/model/query" + "dingospeed/internal/service/task" + "dingospeed/pkg/repository" + "errors" + "testing" +) + +func TestPauseWaitsForExecutionAndCancelUpgradesIntent(t *testing.T) { + msConfig(t, "http://unused.invalid") + svc := NewCacheJobService(nil, nil, nil, nil) + defer svc.cachePool.Close() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + key := repository.RepoKey{Namespace: repository.HuggingFace, RepoType: "models", Repo: "owner/model"} + job := &task.PreheatCacheTask{CacheTask: task.CacheTask{TaskNo: 71, Ctx: ctx, CancelFunc: cancel, LocalJob: true}, Sha: &dao.CommitHfSha{Sha: "fixed"}, UsedStorage: 8} + if err := trackCacheJob(job, key, context.Background()); err != nil { + t.Fatal(err) + } + if err := svc.StopCacheJob(&query.JobStatusReq{Id: 71}); err != nil { + t.Fatal(err) + } + first, err := svc.CacheJobStatus(71) + if err != nil || first.State != "pausing" || ctx.Err() == nil { + t.Fatalf("%+v %v", first, err) + } + if err := svc.ResumeCacheJob(nil, &query.ResumeCacheJobReq{Id: 71}); err == nil { + t.Fatal("resume accepted before the execution exited") + } + if err := svc.StopCacheJob(&query.JobStatusReq{Id: 71, Intent: "cancel"}); err != nil { + t.Fatal(err) + } + second, _ := svc.CacheJobStatus(71) + if second.State != "canceling" || second.SnapshotVersion <= first.SnapshotVersion { + t.Fatalf("%+v", second) + } + if err := svc.StopCacheJob(&query.JobStatusReq{Id: 71, Intent: "pause"}); err != nil { + t.Fatal(err) + } + second, _ = svc.CacheJobStatus(71) + if second.State != "canceling" { + t.Fatal("cancel downgraded") + } + job.OnFinish(context.Canceled) + done, err := svc.CacheJobStatus(71) + if err != nil || done.State != "canceled" || done.Error != "" { + t.Fatalf("%+v %v", done, err) + } + again, _ := svc.CacheJobStatus(71) + if done != again { + t.Fatal("same version returned a different snapshot") + } + if err := svc.ResumeCacheJob(nil, &query.ResumeCacheJobReq{Id: 71}); err == nil { + t.Fatal("canceled job resumed") + } +} + +func TestRestartPersistsMonotonicSnapshots(t *testing.T) { + msConfig(t, "http://unused.invalid") + svc := NewCacheJobService(nil, nil, nil, nil) + defer svc.cachePool.Close() + for i, state := range []string{"running", "pausing", "resuming", "canceling", "stopped"} { + original := CacheJobStatus{ID: int64(90 + i), State: state, Generation: 3, StateVersion: 10, SnapshotVersion: 20} + if err := writeCacheJobStatus(original); err != nil { + t.Fatal(err) + } + got, err := svc.CacheJobStatus(original.ID) + if err != nil { + t.Fatal(err) + } + want := "interrupted" + if state == "canceling" { + want = "canceled" + } + if state == "stopped" { + want = "paused" + } + if got.State != want || got.Generation != 3 || got.StateVersion != 11 || got.SnapshotVersion != 21 { + t.Fatalf("%+v", got) + } + again, _ := svc.CacheJobStatus(original.ID) + if got != again { + t.Fatal("restart reconciliation not durable") + } + } +} + +func TestResumeConflictsWithOtherActiveJob(t *testing.T) { + msConfig(t, "http://unused.invalid") + svc := NewCacheJobService(nil, nil, nil, nil) + defer svc.cachePool.Close() + key := repository.RepoKey{Namespace: repository.HuggingFace, RepoType: "models", Repo: "owner/model"} + if err := writeCacheJobStatus(CacheJobStatus{ID: 81, RepoKey: key, State: "paused", Commit: "fixed", Local: true}); err != nil { + t.Fatal(err) + } + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + job := &task.PreheatCacheTask{CacheTask: task.CacheTask{TaskNo: 82, Ctx: ctx, CancelFunc: cancel, LocalJob: true}, Sha: &dao.CommitHfSha{Sha: "fixed"}} + if err := trackCacheJob(job, key, context.Background()); err != nil { + t.Fatal(err) + } + defer job.OnFinish(context.Canceled) + err := svc.ResumeCacheJob(nil, &query.ResumeCacheJobReq{Id: 81}) + var conflict *CacheJobConflict + if !errors.As(err, &conflict) || conflict.ActiveJobID != 82 || conflict.Code != "active_job_conflict" { + t.Fatalf("%v", err) + } +} diff --git a/internal/service/cache_job_resume_identity_test.go b/internal/service/cache_job_resume_identity_test.go new file mode 100644 index 0000000..182699e --- /dev/null +++ b/internal/service/cache_job_resume_identity_test.go @@ -0,0 +1,65 @@ +package service + +import ( + "context" + "testing" + + "dingospeed/internal/dao" + "dingospeed/internal/service/task" + "dingospeed/pkg/config" + "dingospeed/pkg/proto/manager" + "dingospeed/pkg/repository" +) + +func TestResumeRetainsOriginalTaskAuthority(t *testing.T) { + msConfig(t, "http://unused.invalid") + svc := &CacheJobService{} + old := CacheJobStatus{Local: true} + if err := config.SysConfig.SaveRegistration(config.Registration{Enabled: true, NodeID: "node-a", Address: "scheduler:19091", Host: "speed", Port: 8090}); err != nil { + t.Fatal(err) + } + // No scheduler client: old local work must remain independent even offline. + for _, requested := range []string{"", "node-a"} { + local, id, err := svc.resumeCacheJobIdentity(old, requested) + if err != nil || !local || id != "" { + t.Fatalf("local authority lost: %v %q %v", local, id, err) + } + } + if _, _, err := svc.resumeCacheJobIdentity(old, "node-b"); err == nil { + t.Fatal("wrong explicit node accepted") + } + cluster := CacheJobStatus{InstanceID: "node-a"} + if _, _, err := svc.resumeCacheJobIdentity(cluster, ""); err == nil { + t.Fatal("cluster work accepted without scheduler") + } + svc.schedulerDao = &dao.SchedulerDao{Client: manager.NewManagerClient(nil)} + if local, id, err := svc.resumeCacheJobIdentity(cluster, ""); err != nil || local || id != "node-a" { + t.Fatalf("cluster identity changed: %v %q %v", local, id, err) + } + cluster.InstanceID = "node-b" + if _, _, err := svc.resumeCacheJobIdentity(cluster, ""); err == nil { + t.Fatal("foreign cluster job accepted") + } +} + +func TestActiveCacheExclusionAcrossModes(t *testing.T) { + msConfig(t, "http://unused.invalid") + svc := NewCacheJobService(nil, nil, nil, nil) + defer svc.cachePool.Close() + key := repository.RepoKey{Namespace: repository.HuggingFace, RepoType: "models", Repo: "owner/model"} + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + job := &task.PreheatCacheTask{CacheTask: task.CacheTask{TaskNo: 981, Ctx: ctx, CancelFunc: cancel, LocalJob: true}, Sha: &dao.CommitHfSha{Sha: "fixed"}} + if err := trackCacheJob(job, key, context.Background()); err != nil { + t.Fatal(err) + } + got, err := svc.reusableCacheJob(key, job.Sha, false, "node-a") + if err != nil || got == nil || got.ID != 981 { + t.Fatalf("cross-mode writer not excluded: %+v %v", got, err) + } + job.OnFinish(nil) + status, err := svc.CacheJobStatus(981) + if err != nil || status.State != "complete" || !status.Local || status.InstanceID != "" { + t.Fatalf("local completion changed authority: %+v %v", status, err) + } +} diff --git a/internal/service/cache_job_service.go b/internal/service/cache_job_service.go index f58fb33..3255078 100644 --- a/internal/service/cache_job_service.go +++ b/internal/service/cache_job_service.go @@ -2,22 +2,29 @@ package service import ( "context" + "errors" "fmt" + "math" + "os" + "strings" + "sync" + "time" "dingospeed/internal/dao" "dingospeed/internal/model/query" task2 "dingospeed/internal/service/task" "dingospeed/pkg/app" "dingospeed/pkg/common" + "dingospeed/pkg/config" "dingospeed/pkg/consts" "dingospeed/pkg/proto/manager" + "dingospeed/pkg/repository" - "github.com/bytedance/sonic" "github.com/labstack/echo/v4" - "go.uber.org/zap" ) type CacheJobService struct { + createMu sync.Mutex fileDao *dao.FileDao metaDao *dao.MetaDao downloaderDao *dao.DownloaderDao @@ -26,16 +33,52 @@ type CacheJobService struct { } func NewCacheJobService(fileDao *dao.FileDao, metaDao *dao.MetaDao, downloaderDao *dao.DownloaderDao, schedulerDao *dao.SchedulerDao) *CacheJobService { - return &CacheJobService{ + p := &CacheJobService{ fileDao: fileDao, metaDao: metaDao, downloaderDao: downloaderDao, schedulerDao: schedulerDao, cachePool: common.NewPool(30, true), } + registrationCachePool = p.cachePool + return p } func (p *CacheJobService) CreateCacheJob(c echo.Context, jobReq *query.CreateCacheJobReq) (int64, error) { + result, err := p.CreateCacheJobResult(c, jobReq) + return result.ID, err +} + +func (p *CacheJobService) CreateCacheJobResult(c echo.Context, jobReq *query.CreateCacheJobReq) (CacheJobCreateResult, error) { + registrationAdmission.Lock() + defer registrationAdmission.Unlock() + localJob, instanceID, err := p.cacheJobIdentity(jobReq.InstanceId) + if err != nil { + return CacheJobCreateResult{}, err + } + jobReq.InstanceId = instanceID + if jobReq.Namespace == "" { + k, err := legacyJobKey(jobReq.Datatype, jobReq.Org, jobReq.Repo, jobReq.OrgRepo) + if err != nil { + return CacheJobCreateResult{}, err + } + jobReq.Namespace, jobReq.Repo = k.Namespace, k.Repo + } else if jobReq.Org != "" || jobReq.OrgRepo != "" { + return CacheJobCreateResult{}, fmt.Errorf("conflicting job identity") + } + + key := repository.RepoKey{Namespace: jobReq.Namespace, RepoType: jobReq.Datatype, Repo: jobReq.Repo} + if err := key.Validate(); err != nil { + return CacheJobCreateResult{}, err + } + if key.Namespace == repository.HuggingFace || key.Namespace == repository.ModelScope { + if err := repository.Register(config.SysConfig.Repos(), repository.Remote(key)); err != nil { + return CacheJobCreateResult{}, err + } + } else if _, err := repository.Read(config.SysConfig.Repos(), key); err != nil { + return CacheJobCreateResult{}, err + } + appInfo, _ := app.FromContext(c.Request().Context()) ctx, cancelFunc := context.WithCancel(appInfo.Ctx()) var task common.Task @@ -45,44 +88,80 @@ func (p *CacheJobService) CreateCacheJob(c echo.Context, jobReq *query.CreateCac CancelFunc: cancelFunc, SchedulerDao: p.schedulerDao, RunningStatus: consts.RunningStatusJobBreak, + LocalJob: localJob, } req := &manager.CreateCacheJobReq{ Type: jobReq.Type, InstanceId: jobReq.InstanceId, Datatype: jobReq.Datatype, - Org: jobReq.Org, + Org: jobReq.Namespace, Repo: jobReq.Repo, Status: consts.RunningStatusJobIng, } authorization := c.Request().Header.Get("Authorization") if jobReq.Type == consts.CacheTypePreheat { - orgRepo := fmt.Sprintf("%s/%s", jobReq.Org, jobReq.Repo) - metadata, err := p.metaDao.GetMetadata(jobReq.Datatype, orgRepo, "main", "get", authorization) + revision := "main" + if key.Namespace == repository.ModelScope { + revision = "master" + } + sha, err := p.preheatMetadata(ctx, key, revision, authorization) if err != nil { - return 0, err + cancelFunc() + return CacheJobCreateResult{}, err } - var sha dao.CommitHfSha - if err = sonic.Unmarshal(metadata.OriginContent, &sha); err != nil { - zap.S().Errorf("unmarshal content error:%v", err) - return 0, err + // Metadata resolves the requested branch before comparing immutable commits. + // Hold admission until the new job is registered AND submitted to the pool. + p.createMu.Lock() + defer p.createMu.Unlock() + existing, err := p.reusableCacheJob(key, sha, localJob, instanceID) + if err != nil { + cancelFunc() + return CacheJobCreateResult{}, err + } + if existing != nil { + cancelFunc() + disposition := "running" + if existing.State == "complete" { + disposition = "cached" + } + return CacheJobCreateResult{ID: existing.ID, Msg: "success", Disposition: disposition, Commit: sha.Sha}, nil } req.UsedStorage = sha.UsedStorage req.Commit = sha.Sha - cacheJob, err := p.schedulerDao.CreateCacheJob(req) + var jobID int64 + if localJob { + jobID, err = reserveLocalCacheJob() + } else { + var cacheJob *manager.CreateCacheJobResp + cacheJob, err = p.schedulerDao.CreateCacheJob(req) + if err == nil { + jobID = cacheJob.Id + } + } if err != nil { - return 0, err + cancelFunc() + return CacheJobCreateResult{}, err } - cacheTask.TaskNo = int(cacheJob.Id) + cacheTask.TaskNo = int(jobID) task = &task2.PreheatCacheTask{ CacheTask: cacheTask, FileDao: p.fileDao, DownloaderDao: p.downloaderDao, - Sha: &sha, + Sha: sha, Authorization: authorization, UsedStorage: uint64(sha.UsedStorage), } + if err := trackCacheJob(task.(*task2.PreheatCacheTask), key, appInfo.Ctx()); err != nil { + cancelFunc() + return CacheJobCreateResult{}, err + } if err = p.cachePool.SubmitForTimeout(ctx, task); err != nil { - p.schedulerDao.ExecUpdateCacheJobStatus(int(cacheJob.Id), consts.RunningStatusJobWait, jobReq.InstanceId, "", "", consts.TaskMoreErrMsg, 0) + task.(*task2.PreheatCacheTask).OnFinish(err) + if !localJob { + p.schedulerDao.ExecUpdateCacheJobStatus(int(jobID), consts.RunningStatusJobWait, jobReq.InstanceId, "", "", consts.TaskMoreErrMsg, 0) + } + cancelFunc() + return CacheJobCreateResult{}, err } } else if jobReq.Type == consts.CacheTypeMount { cacheTask.TaskNo = int(jobReq.RepositoryId) @@ -95,68 +174,156 @@ func (p *CacheJobService) CreateCacheJob(c echo.Context, jobReq *query.CreateCac } } else { defer cancelFunc() - return 0, fmt.Errorf("cache job type is err,%d", jobReq.Type) + return CacheJobCreateResult{}, fmt.Errorf("cache job type is err,%d", jobReq.Type) } - return int64(cacheTask.TaskNo), nil + return CacheJobCreateResult{ID: int64(cacheTask.TaskNo), Msg: "success", Disposition: "created", Commit: req.Commit}, nil +} + +type CacheJobConflict struct { + Code string + ActiveJobID int64 } -func (p *CacheJobService) StopCacheJob(jobStatusReq *query.JobStatusReq) error { - if task, ok := p.cachePool.GetTask(int(jobStatusReq.Id)); ok { - if t, ok := task.(*task2.PreheatCacheTask); ok { - t.RunningStatus = consts.RunningStatusJobStop - } else if t, ok := task.(*task2.MountCacheTask); ok { - t.RunningStatus = consts.RunningStatusJobStop +func (e *CacheJobConflict) Error() string { return e.Code } + +func (p *CacheJobService) StopCacheJob(req *query.JobStatusReq) error { + if req.Intent == "" { + req.Intent = "pause" + } + if req.Intent != "pause" && req.Intent != "cancel" { + return fmt.Errorf("invalid stop intent") + } + cacheStateMu.Lock() + defer cacheStateMu.Unlock() + status, err := cacheJobSnapshotLocked(req.Id) + if errors.Is(err, os.ErrNotExist) { + if existing, ok := p.cachePool.GetTask(int(req.Id)); ok { + if mount, ok := existing.(*task2.MountCacheTask); ok { + mount.RequestStop() + return nil + } + } + } + if err != nil { + return err + } + before := status + if status.State == "complete" || status.State == "canceled" { + return nil + } + execution := cacheExecutions[cacheJobStatusPath(req.Id)] + if execution == nil { + if req.Intent == "cancel" { + status.State = "canceled" + } else if status.State == "paused" || status.State == "interrupted" { + status.State = "paused" + } else { + return &CacheJobConflict{Code: "invalid_task_state"} } - cancelFun := task.GetCancelFun() - cancelFun() } else { - p.schedulerDao.ExecUpdateCacheJobStatus(int(jobStatusReq.Id), consts.RunningStatusJobStop, jobStatusReq.InstanceId, "", "", "speed未注册该任务，下载已中断。", 0) + if req.Intent == "cancel" || status.State == "canceling" { + status.State = "canceling" + } else { + status.State = "pausing" + } + } + status.Error, status.ErrorCode = "", "" + if status != before { + if err := commitCacheSnapshot(before, &status); err != nil { + return err + } + } + if execution != nil { + if execution.stopStarted.IsZero() { + execution.stopStarted = time.Now() + } + execution.task.CancelFunc() } return nil } -func (p *CacheJobService) ResumeCacheJob(c echo.Context, resumeCacheJobReq *query.ResumeCacheJobReq) error { +func (p *CacheJobService) ResumeCacheJob(c echo.Context, req *query.ResumeCacheJobReq) error { + registrationAdmission.Lock() + defer registrationAdmission.Unlock() + // Create and resume share admission; no second job may write the same commit. + p.createMu.Lock() + defer p.createMu.Unlock() + previous, err := p.CacheJobStatus(req.Id) + if err != nil { + return err + } + if cacheActive(previous.State) { + if previous.State == "pausing" || previous.State == "canceling" { + return &CacheJobConflict{Code: "task_stopping"} + } + return nil + } + if previous.State != "paused" && previous.State != "interrupted" && previous.State != "failed" { + return &CacheJobConflict{Code: "invalid_task_state"} + } + if previous.Commit == "" { + return &CacheJobConflict{Code: "missing_pinned_commit"} + } + local, instance, err := p.resumeCacheJobIdentity(previous, req.InstanceId) + if err != nil { + return err + } + if _, active := p.cachePool.GetTask(int(req.Id)); active { + return &CacheJobConflict{Code: "task_stopping"} + } + existing, err := p.reusableCacheJob(previous.RepoKey, &dao.CommitHfSha{Sha: previous.Commit}, local, instance) + if err != nil { + return err + } + if existing != nil && existing.ID != req.Id && cacheActive(existing.State) { + return &CacheJobConflict{Code: "active_job_conflict", ActiveJobID: existing.ID} + } appInfo, _ := app.FromContext(c.Request().Context()) - ctx, cancelFunc := context.WithCancel(appInfo.Ctx()) - var task common.Task - cacheTask := task2.CacheTask{ - TaskNo: int(resumeCacheJobReq.Id), - Ctx: ctx, - Job: &query.CreateCacheJobReq{ - InstanceId: resumeCacheJobReq.InstanceId, Type: resumeCacheJobReq.Type, - Org: resumeCacheJobReq.Org, Repo: resumeCacheJobReq.Repo, - Datatype: resumeCacheJobReq.Datatype, - }, - CancelFunc: cancelFunc, - SchedulerDao: p.schedulerDao, - } - if resumeCacheJobReq.Type == consts.CacheTypePreheat { - orgRepo := fmt.Sprintf("%s/%s", resumeCacheJobReq.Org, resumeCacheJobReq.Repo) - authorization := c.Request().Header.Get("Authorization") - metadata, err := p.metaDao.GetMetadata(resumeCacheJobReq.Datatype, orgRepo, "main", "get", authorization) + ctx, cancel := context.WithCancel(appInfo.Ctx()) + preheat := &task2.PreheatCacheTask{ + CacheTask: task2.CacheTask{TaskNo: int(req.Id), Ctx: ctx, CancelFunc: cancel, LocalJob: local, SchedulerDao: p.schedulerDao, + Job: &query.CreateCacheJobReq{InstanceId: instance, Type: consts.CacheTypePreheat, Namespace: previous.Namespace, Repo: previous.Repo, Datatype: previous.RepoType}}, + Sha: &dao.CommitHfSha{Sha: previous.Commit}, UsedStorage: previous.TotalBytes, FileDao: p.fileDao, DownloaderDao: p.downloaderDao, + Authorization: c.Request().Header.Get("Authorization"), InitialState: "resuming", + } + if err := trackCacheJob(preheat, previous.RepoKey, appInfo.Ctx()); err != nil { + cancel() + return err + } + go func() { + sha, err := p.preheatMetadata(ctx, previous.RepoKey, previous.Commit, preheat.Authorization) + if err == nil && sha.Sha != previous.Commit { + err = fmt.Errorf("upstream changed pinned commit") + } + if err == nil { + err = ctx.Err() + } if err != nil { - return err + preheat.OnFinish(err) + cancel() + return } - var sha dao.CommitHfSha - if err = sonic.Unmarshal(metadata.OriginContent, &sha); err != nil { - zap.S().Errorf("unmarshal content error:%v", err) - return err + preheat.Sha = sha + preheat.UsedStorage = uint64(sha.UsedStorage) + cacheStateMu.Lock() + current, err := readCacheJobStatus(req.Id) + if err == nil && current.State == "resuming" { + next := current + next.State = "running" + next.TotalBytes = preheat.UsedStorage + err = commitCacheSnapshot(current, &next) + } else if err == nil { + err = context.Canceled } - // 将状态重置为进行中 - p.schedulerDao.ExecUpdateCacheJobStatus(int(resumeCacheJobReq.Id), consts.RunningStatusJobIng, - resumeCacheJobReq.InstanceId, resumeCacheJobReq.Org, resumeCacheJobReq.Repo, "", 0) - task = &task2.PreheatCacheTask{ - CacheTask: cacheTask, - FileDao: p.fileDao, - DownloaderDao: p.downloaderDao, - Sha: &sha, - Authorization: authorization, - UsedStorage: uint64(resumeCacheJobReq.UsedStorage), + cacheStateMu.Unlock() + if err == nil { + err = p.cachePool.SubmitForTimeout(ctx, preheat) } - if err = p.cachePool.SubmitForTimeout(ctx, task); err != nil { - p.schedulerDao.ExecUpdateCacheJobStatus(int(resumeCacheJobReq.Id), consts.RunningStatusJobWait, resumeCacheJobReq.InstanceId, "", "", consts.TaskMoreErrMsg, 0) + if err != nil { + preheat.OnFinish(err) + cancel() } - } + }() return nil } @@ -165,13 +332,103 @@ func (p *CacheJobService) RealtimeCacheJob(realtimeReq *query.RealtimeReq) []*qu for _, jobId := range realtimeReq.CacheJobIds { if task, ok := p.cachePool.GetTask(int(jobId)); ok { if pTask, ok := task.(*task2.PreheatCacheTask); ok { + speed, progress := pTask.RealtimeProgress() ret = append(ret, &query.RealtimeResp{ CacheJobId: jobId, - StockSpeed: pTask.StockSpeed, - StockProcess: pTask.StockProcess, + StockSpeed: speed, + StockProcess: progress, }) } } } return ret } + +func (p *CacheJobService) preheatMetadata(ctx context.Context, key repository.RepoKey, revision, authorization string) (*dao.CommitHfSha, error) { + if key.Namespace == repository.HuggingFace { + meta, err := p.metaDao.RefreshPreheatMetadataContext(ctx, key, revision, authorization) + if err != nil { + return nil, err + } + // HF usedStorage may include multiple versions; count the pinned snapshot. + meta.UsedStorage = 0 + prefix := "" + if key.RepoType != "models" { + prefix = "/" + key.RepoType + } + for _, file := range meta.Siblings { + if err := ctx.Err(); err != nil { + return nil, err + } + if err := repository.Relative(file.Rfilename); err != nil { + return nil, err + } + uri := prefix + "/" + repository.EscapeURLPath(key.Repo) + "/resolve/" + repository.EscapeURLPath(meta.Sha) + "/" + repository.EscapeURLPath(file.Rfilename) + info, err := p.fileDao.GetPathsInfoContext(ctx, uri, key.RepoType, key.ID(), meta.Sha, authorization, file.Rfilename) + if err != nil { + return nil, err + } + if info == nil || info.Size < 0 || info.Size > math.MaxInt64-meta.UsedStorage { + return nil, fmt.Errorf("invalid snapshot file size") + } + meta.UsedStorage += info.Size + } + return meta, nil + } + return task2.ReadRepositoryMetadata(ctx, key, revision, authorization) +} + +func legacyJobKey(typ, org, repo, id string) (repository.RepoKey, error) { + if id != "" { + if org != "" || repo != "" { + if strings.TrimPrefix(org+"/"+repo, "/") != id { + return repository.RepoKey{}, fmt.Errorf("conflicting orgRepo") + } + } else { + parts := strings.SplitN(id, "/", 2) + if len(parts) == 1 { + repo = id + } else { + org, repo = parts[0], parts[1] + } + } + } + return repository.FromStorage(typ, org, repo) +} + +// Task authority follows durable registration even while a cluster is degraded. +func (p *CacheJobService) resumeCacheJobIdentity(previous CacheJobStatus, requested string) (bool, string, error) { + if previous.Local { + // The durable record belongs to this data root. Joining a scheduler + // does not transfer ownership of an existing standalone job. + if requested != "" && requested != config.SysConfig.Registration().NodeID { + return false, "", &CacheJobConflict{Code: "task_node_mismatch"} + } + return true, "", nil + } + local, instance, err := p.cacheJobIdentity(requested) + if err != nil { + return false, "", err + } + if local || previous.InstanceID != instance { + return false, "", &CacheJobConflict{Code: "task_node_mismatch"} + } + return false, instance, nil +} + +func (p *CacheJobService) cacheJobIdentity(requested string) (bool, string, error) { + registration := config.SysConfig.Registration() + if state := config.SysConfig.RegistrationState(); state.State == "error" { + return false, "", fmt.Errorf("invalid scheduler registration: %s", state.Error) + } + if !registration.Enabled { + return true, "", nil + } + if requested != "" && requested != registration.NodeID { + return false, "", fmt.Errorf("cache job node identity mismatch") + } + if registration.NodeID == "" || p.schedulerDao == nil || p.schedulerDao.Client == nil { + return false, "", fmt.Errorf("scheduler unavailable") + } + return false, registration.NodeID, nil +} diff --git a/internal/service/cache_job_standalone_test.go b/internal/service/cache_job_standalone_test.go new file mode 100644 index 0000000..381ffbd --- /dev/null +++ b/internal/service/cache_job_standalone_test.go @@ -0,0 +1,137 @@ +package service + +import ( + "context" + "fmt" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "dingospeed/internal/dao" + "dingospeed/internal/data" + "dingospeed/internal/model/query" + "dingospeed/pkg/app" + "dingospeed/pkg/common" + "dingospeed/pkg/config" + "dingospeed/pkg/consts" + "dingospeed/pkg/proto/manager" + "github.com/labstack/echo/v4" +) + +func TestStandaloneCacheJobLifecycle(t *testing.T) { + for _, fail := range []bool{false, true} { + t.Run(fmt.Sprint("downloadFailure=", fail), func(t *testing.T) { + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch { + case strings.Contains(r.URL.Path, "/revision/"): + fmt.Fprint(w, `{"sha":"commit-one","usedStorage":24,"siblings":[{"rfilename":"model.bin"}]}`) + case strings.Contains(r.URL.Path, "/paths-info/"): + fmt.Fprint(w, `[{"type":"file","path":"model.bin","oid":"file-oid","size":8}]`) + case strings.Contains(r.URL.Path, "/resolve/"): + if fail { + w.WriteHeader(403) + return + } + w.Header().Set("Content-Length", "8") + fmt.Fprint(w, "abcdefgh") + default: + t.Errorf("unexpected upstream %s", r.URL) + w.WriteHeader(404) + } + })) + defer upstream.Close() + msConfig(t, upstream.URL) + config.SysConfig.Server.HfScheme = "http" + config.SysConfig.Server.HfNetLoc = strings.TrimPrefix(upstream.URL, "http://") + config.SysConfig.Scheduler.Mode = "standalone" + base := data.NewBaseData() + locks := dao.NewLockDao(base) + downloads := dao.NewDownloaderDao(nil) + files := dao.NewFileDao(downloads, base, locks) + svc := NewCacheJobService(files, dao.NewMetaDao(files, locks, base), downloads, nil) + defer svc.cachePool.Close() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + req := httptest.NewRequest("POST", "/api/cacheJob/create", nil).WithContext(app.NewContext(ctx, app.New(app.Context(ctx)))) + c := echo.New().NewContext(req, httptest.NewRecorder()) + id, err := svc.CreateCacheJob(c, &query.CreateCacheJobReq{Type: consts.CacheTypePreheat, Namespace: "huggingface", Datatype: "models", Repo: "owner/model"}) + if err != nil { + t.Fatal(err) + } + deadline := time.Now().Add(5 * time.Second) + var status CacheJobStatus + for time.Now().Before(deadline) { + status, err = svc.CacheJobStatus(id) + if err != nil { + t.Fatal(err) + } + if status.State == "complete" || status.State == "failed" { + break + } + time.Sleep(10 * time.Millisecond) + } + want := "complete" + if fail { + want = "failed" + } + if status.State != want || !status.Local || status.TotalBytes != 8 { + t.Fatalf("status=%+v", status) + } + if !fail && status.CachedBytes != 8 { + t.Fatalf("missing bytes: %+v", status) + } + if fail && status.Error == "" { + t.Fatal("missing failure reason") + } + restarted := &CacheJobService{cachePool: common.NewPool(1, true)} + defer restarted.cachePool.Close() + reloaded, err := restarted.CacheJobStatus(id) + if err != nil || reloaded.State != want { + t.Fatalf("restart lost result: %+v %v", reloaded, err) + } + }) + } +} + +func TestCacheJobIdentityClusterAndStandalone(t *testing.T) { + msConfig(t, "http://unused.invalid") + svc := &CacheJobService{} + if local, id, err := svc.cacheJobIdentity(""); !local || id != "" || err != nil { + t.Fatalf("standalone: %v %q %v", local, id, err) + } + config.SysConfig.Scheduler.Mode = consts.SchedulerModeCluster + config.SysConfig.Scheduler.OriginMode = consts.SchedulerModeCluster + config.SysConfig.Scheduler.Discovery.InstanceId = "speed-1" + if err := config.SysConfig.SaveRegistration(config.Registration{Enabled: true, NodeID: "speed-1", Address: "scheduler:19091", Host: "speed", Port: 8090}); err != nil { + t.Fatal(err) + } + if _, _, err := svc.cacheJobIdentity("speed-2"); err == nil { + t.Fatal("mismatched binding accepted") + } + if _, _, err := svc.cacheJobIdentity(""); err == nil { + t.Fatal("missing scheduler accepted") + } + // A degraded cluster must not silently create a standalone job. + config.SysConfig.Scheduler.Mode = "standalone" + if _, _, err := svc.cacheJobIdentity(""); err == nil { + t.Fatal("degraded cluster switched task authority") + } + svc.schedulerDao = &dao.SchedulerDao{Client: manager.NewManagerClient(nil)} + if local, id, err := svc.cacheJobIdentity(""); local || id != "speed-1" || err != nil { + t.Fatalf("cluster: %v %q %v", local, id, err) + } +} + +func TestLocalCacheIDsDoNotOverwriteAfterRestart(t *testing.T) { + msConfig(t, "http://unused.invalid") + seen := map[int64]bool{} + for i := 0; i < 20; i++ { + id, err := reserveLocalCacheJob() + if err != nil || seen[id] { + t.Fatalf("ID collision: %d %v", id, err) + } + seen[id] = true + } +} diff --git a/internal/service/cache_job_status.go b/internal/service/cache_job_status.go new file mode 100644 index 0000000..ceeb2db --- /dev/null +++ b/internal/service/cache_job_status.go @@ -0,0 +1,299 @@ +package service + +import ( + "context" + "dingospeed/pkg/common" + "dingospeed/pkg/consts" + "encoding/json" + "errors" + "fmt" + "os" + "path/filepath" + "sync" + "time" + + "dingospeed/internal/service/task" + "dingospeed/pkg/config" + "dingospeed/pkg/repository" + "dingospeed/pkg/util" + "go.uber.org/zap" +) + +type CacheJobStatus struct { + ID int64 `json:"id"` + ProtocolVersion int `json:"protocolVersion"` + Generation uint64 `json:"generation"` + StateVersion uint64 `json:"stateVersion"` + SnapshotVersion uint64 `json:"snapshotVersion"` + ActiveSeconds float64 `json:"activeSeconds"` + ErrorCode string `json:"errorCode,omitempty"` + repository.RepoKey + Commit string `json:"commit"` + State string `json:"state"` + CachedBytes uint64 `json:"cachedBytes"` + TotalBytes uint64 `json:"totalBytes"` + Error string `json:"error,omitempty"` + UpdatedAt time.Time `json:"updatedAt"` + Local bool `json:"local,omitempty"` + InstanceID string `json:"instanceId,omitempty"` +} + +// Reserve an ID on disk so concurrent submissions and restarts cannot replace +// another job. Millisecond IDs fit exactly in the console's JavaScript number. +func reserveLocalCacheJob() (int64, error) { + for id := time.Now().UnixMilli(); ; id++ { + path := cacheJobStatusPath(id) + if err := util.MakeDirs(path); err != nil { + return 0, err + } + f, err := os.OpenFile(path, os.O_CREATE|os.O_EXCL|os.O_WRONLY, 0600) + if errors.Is(err, os.ErrExist) { + continue + } + if err != nil { + return 0, err + } + err = json.NewEncoder(f).Encode(CacheJobStatus{ID: id, Local: true, State: "interrupted", UpdatedAt: time.Now().UTC()}) + closeErr := f.Close() + if err != nil { + return 0, err + } + if closeErr != nil { + return 0, closeErr + } + return id, nil + } +} + +func cacheJobStatusPath(id int64) string { + return filepath.Join(config.SysConfig.Repos(), "cache-jobs", fmt.Sprintf("%d.json", id)) +} +func writeCacheJobStatus(status CacheJobStatus) error { + if status.UpdatedAt.IsZero() { + status.UpdatedAt = time.Now().UTC() + } + path := cacheJobStatusPath(status.ID) + if err := util.MakeDirs(path); err != nil { + return err + } + return util.WriteDataToFileAtomic(path, status) +} + +// Serializes durable snapshots and the corresponding execution ownership. +// Keys include the data root: separate local nodes never share execution identity. +var cacheStateMu sync.Mutex +var cacheExecutions = map[string]*cacheExecution{} + +type cacheExecution struct { + trace *common.CacheTransferTrace + state string + task *task.PreheatCacheTask + generation uint64 + finished bool + startedWork bool + result error + serviceCtx context.Context + started time.Time + baseSeconds float64 + stopStarted time.Time + finishedAt time.Time +} + +func readCacheJobStatus(id int64) (CacheJobStatus, error) { + var status CacheJobStatus + if id <= 0 { + return status, fmt.Errorf("invalid job id") + } + b, err := os.ReadFile(cacheJobStatusPath(id)) + if err != nil { + return status, err + } + err = json.Unmarshal(b, &status) + return status, err +} + +func commitCacheSnapshot(before CacheJobStatus, next *CacheJobStatus) error { + next.ProtocolVersion = 1 + next.StateVersion = before.StateVersion + if next.State != before.State || next.StateVersion == 0 { + next.StateVersion++ + } + next.SnapshotVersion = before.SnapshotVersion + 1 + next.UpdatedAt = time.Now().UTC() + err := writeCacheJobStatus(*next) + if err == nil { + if execution := cacheExecutions[cacheJobStatusPath(next.ID)]; execution != nil { + execution.state = next.State + } + } + if err != nil { + cachePersistenceFailures.Inc() + } else if before.State != next.State { + cacheTransitions.WithLabelValues(next.State).Inc() + } + return err +} + +func cacheActive(state string) bool { + switch state { + case "waiting", "running", "pausing", "resuming", "canceling": + return true + } + return false +} + +func (p *CacheJobService) CacheJobStatus(id int64) (CacheJobStatus, error) { + cacheStateMu.Lock() + defer cacheStateMu.Unlock() + return cacheJobSnapshotLocked(id) +} + +func cacheJobSnapshotLocked(id int64) (CacheJobStatus, error) { + status, err := readCacheJobStatus(id) + if err != nil { + return status, err + } + before := status + if execution := cacheExecutions[cacheJobStatusPath(id)]; execution != nil { + if execution.finished { + if err := finishCacheExecutionLocked(status, execution); err != nil { + return status, err + } + return readCacheJobStatus(id) + } + status.CachedBytes = execution.task.CachedBytes() + if status.CachedBytes < before.CachedBytes { + status.CachedBytes = before.CachedBytes + } + status.ActiveSeconds = execution.baseSeconds + time.Since(execution.started).Seconds() + } else if cacheActive(status.State) { + if status.State == "canceling" { + status.State = "canceled" + status.Error = "" + } else { + status.State = "interrupted" + status.Error = "Download interrupted because the task is no longer running; resume manually" + } + } else if status.State == "stopped" { + status.State = "paused" + status.Error = "" + } + if status != before { + if err = commitCacheSnapshot(before, &status); err != nil { + return before, err + } + } + return status, nil +} + +func finishCacheExecutionLocked(status CacheJobStatus, execution *cacheExecution) error { + if status.Generation != execution.generation { + return fmt.Errorf("stale cache execution") + } + if execution.trace.Active.Load() != 0 { + return fmt.Errorf("cache source requests still active") + } + execution.trace.Stop() + before := status + status.CachedBytes = execution.task.CachedBytes() + if status.CachedBytes < before.CachedBytes { + status.CachedBytes = before.CachedBytes + } + status.ActiveSeconds = execution.baseSeconds + execution.finishedAt.Sub(execution.started).Seconds() + status.Error, status.ErrorCode = "", "" + switch status.State { + case "pausing": + status.State = "paused" + case "canceling": + status.State = "canceled" + default: + status.State = "complete" + if execution.result != nil { + status.State = "failed" + status.Error = execution.result.Error() + if execution.task.InitialState == "resuming" && !execution.startedWork { + status.State = "paused" + } + if errors.Is(execution.result, context.Canceled) && execution.serviceCtx.Err() != nil { + status.State = "interrupted" + status.Error = "Download interrupted by service shutdown; resume manually" + } + } + } + // A genuine storage failure must remain visible even during an intentional stop. + if execution.result != nil && !errors.Is(execution.result, context.Canceled) && (before.State == "pausing" || before.State == "canceling") { + status.State = "failed" + status.Error = execution.result.Error() + status.ErrorCode = "download_failed" + } + if err := commitCacheSnapshot(before, &status); err != nil { + return err + } + if !execution.task.LocalJob && execution.task.SchedulerDao != nil && execution.task.Job != nil { + code := int32(consts.RunningStatusJobBreak) + if status.State == "complete" { + code = consts.RunningStatusJobComplete + } + if status.State == "paused" || status.State == "canceled" { + code = consts.RunningStatusJobStop + } + job := execution.task.Job + execution.task.SchedulerDao.ExecUpdateCacheJobStatus(execution.task.TaskNo, code, job.InstanceId, job.Namespace, job.Repo, status.Error, execution.task.StockProcess) + } + if !execution.stopStarted.IsZero() { + cacheStopDuration.WithLabelValues(before.State).Observe(execution.finishedAt.Sub(execution.stopStarted).Seconds()) + } + zap.S().Infow("Finished cache execution", "job", status.ID, "generation", execution.generation, "state", status.State, "sourceRequests", execution.trace.Requests.Load(), "sourceCancellations", execution.trace.Canceled.Load(), "cacheWrittenBytes", execution.trace.WrittenBytes.Load()) + delete(cacheExecutions, cacheJobStatusPath(status.ID)) + return nil +} + +func trackCacheJob(preheat *task.PreheatCacheTask, key repository.RepoKey, serviceCtx context.Context) error { + cacheStateMu.Lock() + defer cacheStateMu.Unlock() + previous, err := readCacheJobStatus(int64(preheat.TaskNo)) + if err != nil && !errors.Is(err, os.ErrNotExist) { + return err + } + path := cacheJobStatusPath(int64(preheat.TaskNo)) + if cacheExecutions[path] != nil { + return fmt.Errorf("cache execution already active") + } + status := CacheJobStatus{ID: int64(preheat.TaskNo), RepoKey: key, Commit: preheat.Sha.Sha, State: "running", TotalBytes: preheat.UsedStorage, Local: preheat.LocalJob, + Generation: previous.Generation + 1, CachedBytes: previous.CachedBytes, ActiveSeconds: previous.ActiveSeconds} + if preheat.InitialState != "" { + status.State = preheat.InitialState + } + if preheat.Job != nil { + status.InstanceID = preheat.Job.InstanceId + } + if err := commitCacheSnapshot(previous, &status); err != nil { + return err + } + execution := &cacheExecution{task: preheat, generation: status.Generation, serviceCtx: serviceCtx, started: time.Now(), baseSeconds: previous.ActiveSeconds} + execution.trace = &common.CacheTransferTrace{} + execution.state = status.State + if preheat.Ctx != nil { + preheat.Ctx = common.WithCacheTrace(preheat.Ctx, execution.trace) + } + cacheExecutions[path] = execution + preheat.OnStart = func() { cacheStateMu.Lock(); defer cacheStateMu.Unlock(); execution.startedWork = true } + preheat.OnFinish = func(taskErr error) { + cacheStateMu.Lock() + defer cacheStateMu.Unlock() + if cacheExecutions[path] != execution { + zap.S().Warnw("Ignored stale cache task completion", "job", status.ID, "generation", execution.generation) + return + } + execution.finished, execution.result = true, taskErr + execution.finishedAt = time.Now() + current, err := readCacheJobStatus(status.ID) + if err == nil { + err = finishCacheExecutionLocked(current, execution) + } + if err != nil { + zap.S().Errorw("Could not persist cache task result", "job", status.ID, "error", err) + } + } + return nil +} diff --git a/internal/service/cache_job_status_test.go b/internal/service/cache_job_status_test.go new file mode 100644 index 0000000..219efa9 --- /dev/null +++ b/internal/service/cache_job_status_test.go @@ -0,0 +1,121 @@ +package service + +import ( + "context" + "errors" + "fmt" + "strings" + "testing" + + "dingospeed/internal/dao" + "dingospeed/internal/service/task" + "dingospeed/pkg/common" + "dingospeed/pkg/repository" +) + +func TestCacheJobShutdownRequiresServiceCancellation(t *testing.T) { + msConfig(t, "http://unused.invalid") + for _, tc := range []struct { + name string + stopService bool + err error + want string + }{ + {"shutdown", true, fmt.Errorf("receive: %w", context.Canceled), "interrupted"}, + {"task cancellation", false, context.Canceled, "failed"}, + {"upstream failure", false, errors.New("upstream disconnected"), "failed"}, + {"failure during shutdown", true, errors.New("disk full"), "failed"}, + {"completed before shutdown", true, nil, "complete"}, + } { + t.Run(tc.name, func(t *testing.T) { + serviceCtx, stop := context.WithCancel(context.Background()) + defer stop() + ctx, cancel := context.WithCancel(serviceCtx) + defer cancel() + key := repository.RepoKey{Namespace: repository.HuggingFace, RepoType: "models", Repo: "owner/model"} + job := &task.PreheatCacheTask{CacheTask: task.CacheTask{TaskNo: 18, Ctx: ctx, LocalJob: true}, Sha: &dao.CommitHfSha{Sha: "pinned-commit"}, UsedStorage: 100} + if err := trackCacheJob(job, key, serviceCtx); err != nil { + t.Fatal(err) + } + if tc.stopService { + stop() + } else { + cancel() + } + job.OnFinish(tc.err) + // A new service can inspect the persisted result without resubmitting it. + service := NewCacheJobService(nil, nil, nil, nil) + defer service.cachePool.Close() + for range 2 { + status, err := service.CacheJobStatus(18) + if err != nil || status.State != tc.want || status.Commit != "pinned-commit" || status.TotalBytes != 100 || !status.Local { + t.Fatalf("restarted service: %+v %v", status, err) + } + if tc.want == "interrupted" && !strings.Contains(status.Error, "service shutdown; resume manually") { + t.Fatalf("lost shutdown reason: %+v", status) + } + if tc.want == "failed" && status.Error != tc.err.Error() { + t.Fatalf("lost original failure: %+v", status) + } + if _, running := service.cachePool.GetTask(18); running { + t.Fatal("inspection must not resume downloads") + } + } + }) + } +} + +func TestCacheJobTerminalStatusSurvivesPoolRestart(t *testing.T) { + msConfig(t, "http://unused.invalid") + for _, namespace := range []string{repository.HuggingFace, repository.ModelScope} { + for _, failed := range []bool{false, true} { + key := repository.RepoKey{Namespace: namespace, RepoType: "models", Repo: "owner/model"} + job := &task.PreheatCacheTask{CacheTask: task.CacheTask{TaskNo: 17}, Sha: &dao.CommitHfSha{Sha: "repo-head"}, UsedStorage: 100} + if err := trackCacheJob(job, key, context.Background()); err != nil { + t.Fatal(err) + } + service := &CacheJobService{cachePool: common.NewPool(1, true)} + status, err := service.CacheJobStatus(17) + if err != nil || status.State != "running" { + t.Fatalf("registered task: %+v %v", status, err) + } + want := "complete" + var taskErr error + if failed { + want = "failed" + taskErr = errors.New("upstream unavailable") + } + job.OnFinish(taskErr) + status, err = service.CacheJobStatus(17) + service.cachePool.Close() + if err != nil || status.State != want || status.RepoKey != key || status.Commit != "repo-head" || status.TotalBytes != 100 { + t.Fatalf("terminal task: %+v %v", status, err) + } + if failed && status.Error != taskErr.Error() { + t.Fatalf("lost failure: %+v", status) + } + } + } +} + +func TestOrphanedCacheJobKeepsProgressWithoutResuming(t *testing.T) { + msConfig(t, "http://unused.invalid") + status := CacheJobStatus{ + ID: 19, RepoKey: repository.RepoKey{Namespace: repository.HuggingFace, RepoType: "models", Repo: "owner/model"}, + Commit: "pinned-commit", State: "running", CachedBytes: 50, TotalBytes: 100, Local: true, + } + if err := writeCacheJobStatus(status); err != nil { + t.Fatal(err) + } + service := NewCacheJobService(nil, nil, nil, nil) + defer service.cachePool.Close() + for range 2 { + got, err := service.CacheJobStatus(status.ID) + if err != nil || got.State != "interrupted" || got.CachedBytes != 50 || got.Commit != status.Commit || !strings.Contains(got.Error, "resume manually") { + t.Fatalf("orphaned task: %+v %v", got, err) + } + if _, running := service.cachePool.GetTask(int(status.ID)); running { + t.Fatal("orphaned task was automatically resumed") + } + } +} diff --git a/internal/service/inventory_outbox.go b/internal/service/inventory_outbox.go new file mode 100644 index 0000000..eb631f3 --- /dev/null +++ b/internal/service/inventory_outbox.go @@ -0,0 +1,401 @@ +package service + +import ( + "bytes" + "context" + "dingospeed/internal/dao" + "dingospeed/pkg/config" + "dingospeed/pkg/inventory" + "dingospeed/pkg/repository" + "encoding/json" + "errors" + "fmt" + "go.uber.org/zap" + "io" + "math/rand" + "net/http" + "net/url" + "strings" + "time" +) + +const maxBaselineBuildAttempts = 3 + +// Only local baseline construction consumes the scan budget. Network retries +// reuse a fixed report and must not exhaust it. +type baselineBuildError struct{ error } + +var inventoryConnected = make(chan struct{}, 1) + +func wakeInventoryConnection() { + select { + case inventoryConnected <- struct{}{}: + default: + } +} + +func inventoryPOST(ctx context.Context, path string, in, out any) error { + base := strings.TrimRight(config.SysConfig.Registration().HTTPURL, "/") + if base == "" { + return fmt.Errorf("Scheduler HTTP URL is not configured") + } + b, err := json.Marshal(in) + if err != nil { + return err + } + req, err := http.NewRequestWithContext(ctx, http.MethodPost, base+path, bytes.NewReader(b)) + if err != nil { + return err + } + req.Header.Set("Content-Type", "application/json") + client := http.Client{Timeout: 45 * time.Second, CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }} + resp, err := client.Do(req) + if err != nil { + return err + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + b, _ := io.ReadAll(io.LimitReader(resp.Body, 4096)) + return fmt.Errorf("inventory HTTP %d: %s", resp.StatusCode, b) + } + d := json.NewDecoder(io.LimitReader(resp.Body, 65536)) + if err = d.Decode(out); err != nil { + return err + } + if d.Decode(&struct{}{}) != io.EOF { + return fmt.Errorf("invalid acknowledgement") + } + return nil +} +func inventoryBackoff(attempt int) time.Duration { + if attempt > 8 { + attempt = 8 + } + return time.Second*time.Duration(1<= maxBaselineBuildAttempts { + current.BuildBlocked = true + current.Error = "baseline construction failed three times; repair the cause and explicitly reconcile again: " + err.Error() + status = "needs_attention" + } + } + return nil + }) + if updateErr != nil { + return fmt.Errorf("persist baseline retry state: %w (build error: %v)", updateErr, err) + } + current, readErr := inventory.Read(root) + if readErr == nil && current.ReconcileEpoch == q.ReconcileEpoch { + s.reportReconcileProgress(ctx, q.ReconcileEpoch, status, current.Error) + } + } + return err + } + if q.Epoch == "" { + return fmt.Errorf("inventory baseline not initialized") + } + for id, e := range q.Entries { + if ctx.Err() != nil { + return ctx.Err() + } + if time.Now().Before(e.Due) { + continue + } + err = s.sendRepository(ctx, e.Key) + if err != nil { + _ = inventory.Update(root, func(q *inventory.State) error { + if current := q.Entries[id]; current != nil { + current.Attempts++ + current.Due = time.Now().Add(inventoryBackoff(current.Attempts)) + current.Error = err.Error() + } + return nil + }) + } + } + return nil +} +func repoKey(k inventory.Key) repository.RepoKey { + return repository.RepoKey{Namespace: k.Namespace, RepoType: k.RepoType, Repo: k.Repo} +} +func (s *SchedulerService) recoverInventory(k inventory.Key) error { + return dao.NewUploadDao(s.metaService.fileDao, nil).RecoverInventory(repoKey(k)) +} + +func (s *SchedulerService) sendRepository(ctx context.Context, k inventory.Key) error { + if err := s.recoverInventory(k); err != nil { + return err + } + root := config.SysConfig.Repos() + p, err := func() (*inventory.Report, error) { + inventory.Barrier.Lock() + defer inventory.Barrier.Unlock() + q, err := inventory.Read(root) + if err != nil { + return nil, err + } + e := q.Entries[k.ID()] + if e == nil { + return nil, nil + } + if q.ReconcileEpoch != "" { + return nil, nil + } + if e.Operation != nil { + return nil, fmt.Errorf("repository operation incomplete") + } + if e.Report != nil { + return e.Report, nil + } + p := &inventory.Report{Version: 2, InstanceID: q.NodeID, Epoch: q.Epoch, Sequence: e.Sequence, Key: k, Deleted: e.Deleted, Files: []inventory.File{}} + if !e.Deleted { + d, err := repository.Read(root, repoKey(k)) + if err != nil { + return nil, err + } + files, err := s.scanInventory([]repository.Descriptor{d}) + if err != nil { + return nil, err + } + p.Files = files + } + err = inventory.Update(root, func(q *inventory.State) error { q.Entries[k.ID()].Report = p; return nil }) + return p, err + }() + if err != nil || p == nil { + return err + } + var a inventory.Ack + if err = inventoryPOST(ctx, "/api/v1/upload-inventory/reports", p, &a); err != nil { + return err + } + return inventory.Confirm(root, p, a) +} +func (s *SchedulerService) scanInventory(descriptors []repository.Descriptor) ([]inventory.File, error) { + snap, err := s.scanUploadDescriptors(&UploadInventorySnapshot{Complete: true, Items: []UploadInventoryItem{}}, descriptors) + if err != nil { + return nil, err + } + if !snap.Complete { + return nil, fmt.Errorf("incomplete inventory: %s", snap.Error) + } + files := make([]inventory.File, 0, len(snap.Items)) + for _, f := range snap.Items { + files = append(files, inventory.File{Key: inventory.Key{Namespace: f.Namespace, RepoType: f.RepoType, Repo: f.Repo}, Path: f.Path, SHA256: f.SHA256, Size: f.Size}) + } + return files, nil +} +func (s *SchedulerService) sendBaseline(ctx context.Context) (resultErr error) { + building := false + defer func() { + if building && resultErr != nil { + resultErr = baselineBuildError{resultErr} + } + }() + root := config.SysConfig.Repos() + q, err := inventory.Read(root) + if err != nil { + return err + } + if q.BuildBlocked { + return fmt.Errorf("baseline requires explicit reconciliation") + } + if q.Baseline == nil { + // Authenticate the requested generation before doing expensive disk work. + if err = s.openInventorySession(ctx); err != nil { + return err + } + building = true + s.reportReconcileProgress(ctx, q.ReconcileEpoch, "scanning", "") + q, err = inventory.Read(root) + if err != nil { + return err + } + for _, e := range q.Entries { + if e.Operation != nil { + if err = s.recoverInventory(e.Key); err != nil { + return err + } + } + } + err = func() error { + inventory.Barrier.Lock() + defer inventory.Barrier.Unlock() + q, err = inventory.Read(root) + if err != nil { + return err + } + for _, e := range q.Entries { + if e.Operation != nil { + return fmt.Errorf("local operation recovery pending") + } + } + descriptors, err := repository.ListHosted(root) + if err != nil { + return err + } + files, err := s.scanInventory(descriptors) + if err != nil { + return err + } + // Inventory trusts Speed's published metadata. scanInventory checks + // manifests, blob presence, sizes and completion markers; rebuilding + // the baseline must not read payloads to recompute their SHA256. + p := &inventory.Report{Version: 2, InstanceID: q.NodeID, Epoch: q.ReconcileEpoch, Baseline: true, Files: files} + return inventory.Update(root, func(st *inventory.State) error { + if st.ReconcileEpoch != q.ReconcileEpoch { + return fmt.Errorf("reconciliation changed") + } + st.Cutoff = st.Sequence + st.Baseline = p + return nil + }) + }() + if err != nil { + return err + } + } + building = false + q, err = inventory.Read(root) + if err != nil { + return err + } + if q.Baseline == nil { + return fmt.Errorf("baseline unavailable") + } + var a inventory.Ack + if err = inventoryPOST(ctx, "/api/v1/upload-inventory/reports", q.Baseline, &a); err != nil { + return err + } + return inventory.Confirm(root, q.Baseline, a) +} + +func (s *SchedulerService) reportReconcileProgress(ctx context.Context, epoch, status, message string) { + call, cancel := context.WithTimeout(ctx, 3*time.Second) + defer cancel() + var result map[string]bool + _ = inventoryPOST(call, "/api/v1/upload-inventory/nodes/"+url.PathEscape(config.SysConfig.Registration().NodeID)+"/reconcile-progress", map[string]string{"epoch": epoch, "status": status, "error": message}, &result) +} diff --git a/internal/service/inventory_outbox_test.go b/internal/service/inventory_outbox_test.go new file mode 100644 index 0000000..062e7ad --- /dev/null +++ b/internal/service/inventory_outbox_test.go @@ -0,0 +1,135 @@ +package service + +import ( + "bytes" + "context" + "crypto/sha256" + "dingospeed/internal/dao" + "dingospeed/internal/data" + "dingospeed/pkg/config" + "dingospeed/pkg/inventory" + "dingospeed/pkg/repository" + "encoding/hex" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "testing" +) + +func TestOutboxFixedRetryRecoveryAndNoIdleOrUnrelatedScan(t *testing.T) { + old := config.SysConfig + t.Cleanup(func() { config.SysConfig = old }) + root := t.TempDir() + epoch := "11111111-1111-1111-1111-111111111111" + var reports [][]byte + fail := false + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/api/v1/upload-inventory/nodes/node-a/session" { + _ = json.NewEncoder(w).Encode(inventory.Session{InstanceID: "node-a", PendingEpoch: epoch}) + return + } + if r.URL.Path != "/api/v1/upload-inventory/reports" { + _, _ = w.Write([]byte(`{}`)) + return + } + b, _ := io.ReadAll(r.Body) + reports = append(reports, b) + var p inventory.Report + _ = json.Unmarshal(b, &p) + if fail { + fail = false + w.WriteHeader(503) + return + } + _ = json.NewEncoder(w).Encode(inventory.Ack{Epoch: p.Epoch, Sequence: p.Sequence, Digest: p.Digest(), Status: "accepted"}) + })) + defer server.Close() + config.SysConfig = &config.Config{Server: config.ServerConfig{Repos: root}, Download: config.Download{BlockSize: 16}, Upload: config.Upload{Namespace: "dingo-local"}, Scheduler: config.Scheduler{HTTPURL: server.URL, Discovery: config.Discovery{InstanceId: "node-a"}}} + base := data.NewBaseData() + locks := dao.NewLockDao(base) + files := dao.NewFileDao(nil, base, locks) + uploads := dao.NewUploadDao(files, locks) + svc := &SchedulerService{metaService: NewMetaService(files, nil)} + ctx := context.Background() + if err := svc.openInventorySession(ctx); err != nil { + t.Fatal(err) + } + if err := svc.SyncPublications(ctx); err != nil { + t.Fatal(err) + } + key := inventory.Key{Namespace: "team", RepoType: "models", Repo: "nested/repo"} + upload := func(path string) { + t.Helper() + body := []byte(path) + hash := sha256.Sum256(body) + _, err := uploads.UploadWholeFile(dao.LocalUploadParam{Namespace: key.Namespace, RepoType: key.RepoType, Repo: key.Repo, Revision: "main", FilePath: path, Size: int64(len(body)), Sha256: hex.EncodeToString(hash[:])}, bytes.NewReader(body)) + if err != nil { + t.Fatal(err) + } + } + upload("a") + fail = true + if err := svc.sendRepository(ctx, key); err == nil { + t.Fatal("expected lost acknowledgement") + } + fixed := append([]byte{}, reports[len(reports)-1]...) + upload("b") + if err := svc.sendRepository(ctx, key); err != nil { + t.Fatal(err) + } + if !bytes.Equal(fixed, reports[len(reports)-1]) { + t.Fatal("retry report changed") + } + q, _ := inventory.Read(root) + if len(q.Entries) != 1 { + t.Fatal("ACK cleared concurrent change") + } + // A malformed unrelated repository would fail any accidental full scan. + bad := repository.RepoKey{Namespace: "team", RepoType: "models", Repo: "broken"} + if err := os.MkdirAll(bad.APIRoot(root), 0755); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(bad.APIRoot(root), repository.Marker), []byte("broken"), 0600); err != nil { + t.Fatal(err) + } + if err := svc.sendRepository(ctx, key); err != nil { + t.Fatal("scanned unrelated repository:", err) + } + before := len(reports) + if err := svc.SyncPublications(ctx); err != nil { + t.Fatal(err) + } + if len(reports) != before { + t.Fatal("idle queue sent inventory") + } + q, _ = inventory.Read(root) + if len(q.Entries) != 0 { + t.Fatal("queue not drained") + } + // Simulate process death after a journaled publish started but before its + // revision metadata finished. Recovery uses the fixed target, not a full scan. + commit, manifest, err := dao.ReadInventorySnapshot(repoKey(key), "main") + if err != nil { + t.Fatal(err) + } + payload, _ := json.Marshal(map[string]any{"Revision": "main", "Commit": commit, "Files": manifest}) + if err = inventory.Update(root, func(q *inventory.State) error { + q.Sequence++ + q.Entries[key.ID()] = &inventory.Entry{Key: key, Sequence: q.Sequence, Operation: &inventory.Operation{Kind: "metadata", Data: payload}} + return nil + }); err != nil { + t.Fatal(err) + } + if err = os.Remove(filepath.Join(repoKey(key).Revision(root, "main"), "meta_get.json")); err != nil { + t.Fatal(err) + } + if err = svc.sendRepository(ctx, key); err != nil { + t.Fatal("failed journal recovery:", err) + } + if _, _, err = dao.ReadInventorySnapshot(repoKey(key), "main"); err != nil { + t.Fatal("revision not recovered:", err) + } +} diff --git a/internal/service/local_operation_service.go b/internal/service/local_operation_service.go index 1ae4e1d..b8b9b61 100644 --- a/internal/service/local_operation_service.go +++ b/internal/service/local_operation_service.go @@ -28,6 +28,7 @@ import ( "dingospeed/pkg/config" "dingospeed/pkg/consts" "dingospeed/pkg/proto/manager" + "dingospeed/pkg/repository" "dingospeed/pkg/util" "github.com/bytedance/sonic" @@ -45,7 +46,7 @@ func NewLocalOperationService(schedulerDao *dao.SchedulerDao) *LocalOperationSer } func (l *LocalOperationService) Initialize() { - if config.SysConfig.GetOriginSchedulerModel() == consts.SchedulerModeCluster { + { // 将失败的信息存储到本地 go l.storeLocalOperation() if config.SysConfig.Online() { @@ -158,6 +159,13 @@ func (l *LocalOperationService) readAndSyncFileProcess(filePath string) error { zap.S().Errorf("解析 JSON 失败: %v, 内容: %s", err, jsonStr) return err } + if localOperation.IdentityVersion != 2 && (cacheJobReq.Org != "" || cacheJobReq.Repo != "") { + k, e := repository.FromStorage("models", cacheJobReq.Org, cacheJobReq.Repo) + if e != nil { + return e + } + cacheJobReq.Org, cacheJobReq.Repo = k.Namespace, k.Repo + } if err = l.schedulerDao.UpdateCacheJobStatus(&cacheJobReq); err != nil { return err } @@ -176,6 +184,13 @@ func (l *LocalOperationService) readAndSyncFileProcess(filePath string) error { zap.S().Errorf("解析 JSON 失败: %v, 内容: %s", err, jsonStr) return err } + if localOperation.IdentityVersion != 2 { + k, e := repository.FromStorage(process.Datatype, process.Org, process.Repo) + if e != nil { + return e + } + process.Org, process.Repo = k.Namespace, k.Repo + } processParams = append(processParams, &process) count++ if count == batchSize { @@ -184,6 +199,7 @@ func (l *LocalOperationService) readAndSyncFileProcess(filePath string) error { return err } count = 0 + processParams = nil } } } @@ -208,7 +224,7 @@ func (l *LocalOperationService) SyncFileProcess(processParams []*data.FileProces Repo: param.Repo, Name: param.Name, Etag: param.Etag, - InstanceId: config.SysConfig.Scheduler.Discovery.InstanceId, + InstanceId: config.SysConfig.Registration().NodeID, StartPos: param.StartPos, EndPos: param.EndPos, FileSize: param.FileSize, diff --git a/internal/service/meta_service.go b/internal/service/meta_service.go index df7ba03..ebd0919 100644 --- a/internal/service/meta_service.go +++ b/internal/service/meta_service.go @@ -17,10 +17,13 @@ package service import ( "archive/zip" "crypto/sha256" + "dingospeed/pkg/repository" "encoding/hex" "fmt" "io" + "net/url" "path" + "path/filepath" "sort" "strings" "time" @@ -74,6 +77,53 @@ func (m *MetaService) GetMetadata(repoType, orgRepo, revision, method, authoriza } func (m *MetaService) GetRepoTree(repoType, orgRepo, revision, pathInRepo string, recursive bool, authorization string) ([]RepoTreeItem, error) { + if !dao.IsLocalOrgRepo(orgRepo) { + upstream, err := dao.UpstreamRepo(repoType, orgRepo) + if err != nil { + return nil, err + } + uri := "/api/" + repoType + "/" + escapeRepoPath(upstream) + "/tree/" + url.PathEscape(revision) + if pathInRepo != "" { + uri += "/" + escapeRepoPath(pathInRepo) + } + if recursive { + uri += "?recursive=true" + } + headers := map[string]string{} + if authorization != "" { + headers["Authorization"] = authorization + } + items := []RepoTreeItem{} + seen := map[string]bool{} + for page := 0; page < 1000; page++ { + if seen[uri] { + return nil, fmt.Errorf("upstream tree pagination loop") + } + seen[uri] = true + resp, err := util.Get(uri, headers) + if err != nil { + return nil, err + } + if resp.StatusCode != 200 { + return nil, echo.NewHTTPError(resp.StatusCode, "upstream tree request rejected") + } + var batch []RepoTreeItem + if err = sonic.Unmarshal(resp.Body, &batch); err != nil { + return nil, err + } + items = append(items, batch...) + next, err := hfTreeNextURI(resp.GetKey("Link"), uri) + if err != nil { + return nil, err + } + if next == "" { + return items, nil + } + uri = next + } + return nil, fmt.Errorf("upstream tree pagination exceeds limit") + } + commitSha, err := m.fileDao.GetFileCommitSha(repoType, orgRepo, revision, authorization, "meta") if err != nil { return nil, err @@ -171,6 +221,41 @@ func (m *MetaService) StreamLocalArchive(repoType, orgRepo, revision, repo strin return archive.Close() } +// Only carry a pagination query back to the configured provider and same tree +// path. A Link response cannot redirect a machine/provider credential elsewhere. +func hfTreeNextURI(link, current string) (string, error) { + for _, part := range strings.Split(link, ",") { + if !strings.Contains(part, `rel="next"`) && !strings.Contains(part, "rel=next") { + continue + } + start, end := strings.Index(part, "<"), strings.Index(part, ">") + if start < 0 || end <= start { + return "", fmt.Errorf("invalid upstream pagination link") + } + next, err := url.Parse(part[start+1 : end]) + if err != nil { + return "", err + } + base, err := url.Parse(config.SysConfig.GetHFURLBase()) + if err != nil { + return "", err + } + currentURL, _ := url.Parse(current) + if next.Host != "" && next.Host != base.Host && next.Host != "huggingface.co" { + return "", fmt.Errorf("upstream pagination host mismatch") + } + nextPath := next.Path + if next.Host == base.Host && base.Path != "" { + nextPath = strings.TrimPrefix(nextPath, strings.TrimRight(base.Path, "/")) + } + if nextPath != currentURL.Path { + return "", fmt.Errorf("upstream pagination path mismatch") + } + return repository.EscapeURLPath(nextPath) + "?" + next.RawQuery, nil + } + return "", nil +} + func stableTreeID(path string) string { sum := sha256.Sum256([]byte(path)) return hex.EncodeToString(sum[:]) @@ -198,7 +283,7 @@ func (m *MetaService) RepoRefs(c echo.Context, repoType, org, repo string) error return util.ErrorRepoNotFound(c) } authorization := c.Request().Header.Get("authorization") - localRefsDir := fmt.Sprintf("%s/api/%s/%s/refs", config.SysConfig.Repos(), repoType, orgRepo) + localRefsDir := filepath.Join(dao.RepositoryKey(repoType, orgRepo).APIRoot(config.SysConfig.Repos()), "refs") localRefsPath := fmt.Sprintf("%s/%s", localRefsDir, fmt.Sprintf("refs_get.json")) err := util.MakeDirs(localRefsPath) if err != nil { @@ -255,9 +340,10 @@ func (m *MetaService) ForwardToNewSite(c echo.Context) error { originalLink := strings.Join(v, ", ") newLink := strings.ReplaceAll( originalLink, - "https://huggingface.co", - config.SysConfig.Scheduler.LinkDomain, + config.SysConfig.GetHFURLBase(), + providerLinkBase(c), ) + newLink = strings.ReplaceAll(newLink, "https://huggingface.co", providerLinkBase(c)) response.Header()[k] = []string{newLink} } else { response.Header()[k] = v @@ -272,11 +358,12 @@ func (m *MetaService) ForwardToNewSite(c echo.Context) error { } func (m *MetaService) RepositoryFiles(repoType, orgRepo, commit, filePath string) ([]*FileDescribe, error) { - pathsInfoShaDir := fmt.Sprintf("%s/api/%s/%s/paths-info/%s", config.SysConfig.Repos(), repoType, orgRepo, commit) + pathsInfoShaDir := dao.RepositoryKey(repoType, orgRepo).PathsInfo(config.SysConfig.Repos(), commit, "") if filePath != "" { pathsInfoShaDir += fmt.Sprintf("/%s", filePath) } - downloadLinkRoot := fmt.Sprintf("%s/%s/%s/resolve/%s", config.SysConfig.Scheduler.PublicDomain, repoType, orgRepo, commit) + key, _ := repository.ParseID(repoType, orgRepo) + downloadLinkRoot := strings.TrimRight(config.SysConfig.Scheduler.PublicDomain, "/") + "/api/repositories/" + url.PathEscape(repoType) + "/" + url.PathEscape(key.Namespace) + "/file?repo=" + url.QueryEscape(key.Repo) + "&revision=" + url.QueryEscape(commit) + "&path=" if dao.IsLocalOrgRepo(orgRepo) { return m.localRepositoryFiles(repoType, orgRepo, commit, filePath, downloadLinkRoot) } @@ -302,7 +389,7 @@ func (m *MetaService) RepositoryFiles(repoType, orgRepo, commit, filePath string } else { filePathName = item } - fileDescribe.Link = fmt.Sprintf("%s/%s", downloadLinkRoot, filePathName) + fileDescribe.Link = downloadLinkRoot + url.QueryEscape(filePathName) } fileDescribes = append(fileDescribes, fileDescribe) } @@ -343,7 +430,7 @@ func (m *MetaService) localRepositoryFiles(repoType, orgRepo, commit, filePath, fileDescribes = append(fileDescribes, &FileDescribe{ Name: rest, Size: item.Size, - Link: fmt.Sprintf("%s/%s", downloadLinkRoot, item.Path), + Link: downloadLinkRoot + url.QueryEscape(item.Path), }) } // 清单本身就是空的（新建还没加文件、或者被清空的 revision），那么仓库根目录下 @@ -409,3 +496,24 @@ type FileDescribe struct { IsDir bool `json:"isDir"` Link string `json:"link"` } + +func (m *MetaService) RevisionCommit(repoType, id, revision string) (string, error) { + return m.fileDao.GetCommitHfOffline(repoType, id, revision) +} + +func escapeRepoPath(path string) string { + parts := strings.Split(path, "/") + for i := range parts { + parts[i] = url.PathEscape(parts[i]) + } + return strings.Join(parts, "/") +} + +func providerLinkBase(c echo.Context) string { + base := config.SysConfig.Scheduler.LinkDomain + if base == "" { + base = c.Scheme() + "://" + c.Request().Host + } + prefix, _ := c.Get("providerPrefix").(string) + return strings.TrimRight(base, "/") + prefix +} diff --git a/internal/service/meta_service_local_test.go b/internal/service/meta_service_local_test.go index 39ce7f9..b381b57 100644 --- a/internal/service/meta_service_local_test.go +++ b/internal/service/meta_service_local_test.go @@ -40,13 +40,13 @@ func uploadLocalRepoRevision(t *testing.T, revision string, paths []string) (*Me content := []byte("content of " + p) sum := sha256.Sum256(content) result, err := uploadDao.UploadWholeFile(dao.LocalUploadParam{ - RepoType: "models", - Org: "dingo-local", - Repo: "demo", - Revision: revision, - FilePath: p, - Size: int64(len(content)), - Sha256: hex.EncodeToString(sum[:]), + RepoType: "models", + Namespace: "dingo-local", + Repo: "demo", + Revision: revision, + FilePath: p, + Size: int64(len(content)), + Sha256: hex.EncodeToString(sum[:]), }, bytes.NewReader(content)) if err != nil { t.Fatalf("upload %s failed: %v", p, err) @@ -99,7 +99,7 @@ func TestRepositoryFilesForLocalRepo(t *testing.T) { if root[1].IsDir || root[1].Name != "config.json" || root[1].Size == 0 { t.Fatalf("unexpected root file entry: %+v", root[1]) } - wantLink := "http://speed.local/models/" + orgRepo + "/resolve/" + commit + "/config.json" + wantLink := "http://speed.local/api/repositories/models/dingo-local/file?repo=demo&revision=" + commit + "&path=config.json" if root[1].Link != wantLink { t.Fatalf("unexpected download link:\n got %s\nwant %s", root[1].Link, wantLink) } diff --git a/internal/service/modelscope_file.go b/internal/service/modelscope_file.go new file mode 100644 index 0000000..179fd57 --- /dev/null +++ b/internal/service/modelscope_file.go @@ -0,0 +1,151 @@ +package service + +import ( + "context" + "fmt" + "net/http" + "net/url" + "strings" + + "dingospeed/internal/downloader" + "dingospeed/pkg/common" + "dingospeed/pkg/config" + myerr "dingospeed/pkg/error" + "dingospeed/pkg/repository" + "dingospeed/pkg/util" + "github.com/labstack/echo/v4" + "go.uber.org/zap" +) + +func (m *ModelscopeService) HandleFileDownload(c echo.Context, owner, repo, repoType string) error { + q, err := url.ParseQuery(c.Request().URL.RawQuery) + if err != nil || len(q["Revision"]) > 1 || len(q["FilePath"]) != 1 { + return echo.NewHTTPError(400, "invalid file locator") + } + return m.serveFile(c, repository.RepoKey{Namespace: repository.ModelScope, RepoType: repoType, Repo: owner + "/" + repo}, q.Get("Revision"), q.Get("FilePath")) +} + +func (m *ModelscopeService) serveFile(c echo.Context, k repository.RepoKey, revision, path string) error { + if revision == "" { + revision = "master" + } + if err := k.Validate(); err != nil { + return echo.NewHTTPError(400, err.Error()) + } + if err := repository.Segment(revision); err != nil { + return echo.NewHTTPError(400, err.Error()) + } + if err := repository.Relative(path); err != nil { + return echo.NewHTTPError(400, err.Error()) + } + // A file-list request both resolves the mutable revision and checks the caller's access. + envelope, err := m.repositoryJSON(c, k, "/repo/files", url.Values{"Revision": {revision}, "Recursive": {"true"}}) + if err != nil { + return err + } + var info *common.PathsInfo + commit := "" + for _, entry := range envelope.Data.Files { + name := entry.Path + if name == "" { + name = entry.Name + } + if name != path || entry.Type == "tree" || entry.Type == "directory" { + continue + } + if info != nil { + return echo.NewHTTPError(502, "ambiguous upstream file") + } + oid := entry.Sha256 + if oid == "" { + oid = entry.Sha1 + } + if repository.Segment(oid) != nil || entry.Size < 0 { + return echo.NewHTTPError(502, "missing upstream file identity") + } + commit = entry.Revision + // Some endpoints return the commit on the revision record instead of each file. + if commit == "" { + refs, e := m.RepositoryRevisions(c, k) + if e != nil { + return e + } + for _, ref := range refs { + if (ref.Name == revision || ref.Commit == revision) && ref.Commit != ref.Name { + commit = ref.Commit + break + } + } + } + if repository.Segment(commit) != nil { + return echo.NewHTTPError(502, "upstream did not provide a file revision") + } + info = &common.PathsInfo{Type: "file", Path: path, Oid: oid, Size: entry.Size} + } + if info == nil { + return echo.NewHTTPError(404, "file not found") + } + // Record a repository only when the upstream supplies repository-level + // identity. Missing repository metadata never raises single-file admission. + if envelope.Data.LatestCommitter.Id != "" || envelope.Data.LatestCommitter.ShortId != "" { + if _, err = m.cacheRepositoryMetadata(k, revision, envelope); err != nil { + // An unrelated sibling must not raise admission for this valid file. + // Explicit repository preheating still reports metadata errors. + zap.S().Warnw("Could not record ModelScope repository snapshot", "repo", k.Repo, "error", err) + } + } + // Use the same paths-info envelope as HF, without inventing a full repository snapshot. + if err = m.fileDao.CacheProviderFileInfo(k, commit, revision, info); err != nil { + return err + } + uri := "/api/v1/" + k.RepoType + "/" + repository.EscapeURLPath(k.Repo) + "/repo?" + url.Values{"Revision": {commit}, "FilePath": {path}}.Encode() + headers := map[string]string{} + for _, name := range []string{"Authorization", "Cookie", "User-Agent"} { + if v := c.Request().Header.Get(name); v != "" { + headers[strings.ToLower(name)] = v + } + } + source := &downloader.RemoteSource{Domain: strings.TrimRight(config.SysConfig.Modelscope.OfficialBaseURL, "/"), Headers: headers, Fetch: modelscopeFetch(info.Size)} + return m.fileDao.ServePreparedFile(c, k, commit, path, strings.ToLower(c.Request().Method), info, uri, source) +} + +// ModelScope may return 200 with Content-Range for a partial response. Normalize +// that transport detail, retaining the shared streaming/retry/block machinery. +func modelscopeFetch(size int64) func(context.Context, string, string, map[string]string, func(*http.Response) error) error { + return func(ctx context.Context, domain, uri string, headers map[string]string, consume func(*http.Response) error) error { + req, err := http.NewRequestWithContext(ctx, http.MethodGet, domain+uri, nil) + if err != nil { + return err + } + for k, v := range headers { + req.Header.Set(k, v) + } + req.Header.Set("Accept-Encoding", "identity") + resp, err := util.DoRequestWithRetry(req) + if err != nil { + return err + } + defer resp.Body.Close() + if resp.StatusCode != 200 && resp.StatusCode != 206 { + return myerr.NewAppendCode(resp.StatusCode, "ModelScope download rejected") + } + start, end, err := util.FileRange(req.Header.Get("Range"), size) + if err != nil { + return err + } + if cr := resp.Header.Get("Content-Range"); cr != "" { + var a, b, total int64 + if n, e := fmt.Sscanf(cr, "bytes %d-%d/%d", &a, &b, &total); e != nil || n != 3 || a != start || b != end-1 || total != size { + return fmt.Errorf("ModelScope returned a mismatched range") + } + resp.StatusCode = http.StatusPartialContent + } else if start != 0 || end != size { + return fmt.Errorf("ModelScope ignored the byte range") + } + if enc := resp.Header.Get("Content-Encoding"); enc != "" && enc != "identity" { + return fmt.Errorf("unexpected encoded ModelScope range") + } + resp.Header.Del("Content-Encoding") + return consume(resp) + } +} diff --git a/internal/service/modelscope_file_test.go b/internal/service/modelscope_file_test.go new file mode 100644 index 0000000..02ac8db --- /dev/null +++ b/internal/service/modelscope_file_test.go @@ -0,0 +1,659 @@ +package service + +import ( + "bytes" + "context" + "crypto/sha1" + "crypto/sha256" + "encoding/hex" + "encoding/json" + "fmt" + "io" + "net/http" + "net/http/httptest" + "net/url" + "os" + "os/exec" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "dingospeed/internal/data" + "dingospeed/internal/downloader" + "dingospeed/pkg/config" + "dingospeed/pkg/repository" + "dingospeed/pkg/util" + "github.com/labstack/echo/v4" + "go.uber.org/zap" + "go.uber.org/zap/zaptest/observer" +) + +func msConfig(t *testing.T, endpoint string) { + t.Helper() + old := config.SysConfig + t.Cleanup(func() { config.SysConfig = old }) + config.SysConfig = &config.Config{ + Server: config.ServerConfig{Repos: t.TempDir(), Online: true}, + Download: config.Download{BlockSize: 1024, RespChanSize: 32, RespChunkSize: 512, RemoteFileRangeSize: 4096, RemoteFileBufferSize: 4096, GoroutineMaxNumPerFile: 4, ReqTimeout: 5}, + Modelscope: config.Modelscope{OfficialBaseURL: endpoint, MaxRetry: 1}, + Retry: config.Retry{Attempts: 1}, + } + data.NewBaseData() +} + +func msEngine() *echo.Echo { + s := NewModelscopeService() + e := echo.New() + e.Add(http.MethodGet, "/api/v1/models/:owner/:repo/repo", func(c echo.Context) error { + return s.HandleFileDownload(c, c.Param("owner"), c.Param("repo"), "models") + }) + e.Add(http.MethodHead, "/api/v1/models/:owner/:repo/repo", func(c echo.Context) error { + return s.HandleFileDownload(c, c.Param("owner"), c.Param("repo"), "models") + }) + for _, method := range []string{http.MethodGet, http.MethodHead} { + e.Add(method, "/file", func(c echo.Context) error { + return s.RepositoryFile(c, repository.RepoKey{Namespace: repository.ModelScope, RepoType: "models", Repo: c.QueryParam("repo")}, c.QueryParam("revision"), c.QueryParam("path")) + }) + } + return e +} + +func msRequest(e *echo.Echo, method, path, rng string) *httptest.ResponseRecorder { + req := httptest.NewRequest(method, path, nil) + if rng != "" { + req.Header.Set("Range", rng) + } + rec := httptest.NewRecorder() + e.ServeHTTP(rec, req) + return rec +} + +func msPayload(t *testing.T, key repository.RepoKey, oid string) ([]byte, *downloader.DingCacheHeader) { + t.Helper() + f, err := os.Open(key.Blob(config.SysConfig.Repos(), oid)) + if err != nil { + t.Fatal(err) + } + defer f.Close() + h := &downloader.DingCacheHeader{} + if err = h.Read(f); err != nil { + t.Fatal(err) + } + b, err := io.ReadAll(f) + if err != nil { + t.Fatal(err) + } + return b, h +} + +func TestModelScopeSharedSingleFile(t *testing.T) { + content := bytes.Repeat([]byte("scope-0123456789!"), 1024) + var current atomic.Int64 + var bodies atomic.Int64 + var denied atomic.Bool + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Header.Get("X-Dingo-Service-Token") != "" || r.Header.Get("X-Dingo-Ensure-Cache") != "" { + t.Error("internal credential/control header leaked") + } + if denied.Load() { + w.WriteHeader(403) + return + } + data := content + if current.Load() == 1 { + data = bytes.Repeat([]byte("changed-7654321!!"), 1024) + } + file := r.URL.Query().Get("FilePath") + revision := fmt.Sprintf("commit-%d", current.Load()) + if strings.HasSuffix(r.URL.Path, "/repo/files") { + files := []map[string]any{} + for _, name := range []string{"nested/weights.bin", "copy.bin", "sha1.txt", "empty.txt"} { + b := data + if name == "empty.txt" { + b = nil + } + hash := sha256.Sum256(b) + entry := map[string]any{"Type": "blob", "Path": name, "Revision": revision, "Size": len(b), "Sha256": hex.EncodeToString(hash[:])} + if name == "sha1.txt" { + delete(entry, "Sha256") + h := sha1.Sum(b) + entry["Sha1"] = hex.EncodeToString(h[:]) + } + files = append(files, entry) + } + json.NewEncoder(w).Encode(map[string]any{"Code": 200, "Data": map[string]any{"Files": files}}) + return + } + if r.URL.Query().Get("Revision") != revision { + t.Errorf("download was not pinned: %s", r.URL.RawQuery) + } + bodies.Add(1) + if file == "empty.txt" { + data = nil + } + start, end, err := util.FileRange(r.Header.Get("Range"), int64(len(data))) + if err != nil { + w.WriteHeader(416) + return + } + if r.Header.Get("Range") != "" { + w.Header().Set("Content-Range", fmt.Sprintf("bytes %d-%d/%d", start, end-1, len(data))) + } + // Match the real provider's 200 + Content-Range response. + w.Header().Set("Content-Length", fmt.Sprint(end-start)) + w.Write(data[start:end]) + })) + defer upstream.Close() + msConfig(t, upstream.URL) + // LocalOnly must suppress ALL cluster calls even in a cluster-configured Speed. + config.SysConfig.SetSchedulerModel("cluster") + e := msEngine() + key := repository.RepoKey{Namespace: repository.ModelScope, RepoType: "models", Repo: "Qwen/demo"} + native := "/api/v1/models/Qwen/demo/repo?Revision=master&FilePath=nested%2Fweights.bin" + canonical := "/file?repo=Qwen%2Fdemo&revision=master&path=nested%2Fweights.bin" + for _, tc := range []struct { + method, path, rng string + code int + want []byte + }{ + {"HEAD", native, "", 200, nil}, + {"GET", native, "bytes=0-2047", 206, content[:2048]}, + {"GET", canonical, "", 200, content}, + } { + r := msRequest(e, tc.method, tc.path, tc.rng) + if r.Code != tc.code || !bytes.Equal(r.Body.Bytes(), tc.want) { + t.Fatalf("%s %s: %d %q", tc.method, tc.rng, r.Code, r.Body.Bytes()[:min(150, r.Body.Len())]) + } + } + hash := sha256.Sum256(content) + oid := hex.EncodeToString(hash[:]) + b, h := msPayload(t, key, oid) + if !bytes.Equal(b, content) || int(h.FileSize) != len(content) { + t.Fatal("wrong dingcache payload") + } + for i := uint64(0); i < h.BlockNumber; i++ { + set, _ := h.BlockMask.Test(i) + if !set { + t.Fatalf("missing block %d", i) + } + } + before := bodies.Load() + for _, rng := range []string{"", "bytes=0-0", "bytes=1024-", "bytes=-19"} { + start, end, _ := util.FileRange(rng, int64(len(content))) + r := msRequest(msEngine(), "GET", canonical, rng) // Recreate service and reopen the cached container. + if !bytes.Equal(r.Body.Bytes(), content[start:end]) { + t.Fatalf("cache range %s mismatch", rng) + } + } + if bodies.Load() != before { + t.Fatal("warm read fetched upstream body") + } + copyURL := "/file?repo=Qwen%2Fdemo&revision=master&path=copy.bin" + if r := msRequest(e, "GET", copyURL, ""); !bytes.Equal(r.Body.Bytes(), content) { + t.Fatal("shared content failed") + } + if bodies.Load() != before { + t.Fatal("same OID different path did not reuse blob") + } + for _, path := range []string{"copy.bin", "nested/weights.bin"} { + if _, err := os.Stat(key.Resolve(config.SysConfig.Repos(), "commit-0", path)); err != nil { + t.Fatal(err) + } + } + for _, rng := range []string{"bytes=999999-", "bytes=8-2", "bytes=1-2,4-5"} { + if r := msRequest(e, "GET", canonical, rng); r.Code != 416 { + t.Fatalf("bad range %s returned %d", rng, r.Code) + } + } + if r := msRequest(e, "GET", "/file?repo=Qwen%2Fdemo&revision=master&path=sha1.txt", ""); r.Code != 200 || !bytes.Equal(r.Body.Bytes(), content) { + t.Fatal("SHA1-only file rejected") + } + if r := msRequest(e, "GET", "/file?repo=Qwen%2Fdemo&revision=master&path=empty.txt", ""); r.Code != 200 || r.Body.Len() != 0 { + t.Fatal("zero-byte file rejected") + } + denied.Store(true) + if r := msRequest(e, "GET", canonical, ""); r.Code != 403 { + t.Fatalf("cached bytes bypassed auth: %d", r.Code) + } + denied.Store(false) + current.Store(1) + if r := msRequest(e, "GET", canonical, ""); r.Code != 200 || bytes.Equal(r.Body.Bytes(), content) { + t.Fatal("stale branch content") + } + old, _ := msPayload(t, key, oid) + if !bytes.Equal(old, content) { + t.Fatal("branch update overwrote old blob") + } +} + +func TestModelScopeStreamingAndCacheCompletion(t *testing.T) { + content := bytes.Repeat([]byte("streaming-content"), 1024) + hash := sha256.Sum256(content) + release := make(chan struct{}) + var once sync.Once + defer once.Do(func() { close(release) }) + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if strings.HasSuffix(r.URL.Path, "/repo/files") { + fmt.Fprintf(w, `{"Code":200,"Data":{"Files":[{"Type":"blob","Path":"file.bin","Revision":"pinned","Sha256":"%x","Size":%d}]}}`, hash, len(content)) + return + } + w.Header().Set("Content-Length", fmt.Sprint(len(content))) + w.Write(content[:1024]) + w.(http.Flusher).Flush() + select { + case <-release: + case <-r.Context().Done(): + return + } + w.Write(content[1024:]) + })) + defer upstream.Close() + msConfig(t, upstream.URL) + config.SysConfig.Download.RemoteFileRangeSize = 0 + speed := httptest.NewServer(msEngine()) + defer speed.Close() + req, _ := http.NewRequest("GET", speed.URL+"/file?repo=Qwen%2Fdemo&revision=master&path=file.bin", nil) + req.Header.Set("X-Dingo-Ensure-Cache", "1") + client := &http.Client{Timeout: 10 * time.Second} + resp, err := client.Do(req) + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + first := make([]byte, 512) + if _, err = io.ReadFull(resp.Body, first); err != nil { + t.Fatal(err) + } + if !bytes.Equal(first, content[:512]) { + t.Fatal("did not stream first bytes") + } + once.Do(func() { close(release) }) + rest, err := io.ReadAll(resp.Body) + if err != nil { + t.Fatal(err) + } + if !bytes.Equal(append(first, rest...), content) || resp.Trailer.Get("X-Dingo-Cache-Complete") != "true" { + t.Fatalf("cache completion missing: %v", resp.Trailer) + } +} + +func TestModelScopeLiveSingleFile(t *testing.T) { + if os.Getenv("DINGO_TEST_MODELSCOPE_LIVE") != "1" { + t.Skip("opt-in real provider test") + } + msConfig(t, "https://modelscope.cn") + config.SysConfig.Download.BlockSize = 64 * 1024 + config.SysConfig.Download.RemoteFileRangeSize = 256 * 1024 + config.SysConfig.Download.RespChunkSize = 32 * 1024 + config.SysConfig.Download.RemoteFileBufferSize = 256 * 1024 + e := msEngine() + for _, file := range []string{"config.json", "merges.txt"} { + target := "/file?" + url.Values{"repo": {"Qwen/Qwen2.5-0.5B-Instruct"}, "revision": {"master"}, "path": {file}}.Encode() + r := msRequest(e, "GET", target, "") + if r.Code != 200 { + t.Fatalf("live %s: %d %s", file, r.Code, r.Body.String()) + } + hash := sha256.Sum256(r.Body.Bytes()) + oid := hex.EncodeToString(hash[:]) + key := repository.RepoKey{Namespace: repository.ModelScope, RepoType: "models", Repo: "Qwen/Qwen2.5-0.5B-Instruct"} + payload, h := msPayload(t, key, oid) + if !bytes.Equal(payload, r.Body.Bytes()) { + t.Fatal("live payload mismatch") + } + warm := msRequest(msEngine(), "GET", target, "bytes=1-31") + if warm.Code != 206 || !bytes.Equal(warm.Body.Bytes(), payload[1:32]) { + t.Fatalf("live warm range mismatch: %d", warm.Code) + } + t.Logf("live file=%s commit=%s size=%d sha256=%s blocks=%d", file, r.Header().Get("X-Repo-Commit"), len(payload), oid, h.BlockNumber) + } +} + +func TestModelScopeLiveSDK(t *testing.T) { + sdkPath := os.Getenv("DINGO_TEST_MODELSCOPE_SDK_PATH") + if os.Getenv("DINGO_TEST_MODELSCOPE_LIVE") != "1" || sdkPath == "" { + t.Skip("opt-in installed SDK test") + } + msConfig(t, "https://modelscope.cn") + config.SysConfig.Download.BlockSize = 64 * 1024 + config.SysConfig.Download.RemoteFileRangeSize = 256 * 1024 + config.SysConfig.Download.RespChunkSize = 32 * 1024 + config.SysConfig.Download.RemoteFileBufferSize = 256 * 1024 + speed := httptest.NewServer(msEngine()) + defer speed.Close() + script := `import sys,hashlib,json +sys.path.insert(0,sys.argv[1]) +from modelscope_hub import HubApi, __version__ +api=HubApi(endpoint=sys.argv[2],token='') +expected={'config.json':'18e18afcaccafade98daf13a54092927904649e1dd4eba8299ab717d5d94ff45','merges.txt':'599bab54075088774b1733fde865d5bd747cbcc7a547c5bc12610e874e26f5e3'} +for attempt in range(2): + for name,sha in expected.items(): + p=api.download_file('Qwen/Qwen2.5-0.5B-Instruct',repo_type='model',file_path=name,revision='master',cache_dir=sys.argv[3],local_dir=sys.argv[3],force=True) + b=p.read_bytes();actual=hashlib.sha256(b).hexdigest() + assert actual==sha,(name,actual) + print(json.dumps({'sdk':__version__,'attempt':attempt,'file':name,'size':len(b),'sha256':actual})) +` + ctx, cancel := context.WithTimeout(context.Background(), 90*time.Second) + defer cancel() + cmd := exec.CommandContext(ctx, "python", "-c", script, sdkPath, speed.URL, t.TempDir()) + output, err := cmd.CombinedOutput() + t.Log(string(output)) + if err != nil { + t.Fatal(err) + } +} + +func TestModelScopeFailureDoesNotCompleteCache(t *testing.T) { + for _, mode := range []string{"truncated", "wrong-range", "disk-write"} { + t.Run(mode, func(t *testing.T) { + content := bytes.Repeat([]byte("failure-test-1234"), 1024) + hash := sha256.Sum256(content) + oid := hex.EncodeToString(hash[:]) + var failing atomic.Bool + failing.Store(true) + key := repository.RepoKey{Namespace: repository.ModelScope, RepoType: "models", Repo: "Qwen/demo"} + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if strings.HasSuffix(r.URL.Path, "/repo/files") { + fmt.Fprintf(w, `{"Code":200,"Data":{"Files":[{"Type":"blob","Path":"file.bin","Revision":"pinned","Sha256":"%s","Size":%d}]}}`, oid, len(content)) + return + } + start, end, _ := util.FileRange(r.Header.Get("Range"), int64(len(content))) + if failing.Load() && mode == "disk-write" { + path := key.Blob(config.SysConfig.Repos(), oid) + if err := os.Remove(path); err != nil { + t.Error(err) + } + if err := os.Mkdir(path, 0755); err != nil { + t.Error(err) + } + } + if failing.Load() && mode == "wrong-range" { + w.Header().Set("Content-Range", fmt.Sprintf("bytes 5-%d/%d", end-1, len(content))) + } else if r.Header.Get("Range") != "" { + w.Header().Set("Content-Range", fmt.Sprintf("bytes %d-%d/%d", start, end-1, len(content))) + } + w.Header().Set("Content-Length", fmt.Sprint(end-start)) + if failing.Load() && mode == "truncated" { + w.Write(content[start:min(end, start+1024)]) + return + } + w.Write(content[start:end]) + })) + defer upstream.Close() + msConfig(t, upstream.URL) + config.SysConfig.Download.RemoteFileRangeSize = 0 + speed := httptest.NewServer(msEngine()) + defer speed.Close() + request := func() ([]byte, string) { + req, _ := http.NewRequest("GET", speed.URL+"/file?repo=Qwen%2Fdemo&revision=master&path=file.bin", nil) + req.Header.Set("X-Dingo-Ensure-Cache", "1") + resp, err := (&http.Client{Timeout: 10 * time.Second}).Do(req) + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + b, _ := io.ReadAll(resp.Body) + return b, resp.Trailer.Get("X-Dingo-Cache-Complete") + } + body, complete := request() + if complete == "true" { + t.Fatalf("%s falsely completed cache", mode) + } + if mode == "disk-write" { + if !bytes.Equal(body, content) { + t.Fatal("cache failure interrupted otherwise valid upstream bytes") + } + if err := os.Remove(key.Blob(config.SysConfig.Repos(), oid)); err != nil { + t.Fatal(err) + } // Empty injected directory only. + } + failing.Store(false) + body, complete = request() + if complete != "true" || !bytes.Equal(body, content) { + t.Fatalf("retry did not recover: complete=%s size=%d", complete, len(body)) + } + payload, _ := msPayload(t, key, oid) + if !bytes.Equal(payload, content) { + t.Fatal("retry cache corrupted") + } + }) + } +} + +func TestModelScopeAutomaticRetryBudget(t *testing.T) { + for _, tc := range []struct { + name string + status, failures, calls int + complete bool + }{ + {"temporary-status", 503, 2, 3, true}, + {"short-body", 200, 2, 3, true}, + {"exhausted", 503, 10, 3, false}, + {"unauthorized", 401, 10, 1, false}, + {"forbidden", 403, 10, 1, false}, + {"not-found", 404, 10, 1, false}, + } { + t.Run(tc.name, func(t *testing.T) { + content := bytes.Repeat([]byte("auto-retry-12345!"), 1024) + oid := fmt.Sprintf("%x", sha256.Sum256(content)) + var calls atomic.Int32 + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if strings.HasSuffix(r.URL.Path, "/repo/files") { + fmt.Fprintf(w, `{"Code":200,"Data":{"Files":[{"Type":"blob","Path":"file.bin","Revision":"fixed","Sha256":%q,"Size":%d}]}}`, oid, len(content)) + return + } + n := calls.Add(1) + if int(n) <= tc.failures && tc.status != 200 { + w.WriteHeader(tc.status) + return + } + start, end, _ := util.FileRange(r.Header.Get("Range"), int64(len(content))) + if tc.name == "short-body" && start != int64(n-1)*1024 { + t.Errorf("retry offset=%d on attempt=%d", start, n) + } + w.Header().Set("Content-Length", fmt.Sprint(end-start)) + if start > 0 { + w.Header().Set("Content-Range", fmt.Sprintf("bytes %d-%d/%d", start, end-1, len(content))) + w.WriteHeader(206) + } + if int(n) <= tc.failures { + end = min(end, start+1024) + } + w.Write(content[start:end]) + })) + defer upstream.Close() + msConfig(t, upstream.URL) + config.SysConfig.Retry.Attempts = 3 + config.SysConfig.Download.RemoteFileRangeSize = 0 + speed := httptest.NewServer(msEngine()) + defer speed.Close() + req, _ := http.NewRequest("GET", speed.URL+"/file?repo=Qwen%2Fdemo&revision=fixed&path=file.bin", nil) + req.Header.Set("X-Dingo-Ensure-Cache", "1") + resp, err := (&http.Client{Timeout: 5 * time.Second}).Do(req) + if err != nil { + t.Fatal(err) + } + body, readErr := io.ReadAll(resp.Body) + resp.Body.Close() + if calls.Load() != int32(tc.calls) || (resp.Trailer.Get("X-Dingo-Cache-Complete") == "true") != tc.complete { + t.Fatalf("calls=%d trailers=%v error=%v", calls.Load(), resp.Trailer, readErr) + } + if tc.complete { + if readErr != nil || !bytes.Equal(body, content) { + t.Fatalf("retried response corrupted: bytes=%d error=%v", len(body), readErr) + } + actual, _ := msPayload(t, repository.RepoKey{Namespace: repository.ModelScope, RepoType: "models", Repo: "Qwen/demo"}, oid) + if !bytes.Equal(actual, content) { + t.Fatal("retried cache differs from source") + } + } else if !strings.Contains(resp.Trailer.Get("X-Dingo-Cache-Error"), fmt.Sprint(tc.status)) { + t.Fatalf("failure lost HTTP reason: %v", resp.Trailer) + } + }) + } +} + +func TestModelScopeConcurrentReaders(t *testing.T) { + core, logs := observer.New(zap.WarnLevel) + restore := zap.ReplaceGlobals(zap.New(core)) + defer restore() + defer func() { + if t.Failed() { + for _, entry := range logs.All() { + t.Log(entry.Message) + } + } + }() + content := bytes.Repeat([]byte("concurrent-bytes"), 1024) + hash := sha256.Sum256(content) + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if strings.HasSuffix(r.URL.Path, "/repo/files") { + fmt.Fprintf(w, `{"Code":200,"Data":{"Files":[{"Type":"blob","Path":"file.bin","Revision":"pinned","Sha256":"%x","Size":%d}]}}`, hash, len(content)) + return + } + start, end, _ := util.FileRange(r.Header.Get("Range"), int64(len(content))) + if r.Header.Get("Range") != "" { + w.Header().Set("Content-Range", fmt.Sprintf("bytes %d-%d/%d", start, end-1, len(content))) + } + w.Header().Set("Content-Length", fmt.Sprint(end-start)) + w.Write(content[start:end]) + })) + defer upstream.Close() + msConfig(t, upstream.URL) + e := msEngine() + handler := e.HTTPErrorHandler + e.HTTPErrorHandler = func(err error, c echo.Context) { + if !c.Response().Committed { + t.Logf("concurrent handler error: %v", err) + } + handler(err, c) + } + var wg sync.WaitGroup + for i := 0; i < 8; i++ { + wg.Add(1) + go func() { + defer wg.Done() + r := msRequest(e, "GET", "/file?repo=Qwen%2Fdemo&revision=master&path=file.bin", "") + if r.Code != 200 || !bytes.Equal(r.Body.Bytes(), content) { + t.Errorf("concurrent response mismatch: %d bytes=%d body=%q", r.Code, r.Body.Len(), r.Body.String()) + } + }() + } + wg.Wait() + key := repository.RepoKey{Namespace: repository.ModelScope, RepoType: "models", Repo: "Qwen/demo"} + payload, _ := msPayload(t, key, hex.EncodeToString(hash[:])) + if !bytes.Equal(payload, content) { + t.Fatal("concurrent cache corrupted") + } +} + +func TestModelScopeProcessReopen(t *testing.T) { + if root := os.Getenv("DINGO_MS_REOPEN_ROOT"); root != "" { + msConfig(t, os.Getenv("DINGO_MS_REOPEN_UPSTREAM")) + config.SysConfig.Server.Repos = root + r := msRequest(msEngine(), "GET", "/file?repo=Qwen%2Fdemo&revision=master&path=file.bin", "") + if r.Code != 200 || !bytes.Equal(r.Body.Bytes(), bytes.Repeat([]byte("persisted-bytes!"), 1024)) { + t.Fatal("new process did not read cached payload") + } + return + } + content := bytes.Repeat([]byte("persisted-bytes!"), 1024) + hash := sha256.Sum256(content) + var rejectBody atomic.Bool + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if strings.HasSuffix(r.URL.Path, "/repo/files") { + fmt.Fprintf(w, `{"Code":200,"Data":{"Files":[{"Type":"blob","Path":"file.bin","Revision":"pinned","Sha256":"%x","Size":%d}]}}`, hash, len(content)) + return + } + if rejectBody.Load() { + t.Error("restarted process refetched body") + w.WriteHeader(503) + return + } + start, end, _ := util.FileRange(r.Header.Get("Range"), int64(len(content))) + if r.Header.Get("Range") != "" { + w.Header().Set("Content-Range", fmt.Sprintf("bytes %d-%d/%d", start, end-1, len(content))) + } + w.Header().Set("Content-Length", fmt.Sprint(end-start)) + w.Write(content[start:end]) + })) + defer upstream.Close() + msConfig(t, upstream.URL) + r := msRequest(msEngine(), "GET", "/file?repo=Qwen%2Fdemo&revision=master&path=file.bin", "") + if r.Code != 200 || !bytes.Equal(r.Body.Bytes(), content) { + t.Fatal("cold fill failed") + } + rejectBody.Store(true) + exe, err := os.Executable() + if err != nil { + t.Fatal(err) + } + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + cmd := exec.CommandContext(ctx, exe, "-test.run=^TestModelScopeProcessReopen$", "-test.v") + cmd.Env = append(os.Environ(), "DINGO_MS_REOPEN_ROOT="+config.SysConfig.Repos(), "DINGO_MS_REOPEN_UPSTREAM="+upstream.URL) + out, err := cmd.CombinedOutput() + if err != nil { + t.Fatalf("restart: %v\n%s", err, out) + } +} + +func TestModelScopeCancelledTransferResumes(t *testing.T) { + content := bytes.Repeat([]byte("cancel-and-retry"), 1024) + hash := sha256.Sum256(content) + var first atomic.Bool + first.Store(true) + upstreamDone := make(chan struct{}) + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if strings.HasSuffix(r.URL.Path, "/repo/files") { + fmt.Fprintf(w, `{"Code":200,"Data":{"Files":[{"Type":"blob","Path":"file.bin","Revision":"pinned","Sha256":"%x","Size":%d}]}}`, hash, len(content)) + return + } + start, end, _ := util.FileRange(r.Header.Get("Range"), int64(len(content))) + if r.Header.Get("Range") != "" { + w.Header().Set("Content-Range", fmt.Sprintf("bytes %d-%d/%d", start, end-1, len(content))) + } + w.Header().Set("Content-Length", fmt.Sprint(end-start)) + if first.CompareAndSwap(true, false) { + defer close(upstreamDone) + w.Write(content[:2048]) + w.(http.Flusher).Flush() + <-r.Context().Done() + return + } + w.Write(content[start:end]) + })) + defer upstream.Close() + msConfig(t, upstream.URL) + config.SysConfig.Download.RemoteFileRangeSize = 0 + speed := httptest.NewServer(msEngine()) + defer speed.Close() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + req, _ := http.NewRequestWithContext(ctx, "GET", speed.URL+"/file?repo=Qwen%2Fdemo&revision=master&path=file.bin", nil) + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatal(err) + } + buf := make([]byte, 1024) + if _, err = io.ReadFull(resp.Body, buf); err != nil { + t.Fatal(err) + } + cancel() + resp.Body.Close() + select { + case <-upstreamDone: + case <-time.After(5 * time.Second): + t.Fatal("upstream was not cancelled") + } + r := msRequest(msEngine(), "GET", "/file?repo=Qwen%2Fdemo&revision=master&path=file.bin", "") + if r.Code != 200 || !bytes.Equal(r.Body.Bytes(), content) { + t.Fatalf("retry after cancellation failed: %d %d", r.Code, r.Body.Len()) + } + key := repository.RepoKey{Namespace: repository.ModelScope, RepoType: "models", Repo: "Qwen/demo"} + payload, _ := msPayload(t, key, hex.EncodeToString(hash[:])) + if !bytes.Equal(payload, content) { + t.Fatal("cancel/retry cache mismatch") + } +} diff --git a/internal/service/modelscope_metadata.go b/internal/service/modelscope_metadata.go new file mode 100644 index 0000000..4cc8cae --- /dev/null +++ b/internal/service/modelscope_metadata.go @@ -0,0 +1,63 @@ +package service + +import ( + "encoding/json" + "fmt" + "path/filepath" + + "dingospeed/pkg/common" + "dingospeed/pkg/config" + "dingospeed/pkg/repository" +) + +// LatestCommitter describes the complete requested tree. A Files[i].Revision +// describes only that file's last change and must never substitute for it. +func (m *ModelscopeService) cacheRepositoryMetadata(k repository.RepoKey, revision string, envelope *modelscopeEnvelope) (map[string]interface{}, error) { + commit := envelope.Data.LatestCommitter.Id + if commit == "" { + commit = envelope.Data.LatestCommitter.ShortId + } + if repository.Segment(revision) != nil || repository.Segment(commit) != nil { + return nil, fmt.Errorf("ModelScope response has no repository revision") + } + siblings := []map[string]interface{}{} + infos := []*common.PathsInfo{} + var size int64 + for _, file := range envelope.Data.Files { + if file.Type == "tree" || file.Type == "directory" { + continue + } + name := file.Path + if name == "" { + name = file.Name + } + oid := file.Sha256 + if oid == "" { + oid = file.Sha1 + } + if repository.Relative(name) != nil || repository.Segment(oid) != nil || file.Size < 0 { + return nil, fmt.Errorf("invalid ModelScope file metadata") + } + siblings = append(siblings, map[string]interface{}{"rfilename": name, "size": file.Size, "blobId": oid, "fileRevision": file.Revision}) + infos = append(infos, &common.PathsInfo{Type: "file", Path: name, Oid: oid, Size: file.Size}) + size += file.Size + } + meta := map[string]interface{}{"namespace": k.Namespace, "repoType": k.RepoType, "repo": k.Repo, "id": k.ID(), "sha": commit, "siblings": siblings, "usedStorage": size} + body, err := json.Marshal(meta) + if err != nil { + return nil, err + } + // Publish the revision pointer only after its complete tree is available. + for _, info := range infos { + if err = m.fileDao.CacheProviderFileInfo(k, commit, revision, info); err != nil { + return nil, err + } + } + for _, ref := range []string{commit, revision} { + target := filepath.Join(k.Revision(config.SysConfig.Repos(), ref), "meta_get.json") + if err = m.fileDao.CacheProviderResponse(target, commit, revision, body); err != nil { + return nil, err + } + } + return meta, nil +} diff --git a/internal/service/modelscope_repository.go b/internal/service/modelscope_repository.go new file mode 100644 index 0000000..e7e7708 --- /dev/null +++ b/internal/service/modelscope_repository.go @@ -0,0 +1,163 @@ +package service + +import ( + "encoding/json" + "fmt" + "io" + "net/http" + "net/url" + "strings" + + "dingospeed/pkg/common" + "dingospeed/pkg/config" + "dingospeed/pkg/repository" + "dingospeed/pkg/util" + "github.com/labstack/echo/v4" +) + +type ProviderRevision struct { + Name string `json:"name"` + Commit string `json:"commit"` +} + +type modelscopeEnvelope struct { + Code int + Message string + Data struct { + LatestCommitter struct{ Id, ShortId string } + Files []struct { + Type, Path, Name, Sha256, Sha1, Revision string + Size int64 + } + RevisionMap struct { + Branches, Tags []struct{ Revision, Sha, CommitId string } + } + } +} + +// ModelScope's own protocol remains byte-for-byte proxied. These methods only +// project its published Files/RevisionMap response into the control-plane API. +func (m *ModelscopeService) repositoryJSON(c echo.Context, k repository.RepoKey, operation string, q url.Values) (*modelscopeEnvelope, error) { + if err := k.Validate(); err != nil { + return nil, err + } + if k.Namespace != repository.ModelScope { + return nil, fmt.Errorf("not a ModelScope repository") + } + if err := repository.Register(config.SysConfig.Repos(), repository.Remote(k)); err != nil { + return nil, err + } + uri := strings.TrimRight(config.SysConfig.Modelscope.OfficialBaseURL, "/") + "/api/v1/" + k.RepoType + "/" + escapeRepoPath(k.Repo) + operation + if len(q) > 0 { + uri += "?" + q.Encode() + } + req, err := http.NewRequestWithContext(c.Request().Context(), http.MethodGet, uri, nil) + if err != nil { + return nil, err + } + for _, name := range []string{"Authorization", "Cookie", "User-Agent"} { + if v := c.Request().Header.Get(name); v != "" { + req.Header.Set(name, v) + } + } + resp, err := util.DoRequestWithRetry(req) + if err != nil { + return nil, err + } + defer resp.Body.Close() + if resp.StatusCode != 200 { + return nil, echo.NewHTTPError(resp.StatusCode, "ModelScope request rejected") + } + var envelope modelscopeEnvelope + decoder := json.NewDecoder(io.LimitReader(resp.Body, 32<<20)) + if err = decoder.Decode(&envelope); err != nil { + return nil, err + } + if envelope.Code != 0 && envelope.Code != 200 { + return nil, echo.NewHTTPError(502, "ModelScope application error") + } + return &envelope, nil +} +func (m *ModelscopeService) RepositoryRevisions(c echo.Context, k repository.RepoKey) ([]ProviderRevision, error) { + envelope, err := m.repositoryJSON(c, k, "/revisions", nil) + if err != nil { + return nil, err + } + revisions := []ProviderRevision{} + for _, entry := range append(envelope.Data.RevisionMap.Branches, envelope.Data.RevisionMap.Tags...) { + if repository.Segment(entry.Revision) != nil { + return nil, fmt.Errorf("unsafe ModelScope revision") + } + commit := entry.CommitId + if commit == "" { + commit = entry.Sha + } + if commit == "" { + commit = entry.Revision + } + if repository.Segment(commit) != nil { + return nil, fmt.Errorf("unsafe ModelScope commit") + } + revisions = append(revisions, ProviderRevision{Name: entry.Revision, Commit: commit}) + } + return revisions, nil +} +func (m *ModelscopeService) RepositoryTree(c echo.Context, k repository.RepoKey, revision, path string, recursive bool) ([]RepoTreeItem, error) { + q := url.Values{"Revision": {revision}, "Recursive": {fmt.Sprint(recursive)}} + if path != "" { + q.Set("Root", path) + } + envelope, err := m.repositoryJSON(c, k, "/repo/files", q) + if err != nil { + return nil, err + } + items := []RepoTreeItem{} + for _, entry := range envelope.Data.Files { + name := entry.Path + if name == "" { + name = entry.Name + } + if err := repository.Relative(name); err != nil { + return nil, err + } + typ := "file" + if entry.Type == "tree" || entry.Type == "directory" { + typ = "directory" + } + oid := entry.Sha256 + if oid == "" { + oid = entry.Sha1 + } + item := RepoTreeItem{Type: typ, Path: name, Oid: oid, Size: entry.Size} + if entry.Sha256 != "" { + item.Lfs = &common.Lfs{Oid: entry.Sha256, Size: entry.Size} + } + items = append(items, item) + } + return items, nil +} +func (m *ModelscopeService) RepositoryMetadata(c echo.Context, k repository.RepoKey, revision string) (map[string]interface{}, error) { + if revision == "" { + revision = "master" + } + envelope, err := m.repositoryJSON(c, k, "/repo/files", url.Values{"Revision": {revision}, "Recursive": {"true"}}) + if err != nil { + return nil, err + } + if envelope.Data.LatestCommitter.Id == "" && envelope.Data.LatestCommitter.ShortId == "" { + refs, refErr := m.RepositoryRevisions(c, k) + if refErr != nil { + return nil, refErr + } + for _, ref := range refs { + if (ref.Name == revision || ref.Commit == revision) && ref.Commit != ref.Name { + envelope.Data.LatestCommitter.Id = ref.Commit + break + } + } + } + return m.cacheRepositoryMetadata(k, revision, envelope) +} +func (m *ModelscopeService) RepositoryFile(c echo.Context, k repository.RepoKey, revision, path string) error { + return m.serveFile(c, k, revision, path) +} diff --git a/internal/service/modelscope_repository_loop_test.go b/internal/service/modelscope_repository_loop_test.go new file mode 100644 index 0000000..e6e8cbe --- /dev/null +++ b/internal/service/modelscope_repository_loop_test.go @@ -0,0 +1,128 @@ +package service + +import ( + "crypto/sha256" + "fmt" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "dingospeed/internal/dao" + "dingospeed/pkg/config" + "dingospeed/pkg/hfprojection" + "dingospeed/pkg/repository" +) + +func TestModelScopeRepositoryLoop(t *testing.T) { + payload := "repository loop payload" + oid := fmt.Sprintf("%x", sha256.Sum256([]byte(payload))) + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if strings.HasSuffix(r.URL.Path, "/repo/files") { + fmt.Fprintf(w, `{"Code":200,"Data":{"LatestCommitter":{"ShortId":"repo-head"},"Files":[{"Type":"blob","Path":"one.txt","Size":%d,"Sha256":%q,"Revision":"file-old"},{"Type":"blob","Path":"two.txt","Size":9,"Sha1":"other-object","Revision":"file-other"}]}}`, len(payload), oid) + return + } + if r.URL.Query().Get("Revision") != "file-old" { + t.Errorf("file fetch not pinned: %s", r.URL) + } + fmt.Fprint(w, payload) + })) + defer upstream.Close() + msConfig(t, upstream.URL) + // Exercise the shared repository/cache lifecycle with only an endpoint + // override, as deployed configurations may omit every other provider option. + config.SysConfig.Modelscope = config.Modelscope{OfficialBaseURL: upstream.URL} + config.SysConfig.SetDefaults() + e := msEngine() + response := msRequest(e, "GET", "/file?repo=owner/model&revision=master&path=one.txt", "") + if response.Code != 200 || response.Body.String() != payload { + t.Fatalf("download: %d %s", response.Code, response.Body.String()) + } + reader := hfprojection.Reader{Root: config.SysConfig.Repos(), Namespace: repository.ModelScope} + catalog, err := reader.Catalog("models") + if err != nil || len(catalog.Repos) != 1 { + t.Fatalf("catalog: %+v %v", catalog, err) + } + revisions := catalog.Repos[0].Revisions + if len(revisions) != 1 || revisions[0].Name != "master" || revisions[0].Commit != "repo-head" { + t.Fatalf("file revisions became repository versions: %+v", revisions) + } + manifest, err := reader.Manifest("models", "owner/model", "master") + if err != nil || !manifest.ManifestComplete || len(manifest.Files) != 2 { + t.Fatalf("manifest: %+v %v", manifest, err) + } + f, body, err := reader.OpenFile("models", "owner/model", "master", "one.txt") + if err != nil { + t.Fatal(err) + } + f.Close() + _ = body + if _, _, err = reader.OpenFile("models", "owner/model", "master", "two.txt"); err == nil { + t.Fatal("missing file reported readable") + } + admin := dao.NewCacheAdminDao(NewModelscopeService().fileDao) + rows := admin.ListFiles("models", "modelscope/owner/model") + if len(rows) != 2 { + t.Fatalf("cache rows=%+v", rows) + } + var found bool + for _, row := range rows { + if row.Path == "one.txt" { + found = row.Complete && row.CachedBytes == int64(len(payload)) && len(row.Revisions) == 1 && row.Revisions[0] == "master" + } else if row.Complete || row.BlobExists { + t.Fatal("uncached file reported complete") + } + } + if !found { + t.Fatalf("cached status wrong: %+v", rows) + } + hf, err := (hfprojection.Reader{Root: config.SysConfig.Repos()}).Catalog("models") + if err != nil || len(hf.Repos) != 0 { + t.Fatalf("namespace leak: %+v %v", hf, err) + } + results, err := admin.SoftDelete([]dao.DeleteItem{{Namespace: repository.ModelScope, RepoType: "models", Repo: "owner/model", OrgRepo: "modelscope/owner/model", Path: "one.txt", Sha: oid}}) + if err != nil || len(results) != 1 || results[0].Status != "deleted" { + t.Fatalf("delete: %+v %v", results, err) + } + for _, row := range admin.ListFiles("models", "modelscope/owner/model") { + if row.Path == "one.txt" { + t.Fatal("deleted reference remains") + } + } + if len(admin.ListOrphans("models", "modelscope/owner/model")) != 1 { + t.Fatal("deleted blob missing from recycle bin") + } + // Metadata remains the upstream tree. Deleting cache does not rewrite its version. + after, err := reader.Manifest("models", "owner/model", "master") + if err != nil || after.Commit != "repo-head" || len(after.Files) != 2 { + t.Fatalf("delete changed upstream tree: %+v %v", after, err) + } + +} + +func TestModelScopeIncompleteSiblingDoesNotBlockSingleFile(t *testing.T) { + payload := "valid file" + oid := fmt.Sprintf("%x", sha256.Sum256([]byte(payload))) + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if strings.HasSuffix(r.URL.Path, "/repo/files") { + fmt.Fprintf(w, `{"Code":200,"Data":{"LatestCommitter":{"ShortId":"repo-head"},"Files":[{"Type":"blob","Path":"one.txt","Size":%d,"Sha256":%q,"Revision":"file-old"},{"Type":"blob","Path":"other.txt","Size":9}]}}`, len(payload), oid) + return + } + fmt.Fprint(w, payload) + })) + defer upstream.Close() + msConfig(t, upstream.URL) + response := msRequest(msEngine(), "GET", "/file?repo=owner/model&revision=master&path=one.txt", "") + if response.Code != 200 || response.Body.String() != payload { + t.Fatalf("unrelated sibling blocked valid file: %d %s", response.Code, response.Body.String()) + } + catalog, err := (hfprojection.Reader{Root: config.SysConfig.Repos(), Namespace: repository.ModelScope}).Catalog("models") + if err != nil { + t.Fatal(err) + } + for _, repo := range catalog.Repos { + if len(repo.Revisions) != 0 { + t.Fatalf("incomplete snapshot was published: %+v", repo) + } + } +} diff --git a/internal/service/modelscope_service.go b/internal/service/modelscope_service.go index e81eb39..bfd0685 100644 --- a/internal/service/modelscope_service.go +++ b/internal/service/modelscope_service.go @@ -1,247 +1,31 @@ package service import ( - "context" + "bytes" + "dingospeed/internal/dao" + "dingospeed/internal/data" "dingospeed/pkg/config" + "dingospeed/pkg/repository" "dingospeed/pkg/util" + "encoding/json" "fmt" - "github.com/gofrs/flock" "github.com/labstack/echo/v4" + "github.com/patrickmn/go-cache" "go.uber.org/zap" "io" "net/http" "net/url" - "os" - "path/filepath" - "strconv" "strings" - "sync" "time" ) -type ModelscopeService struct { - downloadingMap sync.Map - mapMu sync.Mutex -} +type ModelscopeService struct{ fileDao *dao.FileDao } func NewModelscopeService() *ModelscopeService { - return &ModelscopeService{ - downloadingMap: sync.Map{}, - mapMu: sync.Mutex{}, - } -} - -// CalculateBlockInfo 计算块信息：块编号、块起始偏移、块结束偏移 -func (m *ModelscopeService) CalculateBlockInfo(offset int64, chunkSize int64) (blockNum int64, blockStart int64, blockEnd int64) { - blockNum = offset / chunkSize - blockStart = blockNum * chunkSize - blockEnd = blockStart + chunkSize - 1 - return -} - -// IsBlockComplete 检查指定块是否完整 -func (m *ModelscopeService) IsBlockComplete(cachePath string, blockNum int64, chunkSize int64) (bool, error) { - fileInfo, err := os.Stat(cachePath) - if err != nil { - if os.IsNotExist(err) { - return false, nil - } - return false, fmt.Errorf("stat cache file failed: %w", err) - } - - blockStart := blockNum * chunkSize - blockEnd := blockStart + chunkSize - 1 - fileSize := fileInfo.Size() - - if fileSize <= blockStart { - return false, nil // 块为空 - } - if fileSize > blockEnd { - return true, nil // 块完整 - } - return false, nil -} - -// getCacheLock 获取缓存文件对应的锁文件 -func (m *ModelscopeService) getCacheLock(cachePath string) *flock.Flock { - lockPath := cachePath + ".lock" - return flock.New(lockPath) -} - -func (m *ModelscopeService) openCacheFile(cachePath string, c echo.Context) (*os.File, *flock.Flock, error) { - lockPath := cachePath + ".lock" - fileLock := flock.New(lockPath) - - ctx, cancel := context.WithTimeout(c.Request().Context(), 3*time.Second) - defer cancel() - locked, lockErr := fileLock.TryLockContext(ctx, time.Duration(0)) - if lockErr != nil { - zap.S().Errorf("获取缓存文件锁失败: %s, err: %v", lockPath, lockErr) - return nil, nil, c.JSON(http.StatusInternalServerError, map[string]string{ - "code": "500", - "error": "get cache file lock failed", - "msg": lockErr.Error(), - }) - } - if !locked { - zap.S().Warnf("缓存文件锁超时，无法独占写入: %s", cachePath) - return nil, nil, c.JSON(http.StatusTooManyRequests, map[string]string{ - "code": "429", - "error": "cache file is being written by another request", - "msg": "please try again later", - }) - } - - var cacheFile *os.File - var err error - if _, statErr := os.Stat(cachePath); os.IsNotExist(statErr) { - cacheFile, err = os.OpenFile(cachePath, os.O_RDWR|os.O_CREATE|os.O_EXCL, 0664) - } else { - cacheFile, err = os.OpenFile(cachePath, os.O_RDWR, 0664) - } - - if err != nil { - zap.S().Errorf("打开缓存文件失败: %s, err: %v", cachePath, err) - _ = fileLock.Unlock() - return nil, nil, c.JSON(http.StatusInternalServerError, map[string]string{ - "code": "500", - "error": "open cache file failed", - "msg": err.Error(), - }) - } - - return cacheFile, fileLock, nil -} - -func (m *ModelscopeService) copyResponseToCacheAndClient(c echo.Context, resp *http.Response, cacheFile *os.File, cachePath string, totalFileSize int64) error { - chunkSize := config.SysConfig.Modelscope.ChunkSize - if chunkSize <= 0 { - chunkSize = 1024 * 1024 * 8 - } - - currentOffset := int64(0) - contentRange := resp.Header.Get("Content-Range") - if contentRange != "" { - parts := strings.Split(contentRange, " ") - if len(parts) >= 2 { - rangeParts := strings.Split(parts[1], "-") - if len(rangeParts) >= 1 { - parsedOffset, err := strconv.ParseInt(rangeParts[0], 10, 64) - if err == nil { - currentOffset = parsedOffset - } - } - } - } - - buf := make([]byte, chunkSize) - written := int64(0) - - for { - if c.Request().Context().Err() != nil { - zap.S().Warnf("客户端断开连接，停止续传: %s", cachePath) - if cacheFile != nil { - _ = cacheFile.Sync() - closeErr := cacheFile.Close() - if closeErr != nil { - zap.S().Errorf("关闭缓存文件失败: %s, err: %v", cachePath, closeErr) - } - } - return nil - } - - currentWriteOffset := currentOffset + written - blockNum, _, blockEnd := m.CalculateBlockInfo(currentWriteOffset, chunkSize) - - isComplete, err := m.IsBlockComplete(cachePath, blockNum, chunkSize) - if err != nil { - zap.S().Errorf("检查块%d完整性失败: %v", blockNum, err) - return c.JSON(http.StatusInternalServerError, map[string]string{ - "code": "500", - "error": "check block complete failed", - "msg": err.Error(), - }) - } - - if isComplete { - skipBytes := blockEnd - currentWriteOffset + 1 - written += skipBytes - continue - } - - readSize := blockEnd - currentWriteOffset + 1 - if readSize > int64(len(buf)) { - readSize = int64(len(buf)) - } - - n, err := resp.Body.Read(buf[:readSize]) - if n > 0 { - if _, seekErr := cacheFile.Seek(currentWriteOffset, io.SeekStart); seekErr != nil { - zap.S().Errorf("定位缓存文件到偏移%d失败: %v", currentWriteOffset, seekErr) - return c.JSON(http.StatusInternalServerError, map[string]string{ - "code": "500", - "error": "seek cache file failed", - "msg": seekErr.Error(), - }) - } - - if _, writeErr := cacheFile.Write(buf[:n]); writeErr != nil { - zap.S().Errorf("写入缓存块%d失败: %s, err: %v", blockNum, cachePath, writeErr) - return c.JSON(http.StatusInternalServerError, map[string]string{ - "code": "500", - "error": "write cache block failed", - "msg": writeErr.Error(), - }) - } - if syncErr := cacheFile.Sync(); syncErr != nil { - zap.S().Warnf("缓存块%d刷盘失败: %s, err: %v", blockNum, cachePath, syncErr) - } - - if _, writeErr := c.Response().Write(buf[:n]); writeErr != nil { - if strings.Contains(writeErr.Error(), "http2: stream closed") || - strings.Contains(writeErr.Error(), "broken pipe") || - strings.Contains(writeErr.Error(), "connection reset by peer") { - zap.S().Warnf("客户端断开连接，停止返回续传数据: %s, err: %v", cachePath, writeErr) - return nil - } - zap.S().Errorf("返回续传数据失败: %s, err: %v", cachePath, writeErr) - return c.JSON(http.StatusInternalServerError, map[string]string{ - "code": "500", - "error": "write response failed", - "msg": writeErr.Error(), - }) - } - - written += int64(n) - if f, ok := c.Response().Writer.(http.Flusher); ok { - f.Flush() - } - - if written%(100*1024*1024) == 0 { - zap.S().Infof("续传进度: %dMB, 文件: %s", written/(1024*1024), cachePath) - } - } - - if err == io.EOF { - if syncErr := cacheFile.Sync(); syncErr != nil { - zap.S().Warnf("缓存文件最终刷盘失败: %s, err: %v", cachePath, syncErr) - } - zap.S().Infof("续传完成: %s, 共下载%d字节，完整文件%d字节", cachePath, written, totalFileSize) - break - } - if err != nil { - zap.S().Errorf("续传中断: %s, err: %v", cachePath, err) - return c.JSON(http.StatusBadGateway, map[string]string{ - "code": "502", - "error": "download interrupted", - "msg": err.Error(), - }) - } - } - - return nil + // The canonical handler is also built before config loading; do not initialize global data here. + base := &data.BaseData{Cache: cache.New(24*time.Hour, time.Minute)} + return &ModelscopeService{fileDao: dao.NewFileDao(dao.NewDownloaderDao(nil), base, dao.NewLockDao(base))} } - func (m *ModelscopeService) ForwardModelInfo(c echo.Context, owner, repo string, repoType string) error { apiPrefix := util.GetAPIPathPrefix(repoType) officialURL := fmt.Sprintf("%s/api/v1/%s/%s/%s?%s", @@ -301,7 +85,7 @@ func (m *ModelscopeService) ForwardRepoTreeByDatasetId(c echo.Context, datasetId // forwardRequest 通用请求转发逻辑 func (m *ModelscopeService) forwardRequest(c echo.Context, officialURL string) error { - req, err := http.NewRequest(http.MethodGet, officialURL, nil) + req, err := http.NewRequestWithContext(c.Request().Context(), http.MethodGet, officialURL, nil) if err != nil { zap.S().Errorf("构建请求失败: %v", err) return err @@ -310,6 +94,9 @@ func (m *ModelscopeService) forwardRequest(c echo.Context, officialURL string) e util.AddCLIHeaders(req.Header, c.Request().Header.Get("User-Agent")) for k, v := range c.Request().Header { + if strings.EqualFold(k, "X-Dingo-Service-Token") { + continue + } req.Header[k] = v } @@ -320,464 +107,46 @@ func (m *ModelscopeService) forwardRequest(c echo.Context, officialURL string) e } defer resp.Body.Close() - for k, v := range resp.Header { - c.Response().Header()[k] = v - } - c.Response().WriteHeader(resp.StatusCode) - - _, err = io.Copy(c.Response(), resp.Body) - if err != nil { - zap.S().Errorf("复制响应体失败: %v", err) - return err - } - return nil -} - -// HandleFileDownload 处理ModelScope文件下载请求 -func (m *ModelscopeService) HandleFileDownload(c echo.Context, owner, repo, repoType string) error { - repoId := fmt.Sprintf("%s/%s", owner, repo) - revision := c.Request().URL.Query().Get("Revision") - filePath := c.Request().URL.Query().Get("FilePath") - - if revision == "" { - revision = "master" - } - if filePath == "" { - zap.S().Error("请求参数缺失: FilePath为空") - return c.JSON(http.StatusBadRequest, map[string]string{ - "code": "400", - "error": "missing FilePath parameter", - }) - } - - if err := util.EnsureDir(filepath.Join(config.SysConfig.GetModelCacheRoot(), "dummy")); err != nil { - zap.S().Errorf("初始化模型缓存根目录失败: %v", err) - return c.JSON(http.StatusInternalServerError, map[string]string{ - "code": "500", - "error": "init model cache root dir failed", - "msg": err.Error(), - }) - } - - cachePath, cacheExists := util.GetCachePath(repoType, repoId, revision, filePath) - if cachePath == "" { - zap.S().Errorf("生成缓存路径失败: 无效的repoId格式 %s", repoId) - return c.JSON(http.StatusInternalServerError, map[string]string{ - "code": "500", - "error": "get cache path failed", - "msg": "invalid repoId format, require org/repo", - }) - } - zap.S().Infof("生成缓存路径: %s (缓存文件是否存在: %t)", cachePath, cacheExists) - - if err := util.EnsureDir(filepath.Dir(cachePath)); err != nil { - zap.S().Errorf("初始化缓存文件上级目录失败: %s, err: %v", filepath.Dir(cachePath), err) - return c.JSON(http.StatusInternalServerError, map[string]string{ - "code": "500", - "error": "init cache file parent dir failed", - "msg": err.Error(), - }) - } - - cachedSize := util.GetFileSize(cachePath) - zap.S().Infof("缓存文件状态: %s (已下载: %d字节)", cachePath, cachedSize) - if cacheExists && cachedSize == 0 { - zap.S().Warnf("缓存文件存在但大小为0，视为无效缓存: %s", cachePath) - if err := os.Remove(cachePath); err != nil { - zap.S().Errorf("删除空缓存文件失败: %s, err: %v", cachePath, err) - } - cachedSize = 0 - cacheExists = false - } - - clientStart, clientEnd, err := util.ParseRangeHeader(c.Request()) - if err != nil { - zap.S().Errorf("解析Range失败: %v", err) - return c.JSON(http.StatusBadRequest, map[string]string{ - "code": "400", - "error": "parse Range header failed", - "msg": err.Error(), - }) - } - - actualStart := clientStart - if cachedSize > 0 && actualStart < cachedSize { - actualStart = cachedSize - } - zap.S().Infof("续传起始位置: 客户端请求=%d, 缓存末尾=%d, 实际起始=%d", clientStart, cachedSize, actualStart) - - headerWritten := false - c.Response().Header().Set("Transfer-Encoding", "chunked") - c.Response().Header().Set("Content-Type", "application/octet-stream") - c.Response().Header().Set("Access-Control-Expose-Headers", "Content-Range, Content-Type") - - var cacheWritten int64 = 0 - var needRemoteDownload bool - if cachedSize > 0 && clientStart < cachedSize { - cacheWritten, headerWritten, needRemoteDownload, err = m.writeCacheData(c, cachePath, clientStart, clientEnd, cachedSize, headerWritten) - if err != nil { - zap.S().Errorf("读取缓存数据失败: %v", err) - return c.JSON(http.StatusInternalServerError, map[string]string{ - "code": "500", - "error": "write cache data failed", - "msg": err.Error(), - }) - } - - if clientEnd != -1 && (clientStart+cacheWritten-1) >= clientEnd { - zap.S().Infof("缓存数据已满足客户端Range请求，无需续传") - return nil + var body io.Reader = resp.Body + if resp.StatusCode == 200 && strings.HasSuffix(c.Request().URL.Path, "/repo/files") && strings.EqualFold(c.QueryParam("Recursive"), "true") && c.QueryParam("Root") == "" { + raw, readErr := io.ReadAll(resp.Body) + if readErr != nil { + return readErr } - - if needRemoteDownload { - m.mapMu.Lock() - defer m.mapMu.Unlock() - - isDownloading := false - loaded := m.downloadingMap.CompareAndSwap(cachePath, false, true) - if !loaded { - val, exists := m.downloadingMap.Load(cachePath) - if exists && val.(bool) { - isDownloading = true - } else { - m.downloadingMap.Store(cachePath, true) - } + var envelope modelscopeEnvelope + if json.Unmarshal(raw, &envelope) == nil && (envelope.Code == 0 || envelope.Code == 200) && (envelope.Data.LatestCommitter.Id != "" || envelope.Data.LatestCommitter.ShortId != "") { + revision := c.QueryParam("Revision") + if revision == "" { + revision = "master" } - - if !isDownloading { - zap.S().Infof("触发远程下载: 连续10次空轮询，当前偏移=%d, 文件=%s (标记为已发起远程下载)", clientStart+cacheWritten, cachePath) - actualStart = clientStart + cacheWritten - } else { - zap.S().Infof("文件%s已发起远程下载，不重复触发，继续轮询缓存", cachePath) - cacheWritten, headerWritten, needRemoteDownload, err = m.writeCacheData(c, cachePath, clientStart+cacheWritten, clientEnd, cachedSize, headerWritten) - if err != nil { - zap.S().Errorf("再次轮询缓存数据失败: %v", err) - return c.JSON(http.StatusInternalServerError, map[string]string{ - "code": "500", - "error": "retry write cache data failed", - "msg": err.Error(), - }) - } - actualStart = clientStart + cacheWritten - if needRemoteDownload { - zap.S().Warnf("文件%s已标记为下载中，但仍触发远程下载请求，终止流程", cachePath) - return nil - } + typ := "models" + if strings.Contains(c.Request().URL.Path, "/datasets/") { + typ = "datasets" } - } - } - - if err := m.downloadAndWriteRemaining(c, owner, repo, actualStart, clientEnd, cachePath, headerWritten, repoType); err != nil { - zap.S().Errorf("远程下载并写入剩余数据失败: %v", err) - m.downloadingMap.Delete(cachePath) - return c.JSON(http.StatusInternalServerError, map[string]string{ - "code": "500", - "error": "download and write remaining data failed", - "msg": err.Error(), - }) - } - - // 下载完成后清除标识 - m.downloadingMap.Delete(cachePath) - zap.S().Infof("文件%s远程下载完成，清除下载标识", cachePath) - - return nil -} - -func (m *ModelscopeService) writeCacheData(c echo.Context, cachePath string, clientStart, clientEnd, cachedSize int64, headerWritten bool) (int64, bool, bool, error) { - pollInterval := 1 * time.Second // 每秒轮询一次缓存 - maxPollCount := 30 // 最大连续空轮询次数 - - totalWritten := int64(0) - currentCacheOffset := clientStart // 当前读取到的缓存偏移量 - emptyPollCount := 0 // 连续空轮询计数器 - - for { - // 客户端断开连接则直接返回 - if c.Request().Context().Err() != nil { - zap.S().Warnf("客户端断开连接，停止返回缓存数据: %s", cachePath) - return totalWritten, headerWritten, false, nil - } - - cacheFile, err := os.Open(cachePath) - if err != nil { - zap.S().Errorf("打开缓存文件失败: %s, err: %v", cachePath, err) - return totalWritten, headerWritten, false, fmt.Errorf("open cache file failed: %w", err) - } - - fileInfo, err := cacheFile.Stat() - if err != nil { - _ = cacheFile.Close() - zap.S().Errorf("获取缓存文件信息失败: %s, err: %v", cachePath, err) - return totalWritten, headerWritten, false, fmt.Errorf("stat cache file failed: %w", err) - } - latestCacheSize := fileInfo.Size() - - cacheEnd := latestCacheSize - 1 - if clientEnd != -1 && clientEnd < cacheEnd { - cacheEnd = clientEnd - } - - if currentCacheOffset > cacheEnd { - emptyPollCount++ - _ = cacheFile.Close() - - isDownloading := false - if downloadingVal, exists := m.downloadingMap.Load(cachePath); exists { - isDownloading = downloadingVal.(bool) + key := repository.RepoKey{Namespace: repository.ModelScope, RepoType: typ, Repo: c.Param("org") + "/" + c.Param("repo")} + if err = repository.Register(config.SysConfig.Repos(), repository.Remote(key)); err != nil { + return err } - - zap.S().Debugf("缓存暂未更新，第%d次空轮询（最大%d次）: 当前偏移=%d, 缓存末尾=%d, 文件=%s, 有客户端下载中: %t", - emptyPollCount, maxPollCount, currentCacheOffset, cacheEnd, cachePath, isDownloading) - - if emptyPollCount >= maxPollCount && !isDownloading { - zap.S().Warnf("连续%d次空轮询且无客户端下载中，触发当前客户端发起远程下载: %s", maxPollCount, cachePath) - return totalWritten, headerWritten, true, nil + if _, err = m.cacheRepositoryMetadata(key, revision, &envelope); err != nil { + return err } - - zap.S().Debugf("继续轮询缓存: 文件=%s, 原因: %s", cachePath, - func() string { - if isDownloading { - return "有客户端正在远程下载" - } - return fmt.Sprintf("空轮询次数未到上限（当前%d次，最大%d次）", emptyPollCount, maxPollCount) - }()) - time.Sleep(pollInterval) - continue } - - emptyPollCount = 0 - if !headerWritten { - contentRange := fmt.Sprintf("bytes %d-%d/%d", clientStart, cacheEnd, latestCacheSize) - c.Response().Header().Set("Content-Range", contentRange) - c.Response().WriteHeader(http.StatusPartialContent) - headerWritten = true - zap.S().Infof("设置缓存响应头: Content-Range=%s", contentRange) - } - - if _, err := cacheFile.Seek(currentCacheOffset, io.SeekStart); err != nil { - _ = cacheFile.Close() - zap.S().Errorf("定位缓存文件失败: %s, err: %v", cachePath, err) - return totalWritten, headerWritten, false, fmt.Errorf("seek cache file failed: %w", err) - } - - buf := make([]byte, config.SysConfig.Modelscope.ChunkSize) - readableSize := cacheEnd - currentCacheOffset + 1 - writtenInRound := int64(0) - - for writtenInRound < readableSize { - if c.Request().Context().Err() != nil { - _ = cacheFile.Close() - zap.S().Warnf("客户端断开连接，停止返回缓存数据: %s", cachePath) - return totalWritten, headerWritten, false, nil - } - - readSize := readableSize - writtenInRound - if readSize > int64(len(buf)) { - readSize = int64(len(buf)) - } - - n, err := cacheFile.Read(buf[:readSize]) - if n > 0 { - if _, writeErr := c.Response().Write(buf[:n]); writeErr != nil { - _ = cacheFile.Close() - zap.S().Errorf("返回缓存数据失败: %s, err: %v", cachePath, writeErr) - return totalWritten, headerWritten, false, fmt.Errorf("write cache data to response failed: %w", writeErr) - } - - writtenInRound += int64(n) - totalWritten += int64(n) - currentCacheOffset += int64(n) - - if f, ok := c.Response().Writer.(http.Flusher); ok { - f.Flush() - } - } - - if err == io.EOF { - break - } - if err != nil { - _ = cacheFile.Close() - zap.S().Errorf("读取缓存数据失败: %s, err: %v", cachePath, err) - return totalWritten, headerWritten, false, fmt.Errorf("read cache file failed: %w", err) - } - } - _ = cacheFile.Close() - zap.S().Debugf("本轮读取缓存完成: 文件=%s, 读取字节=%d, 累计读取=%d, 当前偏移=%d", - cachePath, writtenInRound, totalWritten, currentCacheOffset) - - if clientEnd == -1 || currentCacheOffset > clientEnd { - zap.S().Infof("缓存数据已满足客户端Range请求: 文件=%s, 累计读取=%d字节", cachePath, totalWritten) - return totalWritten, headerWritten, false, nil - } - - time.Sleep(pollInterval) + body = bytes.NewReader(raw) } -} - -// downloadAndWriteRemaining 下载剩余部分并写入响应+缓存 -func (m *ModelscopeService) downloadAndWriteRemaining(c echo.Context, owner, repo string, actualStart, clientEnd int64, cachePath string, headerWritten bool, repoType string) error { - req, err := m.buildDownloadRequest(c, owner, repo, actualStart, clientEnd, repoType) - if err != nil { - return err - } - - resp, err := m.sendDownloadRequest(req, c, owner, repo, repoType) - if err != nil { - return err - } - defer resp.Body.Close() - - totalFileSize := m.parseTotalFileSize(resp) - - cacheFile, fileLock, err := m.openCacheFile(cachePath, c) - if err != nil { - return err - } - defer func() { - if cacheFile != nil { - if closeErr := cacheFile.Close(); closeErr != nil { - zap.S().Errorf("关闭缓存文件失败: %s, err: %v", cachePath, closeErr) - } - } - if fileLock != nil { - if unlockErr := fileLock.Unlock(); unlockErr != nil { - zap.S().Errorf("解锁缓存文件锁失败: %s, err: %v", cachePath+".lock", unlockErr) + for k, v := range resp.Header { + if k == "Link" || k == "Location" { + for i := range v { + v[i] = strings.ReplaceAll(v[i], config.SysConfig.Modelscope.OfficialBaseURL, providerLinkBase(c)) } } - }() - - headerWritten = m.writeResponseHeader(c, resp, headerWritten) + c.Response().Header()[k] = v + } + c.Response().WriteHeader(resp.StatusCode) - err = m.copyResponseToCacheAndClient(c, resp, cacheFile, cachePath, totalFileSize) + _, err = io.Copy(c.Response(), body) if err != nil { + zap.S().Errorf("复制响应体失败: %v", err) return err } - return nil } - -// buildDownloadRequest 构建下载剩余部分的HTTP请求 -func (m *ModelscopeService) buildDownloadRequest(c echo.Context, owner, repo string, actualStart, clientEnd int64, repoType string) (*http.Request, error) { - apiPrefix := util.GetAPIPathPrefix(repoType) - query := c.Request().URL.RawQuery - officialURL := fmt.Sprintf("%s/api/v1/%s/%s/%s/repo?%s", - config.SysConfig.Modelscope.OfficialBaseURL, - apiPrefix, - url.PathEscape(owner), - url.PathEscape(repo), - query, - ) - zap.S().Infof("请求ModelScope官方地址: %s", officialURL) - - req, err := http.NewRequest(http.MethodGet, officialURL, nil) - if err != nil { - zap.S().Errorf("构建请求失败: %v", err) - return nil, c.JSON(http.StatusInternalServerError, map[string]string{ - "code": "500", - "error": "build request failed", - "msg": err.Error(), - }) - } - - skipHeaders := map[string]bool{ - "Range": true, - "User-Agent": true, - "Host": true, - } - for k, v := range c.Request().Header { - key := strings.ToLower(k) - if !skipHeaders[key] { - req.Header[k] = v - } - } - - util.AddCLIHeaders(req.Header, c.Request().Header.Get("User-Agent")) - - rangeHeader := fmt.Sprintf("bytes=%d-", actualStart) - if clientEnd != -1 { - rangeHeader = fmt.Sprintf("bytes=%d-%d", actualStart, clientEnd) - } - req.Header.Set("Range", rangeHeader) - zap.S().Infof("向官方请求剩余部分: %s", rangeHeader) - - return req, nil -} - -// sendDownloadRequest 发送下载请求并校验响应状态码 -func (m *ModelscopeService) sendDownloadRequest(req *http.Request, c echo.Context, owner, repo, repoType string) (*http.Response, error) { - resp, err := util.DoRequestWithRetry(req) - if err != nil { - zap.S().Errorf("下载剩余部分失败: %v", err) - return nil, c.JSON(http.StatusBadGateway, map[string]string{ - "code": "502", - "error": "download remaining failed", - "msg": err.Error(), - }) - } - - if resp.StatusCode < 200 || resp.StatusCode >= 300 { - officialURL := req.URL.String() - zap.S().Errorf("ModelScope返回错误状态码: %d, URL: %s", resp.StatusCode, officialURL) - errorMsg := fmt.Sprintf("modelscope server return status code: %d", resp.StatusCode) - - var respJSON error - switch resp.StatusCode { - case http.StatusNotFound: - respJSON = c.JSON(http.StatusNotFound, map[string]string{ - "code": "404", - "error": "resource not found", - "msg": "model or file does not exist on ModelScope", - }) - case http.StatusForbidden: - respJSON = c.JSON(http.StatusForbidden, map[string]string{ - "code": "403", - "error": "forbidden", - "msg": "no permission to access the resource", - }) - default: - respJSON = c.JSON(http.StatusBadGateway, map[string]string{ - "code": "502", - "error": "modelscope server error", - "msg": errorMsg, - }) - } - return nil, respJSON - } - - return resp, nil -} - -// parseTotalFileSize 从响应头解析Content-Range获取总文件大小 -func (m *ModelscopeService) parseTotalFileSize(resp *http.Response) int64 { - totalFileSize := int64(-1) - contentRange := resp.Header.Get("Content-Range") - if contentRange != "" { - parts := strings.Split(contentRange, "/") - if len(parts) == 2 { - parsedSize, err := strconv.ParseInt(parts[1], 10, 64) - if err == nil { - totalFileSize = parsedSize - } else { - zap.S().Warnf("解析Content-Range失败: %s, err: %v", contentRange, err) - } - } - } - return totalFileSize -} - -// writeResponseHeader 写入续传响应头 -func (m *ModelscopeService) writeResponseHeader(c echo.Context, resp *http.Response, headerWritten bool) bool { - if !headerWritten { - if resp.StatusCode == http.StatusPartialContent { - c.Response().Header().Set("Content-Range", resp.Header.Get("Content-Range")) - c.Response().WriteHeader(http.StatusPartialContent) - } else { - c.Response().WriteHeader(http.StatusOK) - } - headerWritten = true - zap.S().Infof("设置续传响应头，状态码: %d", resp.StatusCode) - } - return headerWritten -} diff --git a/internal/service/publication_sync.go b/internal/service/publication_sync.go new file mode 100644 index 0000000..cebce5b --- /dev/null +++ b/internal/service/publication_sync.go @@ -0,0 +1,168 @@ +package service + +import ( + "encoding/json" + "fmt" + "os" + "path/filepath" + "sort" + "strings" + "sync" + "time" + + "dingospeed/internal/dao" + "dingospeed/pkg/config" + "dingospeed/pkg/repository" + "github.com/google/uuid" +) + +type UploadInventoryItem struct { + Namespace string `json:"namespace"` + RepoType string `json:"repoType"` + Repo string `json:"repo"` + Path string `json:"path"` + SHA256 string `json:"sha256"` + Size int64 `json:"size"` +} + +type UploadInventorySnapshot struct { + Version int `json:"version"` + InstanceID string `json:"instanceId"` + Epoch string `json:"epoch"` + EpochStartedAt time.Time `json:"epochStartedAt"` + Sequence uint64 `json:"sequence"` + GeneratedAt time.Time `json:"generatedAt"` + Complete bool `json:"complete"` + Error string `json:"error,omitempty"` + Items []UploadInventoryItem `json:"items"` + ReportToken string `json:"reportToken,omitempty"` +} + +var inventorySnapshotMu sync.Mutex + +func uploadInventoryPath() string { + return filepath.Join(config.SysConfig.Repos(), ".upload-inventory", "snapshot.json") +} + +func LoadUploadInventorySnapshot() (*UploadInventorySnapshot, error) { + inventorySnapshotMu.Lock() + defer inventorySnapshotMu.Unlock() + b, err := os.ReadFile(uploadInventoryPath()) + if err != nil { + return nil, err + } + var snap UploadInventorySnapshot + if err = json.Unmarshal(b, &snap); err != nil { + return nil, err + } + return &snap, nil +} + +func (s *SchedulerService) buildUploadInventory() (*UploadInventorySnapshot, error) { + previous, loadErr := LoadUploadInventorySnapshot() + now := time.Now().UTC() + snap := &UploadInventorySnapshot{Version: 1, InstanceID: config.SysConfig.Registration().NodeID, GeneratedAt: now, Complete: true, Items: make([]UploadInventoryItem, 0)} + snap.ReportToken = uuid.NewString() + if loadErr == nil && previous.Epoch != "" && previous.InstanceID == snap.InstanceID { + snap.Epoch, snap.EpochStartedAt, snap.Sequence = previous.Epoch, previous.EpochStartedAt, previous.Sequence+1 + } else { + snap.Epoch, snap.EpochStartedAt, snap.Sequence = uuid.NewString(), now, 1 + } + if snap.InstanceID == "" { + return nil, fmt.Errorf("scheduler discovery instanceId is empty") + } + // Only hosted descriptors become inventory items; ListHosted walks just the + // hosted roots instead of the whole api/ tree, which is mostly remote caches. + descriptors, err := repository.ListHosted(config.SysConfig.Repos()) + if err != nil { + snap.Complete = false + snap.Error = err.Error() + return snap, nil + } + return s.scanUploadDescriptors(snap, descriptors) +} + +func (s *SchedulerService) scanUploadDescriptors(snap *UploadInventorySnapshot, descriptors []repository.Descriptor) (*UploadInventorySnapshot, error) { + seen := make(map[string]struct{}) + errorsFound := make([]string, 0) + for _, d := range descriptors { + if d.Source != "hosted" { + continue + } + entries, readErr := os.ReadDir(filepath.Join(d.APIRoot(config.SysConfig.Repos()), "revision")) + if os.IsNotExist(readErr) { + continue + } + if readErr != nil { + errorsFound = append(errorsFound, d.ID()+": "+readErr.Error()) + continue + } + for _, entry := range entries { + if !entry.IsDir() { + continue + } + commit, files, readErr := dao.ReadInventorySnapshot(d.RepoKey, entry.Name()) + if readErr != nil { + errorsFound = append(errorsFound, d.ID()+"/"+entry.Name()+": "+readErr.Error()) + continue + } + if entry.Name() == commit { + continue + } // immutable snapshot directory, not a live local revision + manifest := make([]dao.LocalManifestFile, len(files)) + for i, file := range files { + manifest[i] = dao.LocalManifestFile{Path: file.Path, Sha256: file.Sha256, Size: file.Size} + } + if verifyErr := dao.VerifyPublishedFiles(d.RepoType, d.ID(), commit, manifest); verifyErr != nil { + errorsFound = append(errorsFound, d.ID()+"/"+entry.Name()+": "+verifyErr.Error()) + continue + } + for _, file := range files { + key := strings.Join([]string{d.Namespace, d.RepoType, d.Repo, file.Path, file.Sha256}, "\x00") + if _, ok := seen[key]; ok { + continue + } + seen[key] = struct{}{} + snap.Items = append(snap.Items, UploadInventoryItem{Namespace: d.Namespace, RepoType: d.RepoType, Repo: d.Repo, Path: file.Path, SHA256: file.Sha256, Size: file.Size}) + } + } + } + sort.Slice(snap.Items, func(i, j int) bool { + a, b := snap.Items[i], snap.Items[j] + return strings.Join([]string{a.Namespace, a.RepoType, a.Repo, a.Path, a.SHA256}, "\x00") < strings.Join([]string{b.Namespace, b.RepoType, b.Repo, b.Path, b.SHA256}, "\x00") + }) + if len(errorsFound) > 0 { + snap.Complete = false + snap.Error = strings.Join(errorsFound, "; ") + } + return snap, nil +} + +func persistUploadInventory(snap *UploadInventorySnapshot) error { + inventorySnapshotMu.Lock() + defer inventorySnapshotMu.Unlock() + dest := uploadInventoryPath() + if err := os.MkdirAll(filepath.Dir(dest), 0o755); err != nil { + return err + } + b, err := json.Marshal(snap) + if err != nil { + return err + } + tmp, err := os.CreateTemp(filepath.Dir(dest), ".snapshot-*.tmp") + if err != nil { + return err + } + name := tmp.Name() + defer os.Remove(name) + if _, err = tmp.Write(b); err == nil { + err = tmp.Sync() + } + if closeErr := tmp.Close(); err == nil { + err = closeErr + } + if err != nil { + return err + } + return os.Rename(name, dest) +} diff --git a/internal/service/registration_settings.go b/internal/service/registration_settings.go new file mode 100644 index 0000000..24713e3 --- /dev/null +++ b/internal/service/registration_settings.go @@ -0,0 +1,23 @@ +package service + +import ( + "dingospeed/pkg/common" + "dingospeed/pkg/config" + "errors" + "sync" +) + +var registrationAdmission sync.Mutex +var registrationCachePool *common.Pool // wired before servers start + +func SaveSchedulerRegistration(r config.Registration) error { + registrationAdmission.Lock() + defer registrationAdmission.Unlock() + cacheStateMu.Lock() + hasExecutions := len(cacheExecutions) > 0 + cacheStateMu.Unlock() + if config.SysConfig.Registration() != r && (hasExecutions || (registrationCachePool != nil && registrationCachePool.ActiveCount() > 0)) { + return errors.New("pause or cancel active cache jobs before changing scheduler registration") + } + return config.SysConfig.SaveRegistration(r) +} diff --git a/internal/service/scheduler_service.go b/internal/service/scheduler_service.go index 59c8bbd..f00369d 100644 --- a/internal/service/scheduler_service.go +++ b/internal/service/scheduler_service.go @@ -17,25 +17,64 @@ type SchedulerService struct { Client manager.ManagerClient Ctx context.Context schedulerDao *dao.SchedulerDao + metaService *MetaService } -func NewSchedulerService(schedulerDao *dao.SchedulerDao) *SchedulerService { +func NewSchedulerService(schedulerDao *dao.SchedulerDao, metaService *MetaService) *SchedulerService { return &SchedulerService{ schedulerDao: schedulerDao, + metaService: metaService, } } +func (s *SchedulerService) BindClient(client manager.ManagerClient) { + s.Client = client + s.schedulerDao.Client = client +} func (s *SchedulerService) Register() { - s.schedulerDao.Client = s.Client - response, err := s.schedulerDao.Register() - if err != nil { - runModeChange(consts.SchedulerModeStandalone) - return - } - zap.S().Infof("enter cluster mode......") - config.SysConfig.Id = response.Id - go s.Heartbeat() go s.ReportFileProcess() + go s.ReconcilePublications() + var previous config.Registration + registered := false + ticker := time.NewTicker(2 * time.Second) + defer ticker.Stop() + for { + r := config.SysConfig.Registration() + if r != previous { + registered = false + previous = r + } + if !r.Enabled { + registered = false + runModeChange(consts.SchedulerModeStandalone) + config.SysConfig.SetRegistrationStatus(r, "disabled", nil) + } else { + var err error + if !registered { + var response *manager.RegisterResponse + response, err = s.schedulerDao.Register() + if err == nil { + config.SysConfig.SetSchedulerID(response.Id) + registered = true + wakeInventoryConnection() + } + } else { + err = s.schedulerDao.Heartbeat() + } + if err != nil { + registered = false + runModeChange(consts.SchedulerModeStandalone) + config.SysConfig.SetRegistrationStatus(r, "error", err) + } else if config.SysConfig.Registration() == r { + config.SysConfig.SetRegistrationStatus(r, "connected", nil) + } + } + select { + case <-s.Ctx.Done(): + return + case <-ticker.C: + } + } } func (s *SchedulerService) Heartbeat() { @@ -46,7 +85,7 @@ func (s *SchedulerService) Heartbeat() { case <-ticker.C: err := s.schedulerDao.Heartbeat() if err != nil { - zap.S().Errorf("speed:%s connect err.%v", config.SysConfig.Scheduler.Discovery.InstanceId, err) + zap.S().Errorf("speed:%s connect err.%v", config.SysConfig.Registration().NodeID, err) runModeChange(consts.SchedulerModeStandalone) break } diff --git a/internal/service/sys_service.go b/internal/service/sys_service.go index db92b87..80b3274 100644 --- a/internal/service/sys_service.go +++ b/internal/service/sys_service.go @@ -14,6 +14,7 @@ import ( "dingospeed/internal/dao" "dingospeed/pkg/config" "dingospeed/pkg/proto/manager" + "dingospeed/pkg/repository" "dingospeed/pkg/util" "github.com/shirou/gopsutil/mem" @@ -106,6 +107,11 @@ func (s *SysService) checkDiskUsage() { zap.S().Infof("Cache size exceeded! Limit: %s, Current: %s.\n", limitSizeH, currentSizeH) zap.S().Infof("Cleaning...") + registered, err := repository.List(baseRepoPath) + if err != nil { + zap.S().Errorf("Skip disk clean: repository registry unavailable: %v", err) + return + } filesPath := filepath.Join(config.SysConfig.Repos(), "files") var allFiles []util.FileWithPath switch config.SysConfig.CacheCleanStrategy() { @@ -132,20 +138,21 @@ func (s *SysService) checkDiskUsage() { return } - instanceID := config.SysConfig.Scheduler.Discovery.InstanceId + instanceID := config.SysConfig.Registration().NodeID for _, file := range allFiles { if currentSize < limitSize { break } filePath := file.Path fileSize := file.Info.Size() - if isProtectedLocalUploadCacheFile(baseRepoPath, filePath) { + descriptor, relative, locateErr := repository.LocateRegistered(baseRepoPath, filePath, registered) + if locateErr != nil || descriptor.Persistent || descriptor.Source != "remote" || repository.SafePath(baseRepoPath, filepath.Dir(filePath)) != nil { zap.S().Infof("Skip local upload cache file during disk clean: %s", filePath) continue } if s.Client != nil { - s.deleteRecordByFilePath(baseRepoPath, filePath, instanceID) + s.deleteRecordForRepository(descriptor, relative, instanceID) } err := os.Remove(filePath) @@ -167,58 +174,27 @@ func (s *SysService) checkDiskUsage() { } func isProtectedLocalUploadCacheFile(baseRepoPath, filePath string) bool { - relPath, err := filepath.Rel(baseRepoPath, filePath) - if err != nil { - return false - } - parts := strings.Split(relPath, string(filepath.Separator)) - if len(parts) < 4 || parts[0] != "files" { - return false - } - if parts[1] != "models" && parts[1] != "datasets" { - return false - } - namespace := config.SysConfig.Upload.Namespace - if namespace == "" { - namespace = "dingo-local" - } - return parts[2] == namespace + d, _, err := repository.Locate(baseRepoPath, filePath) + return err != nil || d.Persistent || d.Source != "remote" } - func (s *SysService) deleteRecordByFilePath(baseRepoPath, filePath, instanceID string) { - relPath, err := filepath.Rel(baseRepoPath, filePath) + d, rel, err := repository.Locate(baseRepoPath, filePath) if err != nil { - zap.S().Errorf("Failed to get relative path for %s: %v", filePath, err) return } + s.deleteRecordForRepository(d, rel, instanceID) +} - parts := strings.Split(relPath, string(filepath.Separator)) - req := &manager.DeleteByEtagsAndFieldsRequest{ - InstanceID: instanceID, // 设置实例ID - } - - if len(parts) >= 6 && parts[0] == "files" && (parts[1] == "datasets" || parts[1] == "models") { - req.Datatype = parts[1] - req.Org = parts[2] - req.Repo = parts[3] - - if parts[4] == "blobs" { - req.Etag = parts[5] - zap.S().Debugf("Deleting record by etag: %s (path type: %s, org: %s, repo: %s) for file %s", - req.Etag, req.Datatype, req.Org, req.Repo, filePath) - } else if parts[4] == "resolve" { - req.Name = parts[len(parts)-1] - zap.S().Debugf("Deleting record by fields - datatype: %s, org: %s, repo: %s, name: %s", - req.Datatype, req.Org, req.Repo, req.Name) - } else { - zap.S().Warnf("Unrecognized subpath: %s in path %s", parts[4], filePath) - return - } +func (s *SysService) deleteRecordForRepository(d repository.Descriptor, rel, instanceID string) { + req := &manager.DeleteByEtagsAndFieldsRequest{InstanceID: instanceID, Datatype: d.RepoType, Org: d.Namespace, Repo: d.Repo} + p := strings.SplitN(rel, "/", 3) + if len(p) == 2 && p[0] == "blobs" { + req.Etag = p[1] + } else if len(p) == 3 && p[0] == "resolve" { + req.Name = p[2] } else { - zap.S().Warnf("Unrecognized file path structure: %s, cannot determine delete parameters", filePath) return } - s.schedulerDao.DeleteByEtagsAndFields(req) } diff --git a/internal/service/task/mount_cache_task.go b/internal/service/task/mount_cache_task.go index 27b400b..6dd7538 100644 --- a/internal/service/task/mount_cache_task.go +++ b/internal/service/task/mount_cache_task.go @@ -9,8 +9,11 @@ import ( "strings" "time" + "dingospeed/internal/dao" "dingospeed/pkg/config" "dingospeed/pkg/consts" + "dingospeed/pkg/hfprojection" + "dingospeed/pkg/repository" "go.uber.org/zap" ) @@ -21,7 +24,20 @@ type MountCacheTask struct { } func (m *MountCacheTask) DoTask() { - orgRepo := fmt.Sprintf("%s/%s", m.Job.Org, m.Job.Repo) + orgRepo := fmt.Sprintf("%s/%s", m.Job.Namespace, m.Job.Repo) + key := repository.RepoKey{Namespace: m.Job.Namespace, RepoType: m.Job.Datatype, Repo: m.Job.Repo} + if key.Namespace != repository.HuggingFace { + revision := "main" + if key.Namespace == repository.ModelScope { + revision = "master" + } + if err := m.mountViaRepositoryAPI(key, revision); err != nil { + m.SchedulerDao.ExecUpdateRepositoryMountStatus(m.TaskNo, m.StopStatus(), err.Error()) + } else { + m.SchedulerDao.ExecUpdateRepositoryMountStatus(m.TaskNo, consts.RunningStatusJobComplete, "") + } + return + } var repoType string if m.Job.Datatype == consts.RepoTypeModel.Value() { repoType = "model" @@ -33,7 +49,7 @@ func (m *MountCacheTask) DoTask() { } modelDirName := filepath.Base(orgRepo) mountDir := config.SysConfig.Cache.MountModelDir - localModelDir := filepath.Join(mountDir, m.Job.Datatype, orgRepo) + localModelDir := filepath.Join(mountDir, m.Job.Datatype, filepath.FromSlash(key.Repo)) logDir := filepath.Join(config.SysConfig.Server.Repos, "mount_download_logs") if err := os.MkdirAll(logDir, 0755); err != nil { @@ -48,15 +64,20 @@ func (m *MountCacheTask) DoTask() { return } defer logF.Close() - hfEndpoint := fmt.Sprintf("http://%s:%d", config.SysConfig.Server.Host, config.SysConfig.Server.Port) + hfEndpoint := fmt.Sprintf("http://%s:%d/huggingface", config.SysConfig.Server.Host, config.SysConfig.Server.Port) + upstream, upErr := dao.UpstreamRepo(m.Job.Datatype, orgRepo) + if upErr != nil { + m.SchedulerDao.ExecUpdateRepositoryMountStatus(m.TaskNo, m.StopStatus(), upErr.Error()) + return + } token := getToken(m.Authorization) var cmd *exec.Cmd if token != "" { cmd = exec.Command("hf", "download", "--repo-type", - repoType, orgRepo, "--local-dir", localModelDir, "--token", token) + repoType, upstream, "--local-dir", localModelDir, "--token", token) } else { cmd = exec.Command("hf", "download", "--repo-type", - repoType, orgRepo, "--local-dir", localModelDir) + repoType, upstream, "--local-dir", localModelDir) } cmd.Env = append(os.Environ(), fmt.Sprintf("HF_ENDPOINT=%s", hfEndpoint)) cmd.Stdout = logF @@ -82,9 +103,10 @@ func (m *MountCacheTask) DoTask() { errMsg = strings.Join(lines, "\n") } } - m.SchedulerDao.ExecUpdateRepositoryMountStatus(m.TaskNo, m.RunningStatus, errMsg) + m.SchedulerDao.ExecUpdateRepositoryMountStatus(m.TaskNo, m.StopStatus(), errMsg) } else { zap.S().Infof("command success.%d", m.TaskNo) + hfprojection.DefaultIndex.Invalidate(config.SysConfig.Repos(), m.Job.Datatype) m.SchedulerDao.ExecUpdateRepositoryMountStatus(m.TaskNo, consts.RunningStatusJobComplete, "") } } diff --git a/internal/service/task/namespace_mount.go b/internal/service/task/namespace_mount.go new file mode 100644 index 0000000..3d692ed --- /dev/null +++ b/internal/service/task/namespace_mount.go @@ -0,0 +1,213 @@ +package task + +import ( + "context" + "encoding/json" + "fmt" + "io" + "net" + "net/http" + "net/url" + "os" + "path/filepath" + "strconv" + + "dingospeed/internal/dao" + "dingospeed/pkg/common" + "dingospeed/pkg/config" + "dingospeed/pkg/repository" +) + +// Repository API requests only target this configured service without redirects. +func repositoryRequest(ctx context.Context, k repository.RepoKey, authorization, operation string, q url.Values) (*http.Response, error) { + if err := k.Validate(); err != nil { + return nil, err + } + host := config.SysConfig.Server.Host + if host == "" || host == "0.0.0.0" || host == "::" { + host = "127.0.0.1" + } + base := "http://" + net.JoinHostPort(host, strconv.Itoa(config.SysConfig.Server.Port)) + "/api/repositories/" + url.PathEscape(k.RepoType) + "/" + url.PathEscape(k.Namespace) + "/" + req, err := http.NewRequestWithContext(ctx, http.MethodGet, base+operation+"?"+q.Encode(), nil) + if err != nil { + return nil, err + } + if k.Namespace == repository.ModelScope && operation == "file" { + req.Header.Set("X-Dingo-Ensure-Cache", "1") + } + if authorization != "" { + req.Header.Set("Authorization", authorization) + } + client := &http.Client{CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }} + transfer := common.NewLocalTransfer(common.CacheTrace(ctx)) + req.Header.Set(common.LocalTransferHeader, transfer.ID()) + resp, err := client.Do(req) + if err != nil { + transfer.Wait() + } else { + resp.Body = &joinedLocalBody{ReadCloser: resp.Body, transfer: transfer} + } + if err != nil { + return nil, err + } + if resp.StatusCode == http.StatusPartialContent && operation == "file" { + var start, end, total int64 + if n, err := fmt.Sscanf(resp.Header.Get("Content-Range"), "bytes %d-%d/%d", &start, &end, &total); err == nil && n == 3 && start == 0 && total > 0 && end == total-1 { + resp.Body = &completeRepositoryBody{ReadCloser: resp.Body, remaining: total} + return resp, nil + } + } + if resp.StatusCode != http.StatusOK { + resp.Body.Close() + return nil, fmt.Errorf("repository %s returned %d", operation, resp.StatusCode) + } + return resp, nil +} + +// A raw provider cache may answer a complete-file range with 206. Verify the +// declared byte count before a mount publishes its temporary file. +type completeRepositoryBody struct { + io.ReadCloser + remaining int64 +} + +func (b *completeRepositoryBody) Read(p []byte) (int, error) { + n, err := b.ReadCloser.Read(p) + b.remaining -= int64(n) + if b.remaining < 0 || (err == io.EOF && b.remaining != 0) { + return n, io.ErrUnexpectedEOF + } + return n, err +} + +// ReadRepositoryMetadata shares the authenticated source dispatch used by mounts +// and preheat. The provider adapter retains its native cache representation. +func ReadRepositoryMetadata(ctx context.Context, k repository.RepoKey, revision, authorization string) (*dao.CommitHfSha, error) { + resp, err := repositoryRequest(ctx, k, authorization, "metadata", url.Values{"repo": {k.Repo}, "revision": {revision}}) + if err != nil { + return nil, err + } + defer resp.Body.Close() + var meta dao.CommitHfSha + if err = json.NewDecoder(io.LimitReader(resp.Body, 32<<20)).Decode(&meta); err != nil { + return nil, err + } + if err = repository.Segment(meta.Sha); err != nil { + return nil, err + } + for _, file := range meta.Siblings { + if err := repository.Relative(file.Rfilename); err != nil { + return nil, err + } + } + return &meta, nil +} + +// Hosted and ModelScope repositories cannot be passed as HF repo_id values. +func (m *MountCacheTask) mountViaRepositoryAPI(k repository.RepoKey, revision string) error { + meta, err := ReadRepositoryMetadata(m.Ctx, k, revision, m.Authorization) + if err != nil { + return err + } + q := url.Values{"repo": {k.Repo}, "revision": {meta.Sha}} + mountRoot := config.SysConfig.Cache.MountModelDir + for _, file := range meta.Siblings { + if err := repository.Relative(file.Rfilename); err != nil { + return err + } + target := filepath.Join(mountRoot, k.RepoType, k.Namespace, filepath.FromSlash(k.Repo), filepath.FromSlash(file.Rfilename)) + if err := repository.SafePath(mountRoot, target); err != nil { + return err + } + if err := os.MkdirAll(filepath.Dir(target), 0755); err != nil { + return err + } + q.Set("revision", meta.Sha) + q.Set("path", file.Rfilename) + resp, err := repositoryRequest(m.Ctx, k, m.Authorization, "file", q) + if err != nil { + return err + } + err = writeMountedFile(target, resp.Body) + resp.Body.Close() + if err != nil { + return err + } + } + return nil +} + +func writeMountedFile(target string, body io.Reader) error { + f, err := os.CreateTemp(filepath.Dir(target), ".mount-*") + if err != nil { + return err + } + defer os.Remove(f.Name()) + if _, err = io.Copy(f, body); err != nil { + f.Close() + return err + } + if err = f.Sync(); err != nil { + f.Close() + return err + } + if err = f.Close(); err != nil { + return err + } + return os.Rename(f.Name(), target) +} + +func (p *PreheatCacheTask) preheatViaRepositoryAPI(k repository.RepoKey) error { + q := url.Values{"repo": {k.Repo}, "revision": {p.Sha.Sha}} + for _, file := range p.Sha.Siblings { + if err := repository.Relative(file.Rfilename); err != nil { + return err + } + q.Set("path", file.Rfilename) + resp, err := repositoryRequest(p.Ctx, k, p.Authorization, "file", q) + if err != nil { + return err + } + // Count bytes as the provider response is consumed. Updating only after + // io.Copy returned made one large ModelScope file look stuck for its + // entire download even though data was arriving normally. + _, err = io.Copy(io.Discard, &cacheProgressReader{reader: resp.Body, add: p.stockLen.Add}) + resp.Body.Close() + if ctxErr := p.Ctx.Err(); ctxErr != nil { + return ctxErr + } + if k.Namespace == repository.ModelScope && resp.Trailer.Get("X-Dingo-Cache-Complete") != "true" { + if reason := resp.Trailer.Get("X-Dingo-Cache-Error"); reason != "" { + return fmt.Errorf("ModelScope file cache incomplete: %s: %s", file.Rfilename, reason) + } + if err != nil { + return fmt.Errorf("ModelScope file transfer failed: %s: %w", file.Rfilename, err) + } + return fmt.Errorf("ModelScope file transfer did not complete its cache: %s", file.Rfilename) + } + if err != nil { + return err + } + } + return nil +} + +type cacheProgressReader struct { + reader io.Reader + add func(uint64) uint64 +} + +func (r *cacheProgressReader) Read(buffer []byte) (int, error) { + n, err := r.reader.Read(buffer) + if n > 0 { + r.add(uint64(n)) + } + return n, err +} + +type joinedLocalBody struct { + io.ReadCloser + transfer *common.LocalTransfer +} + +func (b *joinedLocalBody) Close() error { err := b.ReadCloser.Close(); b.transfer.Wait(); return err } diff --git a/internal/service/task/namespace_mount_test.go b/internal/service/task/namespace_mount_test.go new file mode 100644 index 0000000..61a1e9a --- /dev/null +++ b/internal/service/task/namespace_mount_test.go @@ -0,0 +1,134 @@ +package task + +import ( + "context" + "encoding/json" + "fmt" + "io" + "net" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strconv" + "strings" + "sync/atomic" + "testing" + + "dingospeed/internal/model/query" + "dingospeed/pkg/config" + "dingospeed/pkg/repository" +) + +func TestCacheProgressReaderReportsBytesBeforeFileCompletes(t *testing.T) { + var progress atomic.Uint64 + reader := &cacheProgressReader{reader: strings.NewReader("first-second"), add: progress.Add} + buffer := make([]byte, 5) + + n, err := reader.Read(buffer) + if err != nil || n != 5 || progress.Load() != 5 { + t.Fatalf("first read n=%d progress=%d err=%v", n, progress.Load(), err) + } + if _, err := io.Copy(io.Discard, reader); err != nil { + t.Fatal(err) + } + if progress.Load() != uint64(len("first-second")) { + t.Fatalf("final progress=%d", progress.Load()) + } +} + +func TestMountAndPreheatUseCanonicalNamespaceAndMachineToken(t *testing.T) { + previous := config.SysConfig + t.Cleanup(func() { config.SysConfig = previous }) + config.SysConfig = &config.Config{Server: config.ServerConfig{}, Cache: config.Cache{MountModelDir: t.TempDir()}} + repo, name := "team/resolve/model-a", "weights/50%#.bin" + fileRequests := map[string]int{} + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + parts := strings.Split(r.URL.Path, "/") + if len(parts) != 6 || parts[1] != "api" || parts[2] != "repositories" || parts[3] != "models" { + t.Errorf("unexpected route %q", r.URL.Path) + w.WriteHeader(400) + return + } + ns, operation := parts[4], parts[5] + if r.Header.Get("X-Dingo-Service-Token") != "" || r.Header.Get("Authorization") != "Bearer provider" || r.URL.Query().Get("repo") != repo { + t.Errorf("lost credentials or repository: %s", r.URL.String()) + w.WriteHeader(401) + return + } + switch operation { + case "metadata": + _ = json.NewEncoder(w).Encode(map[string]interface{}{"sha": "commit-1", "siblings": []map[string]string{{"rfilename": name}}}) + case "file": + if r.URL.Query().Get("revision") != "commit-1" || r.URL.Query().Get("path") != name { + t.Errorf("lost file locator: %s", r.URL.String()) + w.WriteHeader(400) + return + } + fileRequests[ns]++ + if ns == repository.ModelScope { + w.Header().Set("Trailer", "X-Dingo-Cache-Complete") + w.Header().Set("Content-Range", fmt.Sprintf("bytes 0-%d/%d", len(ns+" bytes")-1, len(ns+" bytes"))) + w.WriteHeader(206) + } + fmt.Fprint(w, ns+" bytes") + if ns == repository.ModelScope { + w.Header().Set("X-Dingo-Cache-Complete", "true") + } + default: + w.WriteHeader(404) + } + })) + t.Cleanup(server.Close) + host, port, err := net.SplitHostPort(strings.TrimPrefix(server.URL, "http://")) + if err != nil { + t.Fatal(err) + } + config.SysConfig.Server.Host = host + config.SysConfig.Server.Port, _ = strconv.Atoi(port) + for _, ns := range []string{"alice", "用户#50%", repository.ModelScope} { + key := repository.RepoKey{Namespace: ns, RepoType: "models", Repo: repo} + job := &query.CreateCacheJobReq{Namespace: ns, Datatype: "models", Repo: repo} + base := CacheTask{Ctx: context.Background(), Job: job} + mount := MountCacheTask{CacheTask: base, Authorization: "Bearer provider"} + if err := mount.mountViaRepositoryAPI(key, "main"); err != nil { + t.Fatal(err) + } + content, err := os.ReadFile(filepath.Join(config.SysConfig.Cache.MountModelDir, "models", ns, filepath.FromSlash(repo), filepath.FromSlash(name))) + if err != nil || string(content) != ns+" bytes" { + t.Fatalf("mount %s: %q, %v", ns, content, err) + } + meta, err := ReadRepositoryMetadata(context.Background(), key, "main", "Bearer provider") + if err != nil { + t.Fatal(err) + } + preheat := PreheatCacheTask{CacheTask: base, Sha: meta, Authorization: "Bearer provider"} + if err := preheat.preheatProcess(key.ID()); err != nil { + t.Fatal(err) + } + if preheat.stockLen.Load() != uint64(len(ns+" bytes")) || fileRequests[ns] != 2 { + t.Fatalf("preheat did not read canonical bytes for %s", ns) + } + } +} + +func TestCanonicalTaskRequestsRefuseRedirectWithoutServiceToken(t *testing.T) { + previous := config.SysConfig + t.Cleanup(func() { config.SysConfig = previous }) + config.SysConfig = &config.Config{Server: config.ServerConfig{}} + hits := 0 + untrusted := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { hits++; w.WriteHeader(500) })) + defer untrusted.Close() + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { http.Redirect(w, r, untrusted.URL, 302) })) + defer server.Close() + host, port, _ := net.SplitHostPort(strings.TrimPrefix(server.URL, "http://")) + config.SysConfig.Server.Host = host + config.SysConfig.Server.Port, _ = strconv.Atoi(port) + key := repository.RepoKey{Namespace: "alice", RepoType: "models", Repo: "team/model"} + if _, err := ReadRepositoryMetadata(context.Background(), key, "main", ""); err == nil { + t.Fatal("redirect accepted") + } + if hits != 0 { + t.Fatal("machine request followed redirect") + } +} diff --git a/internal/service/task/preheat_cache_task.go b/internal/service/task/preheat_cache_task.go index cd69a41..00ddffa 100644 --- a/internal/service/task/preheat_cache_task.go +++ b/internal/service/task/preheat_cache_task.go @@ -2,6 +2,7 @@ package task import ( "context" + "errors" "fmt" "math" "sync" @@ -13,9 +14,10 @@ import ( "dingospeed/internal/model/query" "dingospeed/pkg/config" "dingospeed/pkg/consts" + "dingospeed/pkg/hfprojection" + "dingospeed/pkg/repository" "go.uber.org/zap" - "golang.org/x/sync/errgroup" ) type CacheTask struct { @@ -25,6 +27,16 @@ type CacheTask struct { Job *query.CreateCacheJobReq SchedulerDao *dao.SchedulerDao RunningStatus int32 + LocalJob bool + stopRequested int32 +} + +func (c *CacheTask) RequestStop() { atomic.StoreInt32(&c.stopRequested, 1); c.CancelFunc() } +func (c *CacheTask) StopStatus() int32 { + if atomic.LoadInt32(&c.stopRequested) != 0 { + return consts.RunningStatusJobStop + } + return c.RunningStatus } func (c *CacheTask) GetTaskNo() int { @@ -43,25 +55,71 @@ type PreheatCacheTask struct { DownloaderDao *dao.DownloaderDao UsedStorage uint64 stockLen atomic.Uint64 + progressMu sync.RWMutex StockSpeed string StockProcess float32 + OnFinish func(error) + OnStart func() + InitialState string +} + +func (p *PreheatCacheTask) RealtimeProgress() (string, float32) { + p.progressMu.RLock() + defer p.progressMu.RUnlock() + return p.StockSpeed, p.StockProcess } +func (p *PreheatCacheTask) CachedBytes() uint64 { return p.stockLen.Load() } + func (p *PreheatCacheTask) DoTask() { - orgRepo := fmt.Sprintf("%s/%s", p.Job.Org, p.Job.Repo) + if p.OnStart != nil { + p.OnStart() + } + orgRepo := fmt.Sprintf("%s/%s", p.Job.Namespace, p.Job.Repo) ctx, cancelFunc := context.WithCancel(p.Ctx) defer cancelFunc() - go p.realTimeSpeed(ctx) + progressDone := make(chan struct{}) + go func() { + defer close(progressDone) + p.realTimeSpeed(ctx) + }() err := p.preheatProcess(orgRepo) + // Stop sampling before recording the terminal progress; a late tick must + // not overwrite 100% with the last in-flight percentage. + cancelFunc() + <-progressDone + if p.OnFinish != nil { + if err == nil && p.Job.Namespace == repository.HuggingFace { + hfprojection.DefaultIndex.Invalidate(config.SysConfig.Repos(), p.Job.Datatype) + } + p.OnFinish(err) + return + } + if p.LocalJob { + if err == nil && p.Job.Namespace == repository.HuggingFace { + hfprojection.DefaultIndex.Invalidate(config.SysConfig.Repos(), p.Job.Datatype) + } + return + } if err != nil { - p.SchedulerDao.ExecUpdateCacheJobStatus(p.TaskNo, p.RunningStatus, p.Job.InstanceId, p.Job.Org, p.Job.Repo, err.Error(), p.StockProcess) + p.SchedulerDao.ExecUpdateCacheJobStatus(p.TaskNo, p.RunningStatus, p.Job.InstanceId, p.Job.Namespace, p.Job.Repo, err.Error(), p.StockProcess) return } p.StockProcess = 100 - p.SchedulerDao.ExecUpdateCacheJobStatus(p.TaskNo, consts.RunningStatusJobComplete, p.Job.InstanceId, p.Job.Org, p.Job.Repo, "", p.StockProcess) + if p.Job.Namespace == repository.HuggingFace { + hfprojection.DefaultIndex.Invalidate(config.SysConfig.Repos(), p.Job.Datatype) + } + p.SchedulerDao.ExecUpdateCacheJobStatus(p.TaskNo, consts.RunningStatusJobComplete, p.Job.InstanceId, p.Job.Namespace, p.Job.Repo, "", p.StockProcess) } func (p *PreheatCacheTask) preheatProcess(orgRepo string) error { + key, err := repository.ParseID(p.Job.Datatype, orgRepo) + if err != nil { + return err + } + if key.Namespace != repository.HuggingFace { + return p.preheatViaRepositoryAPI(key) + } limit := make(chan struct{}, 8) for _, rFile := range p.Sha.Siblings { if p.Ctx.Err() != nil { @@ -69,12 +127,16 @@ func (p *PreheatCacheTask) preheatProcess(orgRepo string) error { } fileName := rFile.Rfilename var hfUri string + upstream, upErr := dao.UpstreamRepo(p.Job.Datatype, orgRepo) + if upErr != nil { + return upErr + } if p.Job.Datatype == "models" { - hfUri = fmt.Sprintf("/%s/resolve/%s/%s", orgRepo, p.Sha.Sha, fileName) + hfUri = fmt.Sprintf("/%s/resolve/%s/%s", repository.EscapeURLPath(upstream), repository.EscapeURLPath(p.Sha.Sha), repository.EscapeURLPath(fileName)) } else { - hfUri = fmt.Sprintf("/%s/%s/resolve/%s/%s", p.Job.Datatype, orgRepo, p.Sha.Sha, fileName) + hfUri = fmt.Sprintf("/%s/%s/resolve/%s/%s", p.Job.Datatype, repository.EscapeURLPath(upstream), repository.EscapeURLPath(p.Sha.Sha), repository.EscapeURLPath(fileName)) } - pathInfo, err := p.FileDao.GetPathsInfo(hfUri, p.Job.Datatype, orgRepo, p.Sha.Sha, + pathInfo, err := p.FileDao.GetPathsInfoContext(p.Ctx, hfUri, p.Job.Datatype, orgRepo, p.Sha.Sha, p.Authorization, fileName) // 获取模型元数据 if err != nil { zap.S().Errorf("RemoteRequestPathsInfo err,%v", err) @@ -89,10 +151,17 @@ func (p *PreheatCacheTask) preheatProcess(orgRepo string) error { } else { etag = pathInfo.Oid } + if err := repository.Segment(etag); err != nil { + return err + } + // A reused blob still needs a file reference under this new commit. + if err := p.FileDao.ConstructBlobsAndFileFile(dao.BlobPath(p.Job.Datatype, orgRepo, etag), dao.ResolvePath(p.Job.Datatype, orgRepo, p.Sha.Sha, fileName)); err != nil { + return err + } if pathInfo.Size == 0 { continue } - offset := p.FileDao.GetFileOffset(p.Job.Datatype, p.Job.Org, p.Job.Repo, etag, pathInfo.Size) + offset := p.FileDao.GetFileOffset(p.Job.Datatype, p.Job.Namespace, p.Job.Repo, etag, pathInfo.Size) if offset > 0 { p.stockLen.Add(uint64(offset)) } @@ -109,7 +178,6 @@ func (p *PreheatCacheTask) preheatProcess(orgRepo string) error { } func (p *PreheatCacheTask) startPreheat(hfUri, orgRepo, fileName, commit, etag, authorization string, fileSize, offset int64) error { - var wg sync.WaitGroup bgCtx := context.WithValue(p.Ctx, consts.PromSource, "localhost") responseChan := make(chan []byte, config.SysConfig.Download.RespChanSize) blobsFile := dao.BlobPath(p.Job.Datatype, orgRepo, etag) @@ -119,7 +187,10 @@ func (p *PreheatCacheTask) startPreheat(hfUri, orgRepo, fileName, commit, etag, return err } taskParam := &downloader.TaskParam{ + LocalOnly: p.LocalJob, TaskNo: 0, + RepoKey: repository.RepoKey{Namespace: p.Job.Namespace, RepoType: p.Job.Datatype, Repo: p.Job.Repo}, + Revision: commit, BlobsFile: blobsFile, FileName: fileName, FileSize: fileSize, @@ -128,22 +199,27 @@ func (p *PreheatCacheTask) startPreheat(hfUri, orgRepo, fileName, commit, etag, Uri: hfUri, DataType: p.Job.Datatype, Etag: etag, + CacheResult: make(chan error, 1), } taskParam.Context = bgCtx taskParam.ResponseChan = responseChan taskParam.Cancel = p.CancelFunc - wg.Add(2) - eg, ctx := errgroup.WithContext(bgCtx) - eg.Go(func() error { - return p.result(ctx, responseChan) - }) - eg.Go(func() error { - if err := p.DownloaderDao.FileDownload(offset, fileSize, false, taskParam); err != nil { - return err - } - return nil - }) - return eg.Wait() + if err := p.DownloaderDao.FileDownload(offset, fileSize, false, taskParam); err != nil { + return err + } + // Consumers may stop immediately, but the cache task must join the producer. + consumeErr := p.result(bgCtx, responseChan) + cacheErr := <-taskParam.CacheResult + if cacheErr != nil && !errors.Is(cacheErr, context.Canceled) { + return cacheErr + } + if consumeErr != nil { + return consumeErr + } + if p.Ctx.Err() != nil { + return p.Ctx.Err() + } + return cacheErr } func (p *PreheatCacheTask) result(ctx context.Context, responseChan chan []byte) error { @@ -169,6 +245,7 @@ func (p *PreheatCacheTask) realTimeSpeed(ctx context.Context) { case <-ticker.C: currentBytes := p.stockLen.Load() delta := currentBytes - lastBytes + p.progressMu.Lock() p.StockSpeed = formatSpeed(delta, 1*time.Second) // 采样间隔 1 秒 // 计算下载进度（百分比） process := float64(currentBytes) / float64(p.UsedStorage) * 100 @@ -176,10 +253,13 @@ func (p *PreheatCacheTask) realTimeSpeed(ctx context.Context) { process = 99.9 } p.StockProcess = float32(math.Round(process*10) / 10) + p.progressMu.Unlock() lastBytes = currentBytes case <-ctx.Done(): zap.S().Debug("speed ctx done") + p.progressMu.Lock() p.StockSpeed = "0 B/s" + p.progressMu.Unlock() return } } diff --git a/internal/service/task/preheat_reuse_test.go b/internal/service/task/preheat_reuse_test.go new file mode 100644 index 0000000..4689131 --- /dev/null +++ b/internal/service/task/preheat_reuse_test.go @@ -0,0 +1,66 @@ +package task + +import ( + "context" + "encoding/binary" + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strings" + "testing" + + "dingospeed/internal/dao" + "dingospeed/internal/data" + "dingospeed/internal/model/query" + "dingospeed/pkg/config" +) + +func TestPreheatReusedBlobGetsNewCommitReference(t *testing.T) { + old := config.SysConfig + defer func() { config.SysConfig = old }() + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != "POST" || !strings.Contains(r.URL.Path, "/paths-info/") { + t.Errorf("unexpected download %s %s", r.Method, r.URL) + w.WriteHeader(500) + return + } + fmt.Fprint(w, `[{"type":"file","path":"model.bin","oid":"shared","size":8}]`) + })) + defer upstream.Close() + config.SysConfig = &config.Config{Server: config.ServerConfig{Online: true, Repos: t.TempDir(), HfScheme: "http", HfNetLoc: strings.TrimPrefix(upstream.URL, "http://")}, Retry: config.Retry{Attempts: 1}, Download: config.Download{BlockSize: 4}} + base := data.NewBaseData() + files := dao.NewFileDao(nil, base, dao.NewLockDao(base)) + blob := dao.BlobPath("models", "huggingface/owner/model", "shared") + content := make([]byte, 45) + copy(content, "OLAH") + binary.LittleEndian.PutUint64(content[4:], 8) + binary.LittleEndian.PutUint64(content[12:], 4) + binary.LittleEndian.PutUint64(content[20:], 8) + binary.LittleEndian.PutUint64(content[28:], 8) + content[36] = 3 + copy(content[37:], "abcdefgh") + if err := os.MkdirAll(filepath.Dir(blob), 0755); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(blob, content, 0600); err != nil { + t.Fatal(err) + } + for _, commit := range []string{"old-commit", "new-commit", "new-commit"} { + var meta dao.CommitHfSha + _ = json.Unmarshal([]byte(fmt.Sprintf(`{"sha":%q,"siblings":[{"rfilename":"model.bin"}]}`, commit)), &meta) + p := PreheatCacheTask{CacheTask: CacheTask{Ctx: context.Background(), Job: &query.CreateCacheJobReq{Namespace: "huggingface", Repo: "owner/model", Datatype: "models"}}, Sha: &meta, FileDao: files} + if err := p.preheatProcess("huggingface/owner/model"); err != nil { + t.Fatal(err) + } + if p.stockLen.Load() != 8 { + t.Fatalf("reused=%d", p.stockLen.Load()) + } + got, err := os.ReadFile(dao.ResolvePath("models", "huggingface/owner/model", commit, "model.bin")) + if err != nil || string(got) != string(content) { + t.Fatalf("reference err=%v bytes=%d", err, len(got)) + } + } +} diff --git a/internal/service/upload_chunk_service.go b/internal/service/upload_chunk_service.go index 2c2e181..fabcb6f 100644 --- a/internal/service/upload_chunk_service.go +++ b/internal/service/upload_chunk_service.go @@ -9,6 +9,7 @@ import ( "dingospeed/internal/dao" "dingospeed/internal/downloader" "dingospeed/pkg/config" + "dingospeed/pkg/transfersettings" "go.uber.org/zap" ) @@ -25,7 +26,7 @@ type chunkLimiter struct { var uploadChunkLimiter chunkLimiter func (l *chunkLimiter) acquire() bool { - limit := config.SysConfig.GetUploadChunkConcurrentLimit() + limit := transfersettings.Current().Upload l.mu.Lock() defer l.mu.Unlock() if l.inFlight >= limit { @@ -50,12 +51,12 @@ func (l *chunkLimiter) release() { func (u *UploadService) UploadChunk(param dao.LocalChunkUploadParam, rawSize, rawOffset, chunkSha string, contentLength int64, body io.Reader) (*dao.LocalChunkUploadResult, error) { if err := validateUploadLocator(dao.LocalUploadParam{ - RepoType: param.RepoType, - Org: param.Org, - Repo: param.Repo, - Revision: param.Revision, - FilePath: param.FilePath, - Sha256: param.Sha256, + RepoType: param.RepoType, + Namespace: param.Namespace, + Repo: param.Repo, + Revision: param.Revision, + FilePath: param.FilePath, + Sha256: param.Sha256, }); err != nil { return nil, chunkArgError(err.Error()) } @@ -113,7 +114,7 @@ func (u *UploadService) UploadChunk(param dao.LocalChunkUploadParam, rawSize, ra return nil, uploadError{status: 500, code: "UPLOAD_INTERNAL_ERROR", msg: err.Error()} } zap.S().Debugf("local chunk upload done: %s/%s/%s/%s offset=%d length=%d status=%s blocks=%d", - param.RepoType, param.Org, param.Repo, param.FilePath, param.Offset, param.Length, result.Status, result.Blocks) + param.RepoType, param.Namespace, param.Repo, param.FilePath, param.Offset, param.Length, result.Status, result.Blocks) return result, nil } diff --git a/internal/service/upload_inventory_test.go b/internal/service/upload_inventory_test.go new file mode 100644 index 0000000..d93ae9a --- /dev/null +++ b/internal/service/upload_inventory_test.go @@ -0,0 +1,87 @@ +package service + +import ( + "bytes" + "crypto/sha256" + "encoding/hex" + "os" + "testing" + + "dingospeed/internal/dao" + "dingospeed/internal/data" + "dingospeed/pkg/config" +) + +func TestBuildUploadInventoryUnionsLiveRevisionsAndRejectsDamage(t *testing.T) { + old := config.SysConfig + config.SysConfig = &config.Config{Server: config.ServerConfig{Repos: t.TempDir()}, Download: config.Download{BlockSize: 16}, Upload: config.Upload{Namespace: "dingo-local"}, Scheduler: config.Scheduler{Discovery: config.Discovery{InstanceId: "node-a"}}} + t.Cleanup(func() { config.SysConfig = old }) + base := data.NewBaseData() + locks := dao.NewLockDao(base) + files := dao.NewFileDao(nil, base, locks) + uploads := dao.NewUploadDao(files, locks) + meta := NewMetaService(files, nil) + svc := &SchedulerService{metaService: meta} + + upload := func(revision, path string, body []byte) string { + t.Helper() + h := sha256.Sum256(body) + result, err := uploads.UploadWholeFile(dao.LocalUploadParam{RepoType: "models", Namespace: "team", Repo: "demo", Revision: revision, FilePath: path, Size: int64(len(body)), Sha256: hex.EncodeToString(h[:])}, bytes.NewReader(body)) + if err != nil { + t.Fatal(err) + } + return result.Commit + } + shared := []byte("shared") + mainCommit := upload("main", "a.bin", shared) + v2Commit := upload("v2", "a.bin", shared) + upload("v2", "b.bin", shared) // same content, different path + upload("other", "a.bin", []byte("different")) // same path, different content + + snap, err := svc.buildUploadInventory() + if err != nil || !snap.Complete { + t.Fatalf("snapshot complete=%v err=%v detail=%s", snap.Complete, err, snap.Error) + } + if len(snap.Items) != 3 { + t.Fatalf("items=%d, want 3: %+v", len(snap.Items), snap.Items) + } + if err = persistUploadInventory(snap); err != nil { + t.Fatal(err) + } + next, err := svc.buildUploadInventory() + if err != nil || next.Epoch != snap.Epoch || next.Sequence != snap.Sequence+1 { + t.Fatalf("next=%+v err=%v", next, err) + } + + // Removing one local revision reference must not remove the node holding while + // another live revision still references the same path/content. + if _, err = uploads.PublishTree(dao.LocalPublishTreeParam{RepoType: "models", Namespace: "team", Repo: "demo", Revision: "main", BaseCommit: mainCommit, Files: []dao.LocalManifestFile{}}); err != nil { + t.Fatal(err) + } + afterOne, err := svc.buildUploadInventory() + if err != nil || !afterOne.Complete { + t.Fatalf("after one complete=%v err=%v %s", afterOne.Complete, err, afterOne.Error) + } + foundSharedA := false + for _, got := range afterOne.Items { + if got.Path == "a.bin" && got.Size == int64(len(shared)) { + foundSharedA = true + } + } + if !foundSharedA { + t.Fatal("shared a.bin disappeared while v2 still referenced it") + } + + // A physically missing/corrupt published blob makes the scan incomplete; it + // must not be emitted as authoritative absence. + _ = v2Commit + sharedHash := sha256.Sum256(shared) + blob := dao.RepositoryKey("models", "team/demo").Blob(config.SysConfig.Repos(), hex.EncodeToString(sharedHash[:])) + if err = os.Remove(blob); err != nil { + t.Fatal(err) + } + damaged, err := svc.buildUploadInventory() + if err != nil || damaged.Complete { + t.Fatalf("damaged complete=%v err=%v", damaged.Complete, err) + } +} diff --git a/internal/service/upload_publish_test.go b/internal/service/upload_publish_test.go index bca9a6e..176c061 100644 --- a/internal/service/upload_publish_test.go +++ b/internal/service/upload_publish_test.go @@ -14,10 +14,10 @@ const testSha = "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcde func validPublishParam() dao.LocalPublishParam { return dao.LocalPublishParam{ - RepoType: "models", - Org: "dingo-local", - Repo: "demo", - Revision: "main", + RepoType: "models", + Namespace: "dingo-local", + Repo: "demo", + Revision: "main", Files: []dao.LocalManifestFile{ {Path: "config.json", Sha256: testSha, Size: 10}, }, @@ -115,7 +115,7 @@ func TestValidatePublishParamRejectionMatrix(t *testing.T) { }) t.Run("org/"+bad.name, func(t *testing.T) { param := validPublishParam() - param.Org = bad.value + param.Namespace = bad.value if err := validatePublishParam(param); err == nil { t.Fatalf("expected org %q to be rejected", bad.value) } @@ -124,7 +124,7 @@ func TestValidatePublishParamRejectionMatrix(t *testing.T) { t.Run("org outside namespace", func(t *testing.T) { param := validPublishParam() - param.Org = "huggingface" + param.Namespace = "huggingface" if err := validatePublishParam(param); err == nil { t.Fatalf("expected org outside the reserved namespace to be rejected") } @@ -214,7 +214,7 @@ func TestValidatePublishParamListShape(t *testing.T) { func validPublishTreeParam() dao.LocalPublishTreeParam { return dao.LocalPublishTreeParam{ RepoType: "models", - Org: "dingo-local", + Namespace: "dingo-local", Repo: "demo", Revision: "main", BaseCommit: testSha, diff --git a/internal/service/upload_service.go b/internal/service/upload_service.go index 9e4c6cf..0bc326c 100644 --- a/internal/service/upload_service.go +++ b/internal/service/upload_service.go @@ -13,6 +13,8 @@ import ( "dingospeed/internal/dao" "dingospeed/internal/downloader" "dingospeed/pkg/config" + "dingospeed/pkg/repository" + "dingospeed/pkg/transfersettings" "go.uber.org/zap" ) @@ -27,6 +29,10 @@ func NewUploadService(uploadDao *dao.UploadDao) *UploadService { return &UploadService{uploadDao: uploadDao} } +func (u *UploadService) DeleteRepository(k repository.RepoKey) (*dao.RepositoryDeleteResult, error) { + return u.uploadDao.DeleteHostedRepository(k) +} + // UploadWholeFile 一次调用上传一个完整文件。rawSize 是调用方声明的完整字节大小， // 由本方法在身份校验之后解析，保证未通过身份校验的请求不会先收到参数错误。 func (u *UploadService) UploadWholeFile(param dao.LocalUploadParam, rawSize, rawStart string, body io.Reader) (*dao.LocalUploadResult, error) { @@ -55,7 +61,7 @@ func (u *UploadService) UploadWholeFile(param dao.LocalUploadParam, rawSize, raw startLog = fmt.Sprintf("%d", *param.Start) } zap.S().Infof("local upload start: %s/%s/%s/%s size=%d sha256=%s overwrite=%t start=%s", - param.RepoType, param.Org, param.Repo, param.FilePath, param.Size, param.Sha256, param.Overwrite, startLog) + param.RepoType, param.Namespace, param.Repo, param.FilePath, param.Size, param.Sha256, param.Overwrite, startLog) result, err := u.uploadDao.UploadWholeFile(param, body) if err != nil { if _, ok := err.(interface{ StatusCode() int }); ok { @@ -64,7 +70,8 @@ func (u *UploadService) UploadWholeFile(param dao.LocalUploadParam, rawSize, raw return nil, uploadError{status: 500, code: "UPLOAD_INTERNAL_ERROR", msg: err.Error()} } zap.S().Infof("local upload done: %s/%s/%s/%s status=%s commit=%s blobReused=%t", - param.RepoType, param.Org, param.Repo, param.FilePath, result.Status, result.Commit, result.BlobReused) + param.RepoType, param.Namespace, param.Repo, param.FilePath, result.Status, result.Commit, result.BlobReused) + dao.NotifyPublished() return result, nil } @@ -74,7 +81,7 @@ func (u *UploadService) PublishFiles(param dao.LocalPublishParam) (*dao.LocalPub return nil, uploadError{status: 400, code: "PUBLISH_INVALID_ARGUMENT", msg: err.Error()} } zap.S().Infof("local publish start: %s/%s/%s revision=%s files=%d overwrite=%t", - param.RepoType, param.Org, param.Repo, param.Revision, len(param.Files), param.Overwrite) + param.RepoType, param.Namespace, param.Repo, param.Revision, len(param.Files), param.Overwrite) result, err := u.uploadDao.PublishFiles(param) if err != nil { if _, ok := err.(interface{ StatusCode() int }); ok { @@ -83,13 +90,14 @@ func (u *UploadService) PublishFiles(param dao.LocalPublishParam) (*dao.LocalPub return nil, uploadError{status: 500, code: "PUBLISH_INTERNAL_ERROR", msg: err.Error()} } zap.S().Infof("local publish done: %s/%s/%s revision=%s commit=%s status=%s added=%d replaced=%d unchanged=%d total=%d", - param.RepoType, param.Org, param.Repo, param.Revision, result.Commit, result.Status, + param.RepoType, param.Namespace, param.Repo, param.Revision, result.Commit, result.Status, result.Added, result.Replaced, result.Unchanged, result.FileCount) + dao.NotifyPublished() return result, nil } func validatePublishParam(param dao.LocalPublishParam) error { - if err := validateRepoLocator(param.RepoType, param.Org, param.Repo, param.Revision); err != nil { + if err := validateRepoLocator(param.RepoType, param.Namespace, param.Repo, param.Revision); err != nil { return err } return validateManifestList(param.Files) @@ -134,7 +142,7 @@ func (u *UploadService) PublishTree(param dao.LocalPublishTreeParam) (*dao.Local return nil, uploadError{status: 400, code: "PUBLISH_TREE_INVALID_ARGUMENT", msg: err.Error()} } zap.S().Infof("local publish tree start: %s/%s/%s revision=%s base=%s files=%d", - param.RepoType, param.Org, param.Repo, param.Revision, param.BaseCommit, len(param.Files)) + param.RepoType, param.Namespace, param.Repo, param.Revision, param.BaseCommit, len(param.Files)) result, err := u.uploadDao.PublishTree(param) if err != nil { if _, ok := err.(interface{ StatusCode() int }); ok { @@ -143,13 +151,14 @@ func (u *UploadService) PublishTree(param dao.LocalPublishTreeParam) (*dao.Local return nil, uploadError{status: 500, code: "PUBLISH_TREE_INTERNAL_ERROR", msg: err.Error()} } zap.S().Infof("local publish tree done: %s/%s/%s revision=%s commit=%s previous=%s status=%s added=%d replaced=%d removed=%d total=%d", - param.RepoType, param.Org, param.Repo, param.Revision, result.Commit, result.PreviousCommit, result.Status, + param.RepoType, param.Namespace, param.Repo, param.Revision, result.Commit, result.PreviousCommit, result.Status, result.Added, result.Replaced, result.Removed, result.FileCount) + dao.NotifyPublished() return result, nil } func validatePublishTreeParam(param dao.LocalPublishTreeParam) error { - if err := validateRepoLocator(param.RepoType, param.Org, param.Repo, param.Revision); err != nil { + if err := validateRepoLocator(param.RepoType, param.Namespace, param.Repo, param.Revision); err != nil { return err } if !sha256Pattern.MatchString(param.BaseCommit) { @@ -198,7 +207,7 @@ func ParseResumeStart(raw string) (*int64, error) { } func validateUploadLocator(param dao.LocalUploadParam) error { - if err := validateRepoLocator(param.RepoType, param.Org, param.Repo, param.Revision); err != nil { + if err := validateRepoLocator(param.RepoType, param.Namespace, param.Repo, param.Revision); err != nil { return err } return validateFileLocator(param.FilePath, param.Sha256) @@ -210,10 +219,13 @@ func validateRepoLocator(repoType, org, repo, revision string) error { if repoType != "models" && repoType != "datasets" { return fmt.Errorf("repoType must be models or datasets") } - if org != config.SysConfig.Upload.Namespace { - return fmt.Errorf("org must be %s", config.SysConfig.Upload.Namespace) + if org == repository.HuggingFace || org == repository.ModelScope { + return fmt.Errorf("remote namespaces do not accept uploads") + } + if _, _, err := (repository.RepoKey{Namespace: org, RepoType: repoType, Repo: repo}).Storage(); err != nil { + return err } - if !validRepoSegment(org) || !validRepoSegment(repo) || !validRepoSegment(revision) { + if !validRepoSegment(revision) || repository.Segment(revision) != nil { return fmt.Errorf("org, repo, and revision must be safe relative path segments") } return nil @@ -254,7 +266,7 @@ func validateUploadParam(param dao.LocalUploadParam) error { const ( maxSegmentLen = 255 - maxFilePathLen = 1024 + maxFilePathLen = 1000 ) // repoSegmentPattern 约束组织名、仓库名与版本标签。这三者会出现在对外 URL 和目录名里， @@ -320,7 +332,7 @@ func (e uploadError) ErrorCode() string { } func (u *UploadService) acquireUploadSlot() bool { - limit := config.SysConfig.Upload.ConcurrentLimit + limit := transfersettings.Current().Upload if limit <= 0 { limit = 1 } diff --git a/internal/service/upload_service_test.go b/internal/service/upload_service_test.go index d829ffe..302f4bf 100644 --- a/internal/service/upload_service_test.go +++ b/internal/service/upload_service_test.go @@ -10,6 +10,8 @@ import ( "dingospeed/internal/dao" "dingospeed/internal/downloader" "dingospeed/pkg/config" + "dingospeed/pkg/repository" + "dingospeed/pkg/transfersettings" ) func withUploadConfig(t *testing.T) { @@ -24,13 +26,13 @@ func withUploadConfig(t *testing.T) { func validParam() dao.LocalUploadParam { return dao.LocalUploadParam{ - RepoType: "models", - Org: "dingo-local", - Repo: "demo", - Revision: "main", - FilePath: "config.json", - Size: 10, - Sha256: "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef", + RepoType: "models", + Namespace: "dingo-local", + Repo: "demo", + Revision: "main", + FilePath: "config.json", + Size: 10, + Sha256: "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef", } } @@ -106,7 +108,7 @@ func TestValidateUploadParamRejectionMatrix(t *testing.T) { }) t.Run("org/"+bad.name, func(t *testing.T) { param := validParam() - param.Org = bad.value + param.Namespace = bad.value if err := validateUploadParam(param); err == nil { t.Fatalf("expected org %q to be rejected", bad.value) } @@ -136,7 +138,7 @@ func TestValidateUploadParamRejectionMatrix(t *testing.T) { }) t.Run("org outside namespace", func(t *testing.T) { param := validParam() - param.Org = "huggingface" + param.Namespace = "huggingface" if err := validateUploadParam(param); err == nil { t.Fatalf("expected org outside the reserved namespace to be rejected") } @@ -182,6 +184,19 @@ func TestParseDeclaredSize(t *testing.T) { func TestUploadConcurrencyLimit(t *testing.T) { withUploadConfig(t) config.SysConfig.Upload.ConcurrentLimit = 2 + // Runtime settings are process-wide; other transfer tests may initialize them first. + config.SysConfig.Server.Repos = t.TempDir() + previous := transfersettings.Current() + settings := previous + settings.Upload = 2 + if err := transfersettings.Save(settings); err != nil { + t.Fatal(err) + } + t.Cleanup(func() { + if err := transfersettings.Save(previous); err != nil { + t.Error(err) + } + }) svc := NewUploadService(nil) if !svc.acquireUploadSlot() || !svc.acquireUploadSlot() { @@ -207,7 +222,10 @@ func errorCodeOf(err error) string { func TestIsProtectedLocalUploadCacheFile(t *testing.T) { withUploadConfig(t) - root := filepath.Join("tmp", "repos") + root := t.TempDir() + if err := repository.Register(root, repository.Remote(repository.RepoKey{Namespace: "huggingface", RepoType: "models", Repo: "demo"})); err != nil { + t.Fatal(err) + } protected := []string{ filepath.Join(root, "files", "models", "dingo-local", "demo", "blobs", "abc"), filepath.Join(root, "files", "datasets", "dingo-local", "demo", "blobs", "abc"), @@ -219,10 +237,11 @@ func TestIsProtectedLocalUploadCacheFile(t *testing.T) { } } - cleanable := []string{ - filepath.Join(root, "files", "models", "huggingface", "demo", "blobs", "abc"), - filepath.Join(root, "files", "datasets", "someorg", "demo", "blobs", "abc"), - filepath.Join(root, "files", "models", "dingo-local-lookalike", "demo", "blobs", "abc"), + cleanable := []string{filepath.Join(root, "files", "models", "demo", "blobs", "abc")} + for _, unknown := range []string{"someorg", "dingo-local-lookalike"} { + if !isProtectedLocalUploadCacheFile(root, filepath.Join(root, "files", "models", unknown, "demo", "blobs", "abc")) { + t.Fatal("unknown source must remain protected") + } } for _, path := range cleanable { if isProtectedLocalUploadCacheFile(root, path) { @@ -265,7 +284,10 @@ func TestDiskCleanKeepsLocalUploadsAndStillCleansPublicCache(t *testing.T) { localBlob := write("files", "models", "dingo-local", "demo", "blobs", "aaa") localResolve := write("files", "models", "dingo-local", "demo", "resolve", "commit", "config.json") localDataset := write("files", "datasets", "dingo-local", "demo", "blobs", "bbb") - publicBlob := write("files", "models", "huggingface", "demo", "blobs", "ccc") + publicBlob := write("files", "models", "demo", "blobs", "ccc") + if err := repository.Register(repos, repository.Remote(repository.RepoKey{Namespace: "huggingface", RepoType: "models", Repo: "demo"})); err != nil { + t.Fatal(err) + } (&SysService{}).checkDiskUsage() diff --git a/pkg/cachemount/mount_linux.go b/pkg/cachemount/mount_linux.go new file mode 100644 index 0000000..6ae69d0 --- /dev/null +++ b/pkg/cachemount/mount_linux.go @@ -0,0 +1,112 @@ +package cachemount + +import ( + "context" + "fmt" + "io" + "os" + "strings" + "syscall" + + "github.com/hanwen/go-fuse/v2/fs" + "github.com/hanwen/go-fuse/v2/fuse" +) + +type fileNode struct { + fs.Inode + payload *Payload +} + +func (n *fileNode) Getattr(ctx context.Context, f fs.FileHandle, out *fuse.AttrOut) syscall.Errno { + out.Mode = syscall.S_IFREG | 0444 + out.Size = uint64(n.payload.Size) + return 0 +} +func (n *fileNode) Open(ctx context.Context, flags uint32) (fs.FileHandle, uint32, syscall.Errno) { + if flags&uint32(syscall.O_ACCMODE|syscall.O_TRUNC|syscall.O_APPEND) != 0 { + return nil, 0, syscall.EROFS + } + return nil, fuse.FOPEN_KEEP_CACHE, 0 +} +func (n *fileNode) Read(ctx context.Context, f fs.FileHandle, dest []byte, off int64) (fuse.ReadResult, syscall.Errno) { + count, err := n.payload.ReadAt(dest, off) + // EOF within the declared range is truncation, never a successful short read. + want := int64(len(dest)) + if off >= n.payload.Size { + want = 0 + } else if want > n.payload.Size-off { + want = n.payload.Size - off + } + if off < 0 || (err != nil && err != io.EOF) || int64(count) != want { + return nil, syscall.EIO + } + return fuse.ReadResultData(dest[:count]), 0 +} + +type rootNode struct { + fs.Inode + snapshot *Snapshot +} + +type unavailableDir struct{ fs.Inode } + +func (n *unavailableDir) Readdir(ctx context.Context) (fs.DirStream, syscall.Errno) { + return nil, syscall.EIO +} +func (n *unavailableDir) Lookup(ctx context.Context, name string, out *fuse.EntryOut) (*fs.Inode, syscall.Errno) { + return nil, syscall.EIO +} + +func ensureDir(ctx context.Context, root *fs.Inode, parts []string) *fs.Inode { + dir := root + for _, part := range parts { + next := dir.GetChild(part) + if next == nil { + next = dir.NewPersistentInode(ctx, &fs.Inode{}, fs.StableAttr{Mode: syscall.S_IFDIR}) + dir.AddChild(part, next, false) + } + dir = next + } + return dir +} + +func (r *rootNode) OnAdd(ctx context.Context) { + for name, reason := range r.snapshot.Directories { + parts := strings.Split(name, "/") + parent := ensureDir(ctx, &r.Inode, parts[:len(parts)-1]) + if reason != "" { + parent.AddChild(parts[len(parts)-1], parent.NewPersistentInode(ctx, &unavailableDir{}, fs.StableAttr{Mode: syscall.S_IFDIR}), false) + } else { + ensureDir(ctx, parent, parts[len(parts)-1:]) + } + } + for name, p := range r.snapshot.Files { + parts := strings.Split(name, "/") + dir := &r.Inode + for _, part := range parts[:len(parts)-1] { + next := dir.GetChild(part) + if next == nil { + next = dir.NewPersistentInode(ctx, &fs.Inode{}, fs.StableAttr{Mode: syscall.S_IFDIR}) + dir.AddChild(part, next, false) + } + dir = next + } + n := dir.NewPersistentInode(ctx, &fileNode{payload: p}, fs.StableAttr{Mode: syscall.S_IFREG}) + dir.AddChild(parts[len(parts)-1], n, false) + } +} + +// Mount requires an existing empty directory, never hides existing user files. +// The caller owns Snapshot until the returned server has finished serving. +func Mount(s *Snapshot, mountpoint string) (*fuse.Server, error) { + entries, err := os.ReadDir(mountpoint) + if err != nil { + return nil, err + } + if len(entries) != 0 { + return nil, fmt.Errorf("mountpoint must be empty") + } + return fs.Mount(mountpoint, &rootNode{snapshot: s}, &fs.Options{ + MountOptions: fuse.MountOptions{Options: []string{"ro", "default_permissions"}, FsName: "dingocache", Name: "dingocache"}, + }) +} diff --git a/pkg/cachemount/mount_linux_test.go b/pkg/cachemount/mount_linux_test.go new file mode 100644 index 0000000..24afa9b --- /dev/null +++ b/pkg/cachemount/mount_linux_test.go @@ -0,0 +1,106 @@ +package cachemount + +import ( + "bytes" + "context" + "os" + "path/filepath" + "sync" + "syscall" + "testing" +) + +// Opt in because CI hosts may lack /dev/fuse. An opted-in failure is fatal, +// never skipped, so a successful report proves an actual kernel mount. +func TestRealMount(t *testing.T) { + if os.Getenv("DINGO_TEST_FUSE") != "1" { + t.Skip("set DINGO_TEST_FUSE=1 for real kernel mount") + } + data := bytes.Repeat([]byte("0123456789abcdef"), 4096) + name := fixture(t, data) + s, err := Pin(map[string]string{"sub/weights.bin": name}) + if err != nil { + t.Fatal(err) + } + defer s.Close() + mount := t.TempDir() + server, err := Mount(s, mount) + if err != nil { + t.Fatal(err) + } + defer func() { + if err := server.Unmount(); err != nil { + t.Error(err) + } + server.Wait() + }() + p := filepath.Join(mount, "sub", "weights.bin") + st, err := os.Stat(p) + if err != nil || st.Size() != int64(len(data)) { + t.Fatalf("stat %v %v", st, err) + } + var wg sync.WaitGroup + for i := 0; i < 8; i++ { + wg.Add(1) + go func() { + defer wg.Done() + for j := 0; j < 10; j++ { + b, err := os.ReadFile(p) + if err != nil || !bytes.Equal(b, data) { + t.Errorf("concurrent read: %v", err) + return + } + } + }() + } + wg.Wait() + f, err := os.Open(p) + if err != nil { + t.Fatal(err) + } + defer f.Close() + b, err := syscall.Mmap(int(f.Fd()), 0, len(data), syscall.PROT_READ, syscall.MAP_SHARED) + if err != nil { + t.Fatal(err) + } + if !bytes.Equal(b, data) { + t.Error("mmap differs") + } + syscall.Munmap(b) + if err := os.WriteFile(p, []byte("oops"), 0600); err == nil { + t.Fatal("write accepted") + } + if err := os.Remove(p); err == nil { + t.Fatal("unlink accepted") + } + if err := os.Mkdir(filepath.Join(mount, "new"), 0700); err == nil { + t.Fatal("mkdir accepted") + } + // Unlink and replace the pathname; pinned descriptor must retain old data. + if err := os.Remove(name); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(name, []byte("replacement"), 0600); err != nil { + t.Fatal(err) + } + got, err := os.ReadFile(p) + if err != nil || !bytes.Equal(got, data) { + t.Fatalf("pin lost after unlink: %v", err) + } +} + +func TestTruncationReturnsIOError(t *testing.T) { + name := fixture(t, []byte("abcdefgh")) + p, err := Open(name) + if err != nil { + t.Fatal(err) + } + defer p.File.Close() + if err := os.Truncate(name, p.Offset+2); err != nil { + t.Fatal(err) + } + n := fileNode{payload: p} + if _, errno := n.Read(context.Background(), nil, make([]byte, 8), 0); errno != syscall.EIO { + t.Fatalf("got %v", errno) + } +} diff --git a/pkg/cachemount/payload.go b/pkg/cachemount/payload.go new file mode 100644 index 0000000..eda6fef --- /dev/null +++ b/pkg/cachemount/payload.go @@ -0,0 +1,139 @@ +// Package cachemount exposes complete DingCache payloads without copying or +// modifying cache files. All backing descriptors remain open until unmount. +package cachemount + +import ( + "encoding/binary" + "fmt" + "io" + "math" + "os" + "path" + "strings" +) + +type Payload struct { + File *os.File + Offset int64 + Size int64 +} + +// Open accepts OLAH v8 only. A damaged container never falls back to raw data. +func Open(name string) (*Payload, error) { + f, err := os.Open(name) + if err != nil { + return nil, err + } + p, err := inspect(f) + if err != nil { + f.Close() + return nil, fmt.Errorf("%s: %w", name, err) + } + return p, nil +} + +func inspect(f *os.File) (*Payload, error) { + st, err := f.Stat() + if err != nil { + return nil, err + } + if !st.Mode().IsRegular() { + return nil, fmt.Errorf("not a regular cache file") + } + var h [36]byte + if _, err := f.ReadAt(h[:], 0); err != nil { + return nil, err + } + v := binary.LittleEndian.Uint64(h[4:12]) + block := binary.LittleEndian.Uint64(h[12:20]) + size := binary.LittleEndian.Uint64(h[20:28]) + bits := binary.LittleEndian.Uint64(h[28:36]) + if string(h[:4]) != "OLAH" || v != 8 || block == 0 || block > math.MaxInt64 || size > math.MaxInt64 || bits == 0 || bits > 8*1024*1024 { + return nil, fmt.Errorf("invalid OLAH v8 header") + } + count := size / block + if size%block != 0 { + count++ + } + if count > bits { + return nil, fmt.Errorf("bitmap too small") + } + mask := make([]byte, (bits+7)/8) + if _, err := f.ReadAt(mask, 36); err != nil { + return nil, err + } + offset := int64(36 + len(mask)) + if st.Size()-offset < int64(size) { + return nil, fmt.Errorf("truncated cache payload") + } + for i := uint64(0); i < count; i++ { + if mask[i/8]&(1<<(i%8)) == 0 { + return nil, fmt.Errorf("missing cache block %d", i) + } + } + return &Payload{File: f, Offset: offset, Size: int64(size)}, nil +} + +func (p *Payload) ReadAt(b []byte, off int64) (int, error) { + if off < 0 { + return 0, fmt.Errorf("negative offset") + } + if len(b) == 0 { + return 0, nil + } + if off >= p.Size { + return 0, io.EOF + } + want := len(b) + if int64(len(b)) > p.Size-off { + b = b[:p.Size-off] + } + n, err := p.File.ReadAt(b, p.Offset+off) + if err == nil && n < want { + err = io.EOF + } + return n, err +} + +type Snapshot struct { + Files map[string]*Payload + Directories map[string]string +} + +// Pin validates the whole tree before exposing any file. It reads headers and +// bitmaps only, not the weights. Failure closes all previously opened files. +func Pin(sources map[string]string) (*Snapshot, error) { + s := &Snapshot{Files: make(map[string]*Payload)} + for name, source := range sources { + if name == "." || !validPath(name) { + s.Close() + return nil, fmt.Errorf("invalid model path %q", name) + } + p, err := Open(source) + if err != nil { + s.Close() + return nil, fmt.Errorf("%s: %w", name, err) + } + s.Files[name] = p + } + for name := range s.Files { + for parent := path.Dir(name); parent != "."; parent = path.Dir(parent) { + if _, ok := s.Files[parent]; ok { + s.Close() + return nil, fmt.Errorf("file/directory collision: %s", parent) + } + } + } + if len(s.Files) == 0 { + return nil, fmt.Errorf("empty model manifest") + } + return s, nil +} +func validPath(p string) bool { + return p != "" && !strings.ContainsAny(p, "\\\x00") && !strings.HasPrefix(p, "/") && path.Clean(p) == p && p != ".." && !strings.HasPrefix(p, "../") +} +func (s *Snapshot) Close() { + for _, p := range s.Files { + p.File.Close() + } +} diff --git a/pkg/cachemount/payload_test.go b/pkg/cachemount/payload_test.go new file mode 100644 index 0000000..a111aba --- /dev/null +++ b/pkg/cachemount/payload_test.go @@ -0,0 +1,147 @@ +package cachemount + +import ( + "bytes" + "encoding/binary" + "io" + "os" + "path/filepath" + "testing" + "time" +) + +func fixture(t *testing.T, data []byte) string { + t.Helper() + name := filepath.Join(t.TempDir(), "blob") + bits := uint64(8 * ((len(data)+3)/4 + 7) / 8) + if bits == 0 { + bits = 8 + } + h := make([]byte, 36+(bits+7)/8) + copy(h, "OLAH") + binary.LittleEndian.PutUint64(h[4:], 8) + binary.LittleEndian.PutUint64(h[12:], 4) + binary.LittleEndian.PutUint64(h[20:], uint64(len(data))) + binary.LittleEndian.PutUint64(h[28:], bits) + for i := 36; i < len(h); i++ { + h[i] = 255 + } + if err := os.WriteFile(name, append(h, data...), 0600); err != nil { + t.Fatal(err) + } + return name +} + +func TestReadRanges(t *testing.T) { + data := []byte("0123456789abcdefg") + name := fixture(t, data) + before, _ := os.ReadFile(name) + p, err := Open(name) + if err != nil { + t.Fatal(err) + } + defer p.File.Close() + for off := 0; off <= len(data)+1; off++ { + for length := 1; length < 24; length++ { + b := make([]byte, length) + n, err := p.ReadAt(b, int64(off)) + expect := []byte{} + if off < len(data) { + expect = data[off:min(off+length, len(data))] + } + if !bytes.Equal(b[:n], expect) || (n < length && err != io.EOF) { + t.Fatalf("off=%d length=%d n=%d err=%v", off, length, n, err) + } + } + } + after, _ := os.ReadFile(name) + if !bytes.Equal(before, after) { + t.Fatal("source changed") + } +} + +func TestRejectMalformedAndIncomplete(t *testing.T) { + for _, kind := range []string{"magic", "version", "zeroBlock", "hugeBitmap", "shortBitmap", "missingBlock", "truncated"} { + t.Run(kind, func(t *testing.T) { + name := fixture(t, []byte("abcdefgh")) + b, _ := os.ReadFile(name) + switch kind { + case "magic": + b[0] = 'X' + case "version": + b[4] = 9 + case "zeroBlock": + binary.LittleEndian.PutUint64(b[12:], 0) + case "hugeBitmap": + binary.LittleEndian.PutUint64(b[28:], 1<<63) + case "shortBitmap": + b = b[:36] + case "missingBlock": + b[36] = 1 + case "truncated": + b = b[:len(b)-1] + } + os.WriteFile(name, b, 0600) + if p, err := Open(name); err == nil { + p.File.Close() + t.Fatal("accepted invalid cache") + } + }) + } +} + +func TestPinRejectsUnsafeTree(t *testing.T) { + name := fixture(t, []byte("abcd")) + for _, p := range []string{"../x", "/x", "a/../x", "a\\x", ".", ""} { + if s, err := Pin(map[string]string{p: name}); err == nil { + s.Close() + t.Fatalf("accepted %q", p) + } + } + if s, err := Pin(map[string]string{"a": name, "a/b": name}); err == nil { + s.Close() + t.Fatal("accepted collision") + } +} + +func TestPinLargeSparsePayload(t *testing.T) { + if os.Getenv("DINGO_TEST_FUSE") != "1" { + t.Skip("opt-in large sparse file on isolated Linux filesystem") + } + const size = int64(512) * 1024 * 1024 * 1024 + const block = 4 * 1024 * 1024 + const bits = size / block + name := filepath.Join(t.TempDir(), "large") + f, err := os.Create(name) + if err != nil { + t.Fatal(err) + } + h := make([]byte, 36+bits/8) + copy(h, "OLAH") + binary.LittleEndian.PutUint64(h[4:], 8) + binary.LittleEndian.PutUint64(h[12:], block) + binary.LittleEndian.PutUint64(h[20:], uint64(size)) + binary.LittleEndian.PutUint64(h[28:], uint64(bits)) + for i := 36; i < len(h); i++ { + h[i] = 255 + } + if _, err = f.Write(h); err != nil { + f.Close() + t.Fatal(err) + } + if err = f.Truncate(int64(len(h)) + size); err != nil { + f.Close() + t.Fatal(err) + } + f.Close() + start := time.Now() + p, err := Open(name) + if err != nil { + t.Fatal(err) + } + defer p.File.Close() + if p.Size != size { + t.Fatal(p.Size) + } + t.Logf("512 GiB logical payload pinned in %s; header bytes=%d (sparse fixture, not a throughput test)", time.Since(start), len(h)) +} diff --git a/pkg/cachemount/tree.go b/pkg/cachemount/tree.go new file mode 100644 index 0000000..3fb068e --- /dev/null +++ b/pkg/cachemount/tree.go @@ -0,0 +1,350 @@ +package cachemount + +import ( + "encoding/hex" + "encoding/json" + "fmt" + "io" + "io/fs" + "os" + "path" + "path/filepath" + "sort" + "strings" + + "dingospeed/pkg/hfprojection" + "dingospeed/pkg/repository" +) + +// Selection is an administrator-selected local model version. This service is +// local filesystem access, not a substitute for the web console's user ACLs. +type Selection struct { + CacheRoot string `json:"cache_root"` + Namespace string `json:"namespace"` + Repo string `json:"repo"` + Revision string `json:"revision"` +} +type TreeFile struct { + Path string `json:"path"` + Size int64 `json:"size"` + ContentID string `json:"content_id"` + Backing string `json:"backing"` +} +type Version struct { + Selection + Commit string `json:"commit"` + SourceFilePath string `json:"source_file_path"` + Status string `json:"status"` + Error string `json:"error,omitempty"` + Files []TreeFile `json:"files"` +} + +func readJSON(root, p string, out any) error { + if err := repository.SafePath(root, p); err != nil { + return err + } + f, err := os.Open(p) + if err != nil { + return err + } + defer f.Close() + b, err := io.ReadAll(io.LimitReader(f, 32*1024*1024+1)) + if err != nil { + return err + } + if len(b) > 32*1024*1024 { + return fmt.Errorf("metadata exceeds 32 MiB") + } + return json.Unmarshal(b, out) +} +func commitAt(root, p string) (string, error) { + var envelope struct { + Status int `json:"status_code"` + Content string `json:"content"` + } + if err := readJSON(root, p, &envelope); err != nil { + return "", err + } + if envelope.Status != 200 { + return "", fmt.Errorf("metadata status %d", envelope.Status) + } + b, err := hex.DecodeString(envelope.Content) + if err != nil { + return "", err + } + var meta struct { + SHA string `json:"sha"` + } + if err = json.Unmarshal(b, &meta); err != nil { + return "", err + } + return meta.SHA, repository.Segment(meta.SHA) +} + +// Discover uses registered identities, legacy upload manifests and HF/ModelScope cached +// metadata. It never registers repositories or downloads content. +func Discover(root string) ([]Selection, error) { + root, err := filepath.Abs(root) + if err != nil { + return nil, err + } + if _, err = os.Stat(root); err != nil { + return nil, err + } + descriptors, err := repository.List(root) + if err != nil { + return nil, err + } + keys := map[string]repository.RepoKey{} + for _, d := range descriptors { + if d.RepoType == "models" && d.Source == "hosted" { + keys[d.ID()] = d.RepoKey + } + } + local := filepath.Join(root, "api", "models", repository.Local) + err = filepath.WalkDir(local, func(p string, e fs.DirEntry, walkErr error) error { + if walkErr != nil { + return walkErr + } + if e.IsDir() { + return nil + } + if e.Name() != "dingo-local-manifest.json" { + return nil + } + repoRoot := filepath.Dir(filepath.Dir(filepath.Dir(p))) + for _, k := range keys { + if filepath.Clean(k.APIRoot(root)) == repoRoot { + return nil + } + } + rel, err := filepath.Rel(local, repoRoot) + if err != nil { + return err + } + k := repository.RepoKey{Namespace: repository.Local, RepoType: "models", Repo: filepath.ToSlash(rel)} + if err = k.Validate(); err != nil { + return err + } + keys[k.ID()] = k + return nil + }) + if err != nil && !os.IsNotExist(err) { + return nil, err + } + out := []Selection{} + for _, k := range keys { + entries, err := os.ReadDir(filepath.Join(k.APIRoot(root), "revision")) + if err != nil { + return nil, err + } + revisions := map[string]string{} + pointed := map[string]bool{} + for _, e := range entries { + if !e.IsDir() { + continue + } + p := k.Revision(root, e.Name()) + if _, err := os.Stat(filepath.Join(p, "dingo-superseded.json")); err == nil { + continue + } else if !os.IsNotExist(err) { + return nil, err + } + commit, err := commitAt(root, filepath.Join(p, "meta_get.json")) + if os.IsNotExist(err) { + continue + } + if err != nil { + return nil, err + } + revisions[e.Name()] = commit + if e.Name() != commit { + pointed[commit] = true + } + } + for rev, c := range revisions { + if rev == c && pointed[c] { + continue + } + out = append(out, Selection{root, k.Namespace, k.Repo, rev}) + } + } + for _, namespace := range []string{repository.HuggingFace, repository.ModelScope} { + catalog, err := (hfprojection.Reader{Root: root, Namespace: namespace}).Catalog("models") + if err != nil { + return nil, err + } + for _, repo := range catalog.Repos { + if repo.Error != "" { + return nil, fmt.Errorf("%s/%s: %s", namespace, repo.Repo, repo.Error) + } + for _, rev := range repo.Revisions { + out = append(out, Selection{root, namespace, repo.Repo, rev.Name}) + } + } + } + sort.Slice(out, func(i, j int) bool { + return out[i].Namespace+"/"+out[i].Repo+"/"+out[i].Revision < out[j].Namespace+"/"+out[j].Repo+"/"+out[j].Revision + }) + return out, nil +} + +func resolve(s Selection) (Version, error) { + v := Version{Selection: s, Status: "unavailable", Files: []TreeFile{}} + k := repository.RepoKey{Namespace: s.Namespace, RepoType: "models", Repo: s.Repo} + if err := k.Validate(); err != nil { + return v, err + } + if err := repository.Segment(s.Revision); err != nil { + return v, err + } + if !filepath.IsAbs(s.CacheRoot) { + return v, fmt.Errorf("cache_root must be absolute") + } + if s.Namespace == repository.HuggingFace || s.Namespace == repository.ModelScope { + m, paths, err := (hfprojection.Reader{Root: s.CacheRoot, Namespace: s.Namespace}).MountSources("models", s.Repo, s.Revision) + if err != nil { + return v, err + } + v.Commit = m.Commit + for _, f := range m.Files { + v.Files = append(v.Files, TreeFile{f.Path, f.Size, f.OID, paths[f.Path]}) + } + } else { + if s.Namespace != repository.Local { + d, err := repository.Read(s.CacheRoot, k) + if err != nil { + return v, err + } + if d.Source != "hosted" { + return v, fmt.Errorf("not a hosted repository") + } + } + commit, err := commitAt(s.CacheRoot, filepath.Join(k.Revision(s.CacheRoot, s.Revision), "meta_get.json")) + if err != nil { + return v, err + } + v.Commit = commit + var manifest []struct { + Path string `json:"path"` + SHA string `json:"sha256"` + Size int64 `json:"size"` + } + if err = readJSON(s.CacheRoot, filepath.Join(k.Revision(s.CacheRoot, commit), "dingo-local-manifest.json"), &manifest); err != nil { + return v, err + } + if manifest == nil { + return v, fmt.Errorf("null manifest") + } + for _, f := range manifest { + digest, err := hex.DecodeString(f.SHA) + if err != nil || len(digest) != 32 || f.Size < 0 { + return v, fmt.Errorf("invalid manifest entry %s", f.Path) + } + p := k.Blob(s.CacheRoot, f.SHA) + if err = repository.SafePath(s.CacheRoot, p); err != nil { + return v, err + } + v.Files = append(v.Files, TreeFile{f.Path, f.Size, f.SHA, p}) + } + } + seen := map[string]bool{} + for _, f := range v.Files { + if err := repository.Relative(f.Path); err != nil { + return v, err + } + if seen[f.Path] { + return v, fmt.Errorf("duplicate path %s", f.Path) + } + seen[f.Path] = true + } + sort.Slice(v.Files, func(i, j int) bool { return v.Files[i].Path < v.Files[j].Path }) + return v, nil +} + +// Historical HF caches can contain ordinary files. Accept those only with a +// known matching size; malformed OLAH must never fall back to ordinary bytes. +func openTreePayload(name string, size int64, allowPlain bool) (*Payload, error) { + p, err := Open(name) + if err == nil || !allowPlain || size < 0 { + return p, err + } + f, openErr := os.Open(name) + if openErr != nil { + return nil, openErr + } + var magic [4]byte + _, _ = f.ReadAt(magic[:], 0) + st, statErr := f.Stat() + if statErr != nil || !st.Mode().IsRegular() || st.Size() != size || string(magic[:]) == "OLAH" { + f.Close() + return nil, err + } + return &Payload{File: f, Size: size}, nil +} + +// PrepareTree pins every ready version atomically. Unavailable versions remain +// visible as error directories, never as misleading partial/empty models. +func PrepareTree(selections []Selection, mountpoint string) (*Snapshot, []Version, error) { + snapshot := &Snapshot{Files: map[string]*Payload{}, Directories: map[string]string{}} + versions := []Version{} + prefixes := []string{} + for _, s := range selections { + k := repository.RepoKey{Namespace: s.Namespace, RepoType: "models", Repo: s.Repo} + if err := k.Validate(); err != nil { + snapshot.Close() + return nil, nil, err + } + if err := repository.Segment(s.Revision); err != nil { + snapshot.Close() + return nil, nil, err + } + prefix := path.Join(s.Namespace, s.Repo, s.Revision) + for _, old := range prefixes { + if prefix == old || strings.HasPrefix(prefix, old+"/") || strings.HasPrefix(old, prefix+"/") { + snapshot.Close() + return nil, nil, fmt.Errorf("ambiguous version paths %s and %s", old, prefix) + } + } + prefixes = append(prefixes, prefix) + v, err := resolve(s) + v.SourceFilePath = filepath.Join(mountpoint, filepath.FromSlash(prefix)) + opened := map[string]*Payload{} + if err == nil { + for _, f := range v.Files { + var p *Payload + p, err = openTreePayload(f.Backing, f.Size, s.Namespace == repository.HuggingFace) + if err != nil { + break + } + opened[f.Path] = p + if f.Size >= 0 && p.Size != f.Size { + err = fmt.Errorf("size mismatch: %s", f.Path) + break + } + } + for name := range opened { + for parent := path.Dir(name); parent != "."; parent = path.Dir(parent) { + if _, ok := opened[parent]; ok { + err = fmt.Errorf("file/directory collision: %s", parent) + } + } + } + } + if err != nil { + for _, p := range opened { + p.File.Close() + } + v.Error = err.Error() + snapshot.Directories[prefix] = v.Error + } else { + v.Status = "ready" + snapshot.Directories[prefix] = "" + for name, p := range opened { + snapshot.Files[prefix+"/"+name] = p + } + } + versions = append(versions, v) + } + return snapshot, versions, nil +} diff --git a/pkg/cachemount/tree_test.go b/pkg/cachemount/tree_test.go new file mode 100644 index 0000000..41af311 --- /dev/null +++ b/pkg/cachemount/tree_test.go @@ -0,0 +1,193 @@ +package cachemount + +import ( + "bytes" + "encoding/hex" + "encoding/json" + "os" + "path/filepath" + "testing" + + "dingospeed/pkg/repository" +) + +func treeJSON(t *testing.T, p string, v any) { + t.Helper() + if err := os.MkdirAll(filepath.Dir(p), 0700); err != nil { + t.Fatal(err) + } + b, err := json.Marshal(v) + if err != nil { + t.Fatal(err) + } + if err = os.WriteFile(p, b, 0600); err != nil { + t.Fatal(err) + } +} +func hostedTree(t *testing.T) (Selection, string) { + t.Helper() + root := t.TempDir() + k := repository.RepoKey{Namespace: "datacanvas", RepoType: "models", Repo: "team/model"} + treeJSON(t, filepath.Join(k.APIRoot(root), repository.Marker), repository.Hosted(k)) + body, _ := json.Marshal(map[string]any{"sha": "commit1", "siblings": []any{map[string]any{"rfilename": "nested/config.json"}}}) + meta := map[string]any{"status_code": 200, "content": hex.EncodeToString(body)} + treeJSON(t, filepath.Join(k.Revision(root, "main"), "meta_get.json"), meta) + treeJSON(t, filepath.Join(k.Revision(root, "commit1"), "meta_get.json"), meta) + sha := "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa" + treeJSON(t, filepath.Join(k.Revision(root, "commit1"), "dingo-local-manifest.json"), []any{map[string]any{"path": "nested/config.json", "sha256": sha, "size": 8}}) + data, err := os.ReadFile(fixture(t, []byte("abcdefgh"))) + if err != nil { + t.Fatal(err) + } + p := k.Blob(root, sha) + if err = os.MkdirAll(filepath.Dir(p), 0700); err != nil { + t.Fatal(err) + } + if err = os.WriteFile(p, data, 0600); err != nil { + t.Fatal(err) + } + return Selection{root, k.Namespace, k.Repo, "main"}, p +} + +func TestTreeDiscoveryAndPinnedVersion(t *testing.T) { + s, _ := hostedTree(t) + choices, err := Discover(s.CacheRoot) + if err != nil { + t.Fatal(err) + } + if len(choices) != 1 || choices[0] != s { + t.Fatalf("unexpected choices: %+v", choices) + } + snapshot, versions, err := PrepareTree(choices, "/mnt/dingo-models") + if err != nil { + t.Fatal(err) + } + defer snapshot.Close() + if versions[0].Status != "ready" || versions[0].SourceFilePath != filepath.Join("/mnt/dingo-models", "datacanvas", "team", "model", "main") { + t.Fatal(versions) + } + // Moving the branch metadata must not change the already mounted snapshot. + k := repository.RepoKey{Namespace: s.Namespace, RepoType: "models", Repo: s.Repo} + body, _ := json.Marshal(map[string]string{"sha": "anothercommit"}) + treeJSON(t, filepath.Join(k.Revision(s.CacheRoot, "main"), "meta_get.json"), map[string]any{"status_code": 200, "content": hex.EncodeToString(body)}) + b := make([]byte, 8) + _, err = snapshot.Files["datacanvas/team/model/main/nested/config.json"].ReadAt(b, 0) + if err != nil || !bytes.Equal(b, []byte("abcdefgh")) { + t.Fatalf("pinned bytes changed: %q %v", b, err) + } +} + +func TestTreeIncompleteVersionIsAtomic(t *testing.T) { + s, p := hostedTree(t) + if err := os.Truncate(p, 38); err != nil { + t.Fatal(err) + } + snapshot, versions, err := PrepareTree([]Selection{s}, "/mnt/dingo-models") + if err != nil { + t.Fatal(err) + } + defer snapshot.Close() + if versions[0].Status != "unavailable" || len(snapshot.Files) != 0 || snapshot.Directories["datacanvas/team/model/main"] == "" { + t.Fatalf("partial version exposed: %+v", versions) + } +} + +func TestTreeRejectsAmbiguousAndUnsafeSelections(t *testing.T) { + s, _ := hostedTree(t) + duplicate := []Selection{s, s} + if snapshot, _, err := PrepareTree(duplicate, "/mnt/dingo-models"); err == nil { + snapshot.Close() + t.Fatal("accepted duplicate") + } + s.Revision = "../escape" + if snapshot, _, err := PrepareTree([]Selection{s}, "/mnt/dingo-models"); err == nil { + snapshot.Close() + t.Fatal("accepted traversal") + } +} + +func TestPlainHFCannotHideCorruptContainer(t *testing.T) { + name := filepath.Join(t.TempDir(), "plain") + os.WriteFile(name, []byte("{}"), 0600) + p, err := openTreePayload(name, 2, true) + if err != nil { + t.Fatal(err) + } + p.File.Close() + if p, err := openTreePayload(name, 3, true); err == nil { + p.File.Close() + t.Fatal("accepted wrong size") + } + os.WriteFile(name, []byte("OLAH"), 0600) + if p, err := openTreePayload(name, 4, true); err == nil { + p.File.Close() + t.Fatal("corrupt OLAH accepted as raw") + } +} + +// Both remote providers use normalized cached metadata, but distinct payload roots. +func TestRemoteProviderTree(t *testing.T) { + for _, namespace := range []string{repository.HuggingFace, repository.ModelScope} { + t.Run(namespace, func(t *testing.T) { + root := t.TempDir() + k := repository.RepoKey{Namespace: namespace, RepoType: "models", Repo: "team/model"} + treeJSON(t, filepath.Join(k.APIRoot(root), repository.Marker), repository.Remote(k)) + body, _ := json.Marshal(map[string]any{"sha": "commit1", "siblings": []any{map[string]any{"rfilename": "nested/config.json", "size": 8, "blobId": "blob1"}}}) + for _, rev := range []string{"master", "commit1"} { + treeJSON(t, filepath.Join(k.Revision(root, rev), "meta_get.json"), map[string]any{"status_code": 200, "content": hex.EncodeToString(body)}) + } + backing := k.Blob(root, "blob1") + data, err := os.ReadFile(fixture(t, []byte("abcdefgh"))) + if err != nil { + t.Fatal(err) + } + if err = os.MkdirAll(filepath.Dir(backing), 0700); err != nil { + t.Fatal(err) + } + if err = os.WriteFile(backing, data, 0600); err != nil { + t.Fatal(err) + } + selections, err := Discover(root) + if err != nil { + t.Fatal(err) + } + if len(selections) != 1 || selections[0].Namespace != namespace || selections[0].Revision != "master" { + t.Fatalf("unexpected discovery: %+v", selections) + } + snapshot, versions, err := PrepareTree(selections, filepath.Join(root, "view")) + if err != nil { + t.Fatal(err) + } + prefix := namespace + "/team/model/master" + if len(versions) != 1 || versions[0].Status != "ready" || versions[0].Commit != "commit1" { + snapshot.Close() + t.Fatalf("not ready: %+v", versions) + } + if versions[0].Files[0].Backing != backing { + t.Fatalf("wrong backing: %+v", versions) + } + payload := snapshot.Files[prefix+"/nested/config.json"] + if payload == nil { + snapshot.Close() + t.Fatal("missing virtual file") + } + b := make([]byte, 8) + _, err = payload.ReadAt(b, 0) + snapshot.Close() + if err != nil || string(b) != "abcdefgh" { + t.Fatalf("payload: %q %v", b, err) + } + if err = os.Remove(backing); err != nil { + t.Fatal(err) + } + snapshot, versions, err = PrepareTree(selections, filepath.Join(root, "view")) + if err != nil { + t.Fatal(err) + } + defer snapshot.Close() + if versions[0].Status != "unavailable" || len(snapshot.Files) != 0 { + t.Fatalf("exposed incomplete version: %+v", versions) + } + }) + } +} diff --git a/pkg/common/cache_transfer_trace.go b/pkg/common/cache_transfer_trace.go new file mode 100644 index 0000000..11e2781 --- /dev/null +++ b/pkg/common/cache_transfer_trace.go @@ -0,0 +1,69 @@ +package common + +import ( + "context" + "github.com/prometheus/client_golang/prometheus" + "github.com/prometheus/client_golang/prometheus/promauto" + "sync/atomic" +) + +type cacheTraceKey struct{} +type CacheTransferTrace struct { + Requests atomic.Int64 + Canceled atomic.Int64 + Active atomic.Int64 + WrittenBytes atomic.Int64 + stopped atomic.Bool +} + +var cacheSourceCanceled = promauto.NewCounter(prometheus.CounterOpts{Name: "dingospeed_cache_source_cancellations_total", Help: "Source requests canceled by cache tasks"}) +var cacheAfterStop = promauto.NewCounterVec(prometheus.CounterOpts{Name: "dingospeed_cache_activity_after_stop_total", Help: "Unexpected cache activity after the completion barrier"}, []string{"operation"}) + +func WithCacheTrace(ctx context.Context, trace *CacheTransferTrace) context.Context { + return context.WithValue(ctx, cacheTraceKey{}, trace) +} +func CacheTrace(ctx context.Context) *CacheTransferTrace { + trace, _ := ctx.Value(cacheTraceKey{}).(*CacheTransferTrace) + return trace +} +func (trace *CacheTransferTrace) Begin() bool { + if trace == nil { + return true + } + if trace.stopped.Load() { + cacheAfterStop.WithLabelValues("request").Inc() + return false + } + trace.Active.Add(1) + if trace.stopped.Load() { + trace.Active.Add(-1) + cacheAfterStop.WithLabelValues("request").Inc() + return false + } + trace.Requests.Add(1) + return true +} +func (trace *CacheTransferTrace) End(canceled bool) { + if trace == nil { + return + } + trace.Active.Add(-1) + if canceled { + trace.Canceled.Add(1) + cacheSourceCanceled.Inc() + } +} +func (trace *CacheTransferTrace) Wrote(bytes int64) { + if trace == nil { + return + } + trace.WrittenBytes.Add(bytes) + if trace.stopped.Load() { + cacheAfterStop.WithLabelValues("write").Inc() + } +} +func (trace *CacheTransferTrace) Stop() { + if trace != nil { + trace.stopped.Store(true) + } +} diff --git a/pkg/common/local_transfer.go b/pkg/common/local_transfer.go new file mode 100644 index 0000000..168c36c --- /dev/null +++ b/pkg/common/local_transfer.go @@ -0,0 +1,73 @@ +package common + +import ( + "crypto/rand" + "encoding/hex" + "sync" +) + +// LocalTransfer joins a loopback handler, not merely its HTTP client. A canceled +// client must still wait for the handler's cache producers to release their files. +const LocalTransferHeader = "X-Dingo-Local-Transfer" + +var localTransfers sync.Map + +type LocalTransfer struct { + Trace *CacheTransferTrace + mu sync.Mutex + claimed, closed bool + done chan struct{} + id string +} + +func NewLocalTransfer(traces ...*CacheTransferTrace) *LocalTransfer { + var b [24]byte + if _, err := rand.Read(b[:]); err != nil { + panic(err) + } + t := &LocalTransfer{id: hex.EncodeToString(b[:]), done: make(chan struct{})} + if len(traces) > 0 { + t.Trace = traces[0] + } + localTransfers.Store(t.id, t) + return t +} +func (t *LocalTransfer) ID() string { return t.id } +func ClaimLocalTransfer(id string) (func(), bool) { + value, ok := localTransfers.Load(id) + if !ok { + return nil, false + } + t := value.(*LocalTransfer) + t.mu.Lock() + defer t.mu.Unlock() + if t.closed || t.claimed { + return nil, false + } + t.claimed = true + return func() { + t.mu.Lock() + defer t.mu.Unlock() + if !t.closed { + t.closed = true + close(t.done) + } + }, true +} +func (t *LocalTransfer) Wait() { + t.mu.Lock() + if !t.claimed && !t.closed { + t.closed = true + close(t.done) + } + t.mu.Unlock() + <-t.done + localTransfers.Delete(t.id) +} + +func LocalCacheTrace(id string) *CacheTransferTrace { + if value, ok := localTransfers.Load(id); ok { + return value.(*LocalTransfer).Trace + } + return nil +} diff --git a/pkg/common/pool.go b/pkg/common/pool.go index a8aa326..a3babcc 100644 --- a/pkg/common/pool.go +++ b/pkg/common/pool.go @@ -16,6 +16,7 @@ package common import ( "context" + "errors" "sync" "time" @@ -75,7 +76,7 @@ func (p *Pool) worker() { } task.DoTask() if p.persist { - p.taskMap.Delete(task.GetTaskNo()) + p.release(task) } } } @@ -89,43 +90,59 @@ func (p *Pool) GetTask(taskNo int) (Task, bool) { return p.taskMap.Get(taskNo) } -// Submit 提交任务 -func (p *Pool) Submit(ctx context.Context, task Task) error { - p.mu.Lock() - defer p.mu.Unlock() - if p.persist && p.exist(task.GetTaskNo()) { - return nil - } - select { - case p.taskChan <- task: - if p.persist { - p.taskMap.Set(task.GetTaskNo(), task) - } - return nil - case <-ctx.Done(): - return nil +func (p *Pool) ActiveCount() int { + if p.taskMap == nil { + return 0 } + return p.taskMap.Len() } -func (p *Pool) SubmitForTimeout(ctx context.Context, task Task) error { +var ErrTaskActive = errors.New("task already active") + +func (p *Pool) release(task Task) { p.mu.Lock() defer p.mu.Unlock() - if p.persist && p.exist(task.GetTaskNo()) { - return nil + if current, ok := p.taskMap.Get(task.GetTaskNo()); ok && current == task { + p.taskMap.Delete(task.GetTaskNo()) + } +} + +func (p *Pool) submit(ctx context.Context, task Task, timeout <-chan time.Time) error { + if err := ctx.Err(); err != nil { + return err + } + if p.persist { + p.mu.Lock() + if p.exist(task.GetTaskNo()) { + p.mu.Unlock() + return ErrTaskActive + } + p.taskMap.Set(task.GetTaskNo(), task) + p.mu.Unlock() } select { case p.taskChan <- task: + return nil + case <-timeout: if p.persist { - p.taskMap.Set(task.GetTaskNo(), task) + p.release(task) } - return nil - case <-time.After(3 * time.Second): return myerr.New(consts.TaskMoreErrMsg) case <-ctx.Done(): - return nil + if p.persist { + p.release(task) + } + return ctx.Err() } } +func (p *Pool) Submit(ctx context.Context, task Task) error { return p.submit(ctx, task, nil) } +func (p *Pool) SubmitForTimeout(ctx context.Context, task Task) error { + timer := time.NewTimer(3 * time.Second) + defer timer.Stop() + return p.submit(ctx, task, timer.C) +} + // Close 关闭协程池（安全关闭） func (p *Pool) Close() { close(p.taskChan) diff --git a/pkg/common/pool_lifecycle_test.go b/pkg/common/pool_lifecycle_test.go new file mode 100644 index 0000000..2f152d4 --- /dev/null +++ b/pkg/common/pool_lifecycle_test.go @@ -0,0 +1,67 @@ +package common + +import ( + "context" + "errors" + "testing" + "time" +) + +type lifecycleTask struct { + id int + entered, release chan struct{} +} + +func (t *lifecycleTask) GetTaskNo() int { return t.id } +func (t *lifecycleTask) GetCancelFun() context.CancelFunc { return func() {} } +func (t *lifecycleTask) DoTask() { close(t.entered); <-t.release } +func TestPoolRejectsDuplicateAndCanceledSubmission(t *testing.T) { + p := NewPool(1, true) + defer p.Close() + first := &lifecycleTask{1, make(chan struct{}), make(chan struct{})} + if err := p.Submit(context.Background(), first); err != nil { + t.Fatal(err) + } + defer close(first.release) + <-first.entered + second := &lifecycleTask{1, make(chan struct{}), make(chan struct{})} + if err := p.Submit(context.Background(), second); !errors.Is(err, ErrTaskActive) { + t.Fatalf("duplicate: %v", err) + } + ctx, cancel := context.WithCancel(context.Background()) + cancel() + if err := p.Submit(ctx, &lifecycleTask{id: 2}); !errors.Is(err, context.Canceled) { + t.Fatalf("canceled: %v", err) + } + if p.ActiveCount() != 1 { + t.Fatal("failed admission leaked a registration") + } +} +func TestLocalTransferWaitsForClaimedHandler(t *testing.T) { + transfer := NewLocalTransfer() + finish, ok := ClaimLocalTransfer(transfer.ID()) + if !ok { + t.Fatal("claim failed") + } + waited := make(chan struct{}) + go func() { transfer.Wait(); close(waited) }() + select { + case <-waited: + t.Fatal("client returned before handler finished") + default: + } + finish() + select { + case <-waited: + case <-time.After(time.Second): + t.Fatal("completion not delivered") + } + if _, ok := ClaimLocalTransfer(transfer.ID()); ok { + t.Fatal("completed request reclaimed") + } + canceled := NewLocalTransfer() + canceled.Wait() + if _, ok := ClaimLocalTransfer(canceled.ID()); ok { + t.Fatal("late request started after cancellation") + } +} diff --git a/pkg/config/config.go b/pkg/config/config.go index 93cdff4..64c0974 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -17,6 +17,7 @@ package config import ( "errors" "fmt" + "net/url" "os" "path/filepath" "sync" @@ -36,20 +37,31 @@ var SysConfig *Config var SystemInfo *model.SystemInfo type Config struct { - Id int32 - Server ServerConfig `json:"server" yaml:"server"` - Download Download `json:"download" yaml:"download"` - Cache Cache `json:"cache" yaml:"cache"` - Log LogConfig `json:"log" yaml:"log"` - Retry Retry `json:"retry" yaml:"retry"` - TokenBucketLimit TokenBucketLimit `json:"tokenBucketLimit" yaml:"tokenBucketLimit"` - DiskClean DiskClean `json:"diskClean" yaml:"diskClean"` - DynamicProxy DynamicProxy `json:"dynamicProxy" yaml:"dynamicProxy"` - ProxyPool ProxyPool `json:"proxyPool" yaml:"proxyPool"` - Scheduler Scheduler `json:"scheduler" yaml:"scheduler"` - Upload Upload `json:"upload" yaml:"upload"` - mu sync.RWMutex - Modelscope Modelscope `yaml:"modelscope"` + Id int32 + Server ServerConfig `json:"server" yaml:"server"` + Download Download `json:"download" yaml:"download"` + Cache Cache `json:"cache" yaml:"cache"` + Log LogConfig `json:"log" yaml:"log"` + Retry Retry `json:"retry" yaml:"retry"` + TokenBucketLimit TokenBucketLimit `json:"tokenBucketLimit" yaml:"tokenBucketLimit"` + DiskClean DiskClean `json:"diskClean" yaml:"diskClean"` + DynamicProxy DynamicProxy `json:"dynamicProxy" yaml:"dynamicProxy"` + ProxyPool ProxyPool `json:"proxyPool" yaml:"proxyPool"` + Scheduler Scheduler `json:"scheduler" yaml:"scheduler"` + Upload Upload `json:"upload" yaml:"upload"` + StorageProbe StorageProbe `json:"storageProbe" yaml:"storageProbe"` + mu sync.RWMutex + registrationOnce sync.Once + registration *Registration + registrationErr error + registrationPersisted bool + registrationStatus RegistrationStatus + Modelscope Modelscope `yaml:"modelscope"` +} + +// StorageProbe is opt-in because probing adds filesystem I/O. +type StorageProbe struct { + Enabled bool `json:"enabled" yaml:"enabled"` } type ServerConfig struct { @@ -104,6 +116,7 @@ type ReadBlock struct { } type Scheduler struct { + HTTPURL string `json:"httpUrl" yaml:"httpUrl"` OriginMode string `json:"-" yaml:"-"` // 记录原始状态 Mode string `json:"mode" yaml:"mode"` Addr string `json:"addr" yaml:"addr"` @@ -146,7 +159,7 @@ type DiskClean struct { Enabled bool `json:"enabled" yaml:"enabled"` CacheSizeLimit int64 `json:"cacheSizeLimit" yaml:"cacheSizeLimit"` CacheCleanStrategy string `json:"cacheCleanStrategy" yaml:"cacheCleanStrategy"` - CollectTimePeriod int `json:"collectTimePeriod" yaml:"collectTimePeriod" validate:"min=1,max=600"` // 周期采集内存使用量，单位秒 + CollectTimePeriod int `json:"collectTimePeriod" yaml:"collectTimePeriod" validate:"min=1,max=600"` // 周期检测磁盘使用量，单位小时 InstanceID string `json:"instanceID" yaml:"instanceID"` } @@ -200,7 +213,36 @@ type Modelscope struct { RetryDelay int `yaml:"retryDelay"` } +// ModelScope is available without a provider-specific YAML section. Explicit +// overrides are preserved; invalid values fail at startup instead of on download. +func (m *Modelscope) setDefaults() { + if m.OfficialBaseURL == "" { + m.OfficialBaseURL = "https://www.modelscope.cn" + } + if m.ChunkSize == 0 { + m.ChunkSize = 8388608 + } + if m.MaxRetry == 0 { + m.MaxRetry = 5 + } + if m.RetryDelay == 0 { + m.RetryDelay = 3 + } +} + +func (m Modelscope) validate() error { + u, err := url.Parse(m.OfficialBaseURL) + if err != nil || u.Hostname() == "" || (u.Scheme != "http" && u.Scheme != "https") || u.RawQuery != "" || u.Fragment != "" { + return fmt.Errorf("modelscope.officialBaseURL must be an absolute HTTP(S) base URL without query or fragment") + } + if m.ChunkSize < 1 || m.MaxRetry < 1 || m.RetryDelay < 0 { + return fmt.Errorf("modelscope requires chunkSize > 0, maxRetry > 0 and retryDelay >= 0") + } + return nil +} + type Upload struct { + AdvertiseURL string `json:"advertiseUrl" yaml:"advertiseUrl"` Host string `json:"host" yaml:"host"` Port int `json:"port" yaml:"port"` Namespace string `json:"namespace" yaml:"namespace"` @@ -464,7 +506,10 @@ func (c *Config) GetSchedulerModel() string { } func (c *Config) GetOriginSchedulerModel() string { - return c.Scheduler.OriginMode + if c.Registration().Enabled { + return consts.SchedulerModeCluster + } + return consts.SchedulerModeStandalone } func (c *Config) GetModelCacheRoot() string { @@ -476,6 +521,7 @@ func (c *Config) GetDatasetCacheRoot() string { } func (c *Config) SetDefaults() { + c.Modelscope.setDefaults() if c.Server.Port == 0 { c.Server.Port = 8090 } @@ -502,7 +548,7 @@ func (c *Config) SetDefaults() { c.Download.RespChanSize = 30 } if c.DiskClean.CollectTimePeriod == 0 { - c.DiskClean.CollectTimePeriod = 1 + c.DiskClean.CollectTimePeriod = 24 } c.Scheduler.OriginMode = c.Scheduler.Mode if c.Scheduler.LinkDomain == "" { @@ -552,6 +598,9 @@ func Scan(path string) (*Config, error) { return nil, err } c.SetDefaults() + if err := c.Modelscope.validate(); err != nil { + return nil, err + } if c.Download.RemoteFileRangeSize%c.Download.BlockSize != 0 { return nil, myerr.New("RemoteFileRangeSize must be a multiple of BlockSize") diff --git a/pkg/config/modelscope_defaults_test.go b/pkg/config/modelscope_defaults_test.go new file mode 100644 index 0000000..35dccb2 --- /dev/null +++ b/pkg/config/modelscope_defaults_test.go @@ -0,0 +1,71 @@ +package config + +import ( + "os" + "path/filepath" + "strings" + "testing" + + "gopkg.in/yaml.v3" +) + +func TestScanModelScopeOptionalConfiguration(t *testing.T) { + raw, err := os.ReadFile("../../config/config.yaml") + if err != nil { + t.Fatal(err) + } + var base map[string]interface{} + if err := yaml.Unmarshal(raw, &base); err != nil { + t.Fatal(err) + } + old, oldInfo := SysConfig, SystemInfo + t.Cleanup(func() { SysConfig, SystemInfo = old, oldInfo }) + for _, tc := range []struct { + name string + override interface{} + wantURL string + wantRetry int + invalid bool + }{ + {"omitted", nil, "https://www.modelscope.cn", 5, false}, + {"empty", map[string]interface{}{}, "https://www.modelscope.cn", 5, false}, + {"partial", map[string]interface{}{"officialBaseURL": "http://localhost:9000"}, "http://localhost:9000", 5, false}, + {"explicit", map[string]interface{}{"officialBaseURL": "https://modelscope.cn", "maxRetry": 2, "retryDelay": 7, "chunkSize": 4096}, "https://modelscope.cn", 2, false}, + {"negative retries", map[string]interface{}{"maxRetry": -1}, "", 0, true}, + {"negative delay", map[string]interface{}{"retryDelay": -1}, "", 0, true}, + {"negative chunk", map[string]interface{}{"chunkSize": -1}, "", 0, true}, + {"relative URL", map[string]interface{}{"officialBaseURL": "modelscope.cn"}, "", 0, true}, + {"unsupported scheme", map[string]interface{}{"officialBaseURL": "ftp://modelscope.cn"}, "", 0, true}, + } { + t.Run(tc.name, func(t *testing.T) { + delete(base, "modelscope") + if tc.override != nil { + base["modelscope"] = tc.override + } + data, err := yaml.Marshal(base) + if err != nil { + t.Fatal(err) + } + path := filepath.Join(t.TempDir(), "config.yaml") + if err := os.WriteFile(path, data, 0600); err != nil { + t.Fatal(err) + } + cfg, err := Scan(path) + if tc.invalid { + if err == nil || !strings.Contains(err.Error(), "modelscope") { + t.Fatalf("expected actionable startup error, got %v", err) + } + return + } + if err != nil { + t.Fatal(err) + } + if cfg.Modelscope.OfficialBaseURL != tc.wantURL || cfg.Modelscope.MaxRetry != tc.wantRetry || cfg.Modelscope.ChunkSize <= 0 || cfg.Modelscope.RetryDelay <= 0 { + t.Fatalf("unexpected provider config: %+v", cfg.Modelscope) + } + if tc.name == "explicit" && (cfg.Modelscope.RetryDelay != 7 || cfg.Modelscope.ChunkSize != 4096) { + t.Fatalf("override lost: %+v", cfg.Modelscope) + } + }) + } +} diff --git a/pkg/config/registration.go b/pkg/config/registration.go new file mode 100644 index 0000000..b7a19ad --- /dev/null +++ b/pkg/config/registration.go @@ -0,0 +1,185 @@ +package config + +import ( + "encoding/json" + "errors" + "fmt" + "net" + "net/url" + "os" + "path/filepath" + "regexp" + "strconv" + "strings" + "time" +) + +type Registration struct { + Enabled bool `json:"enabled"` + NodeID string `json:"nodeId"` + Address string `json:"address"` + HTTPURL string `json:"httpUrl,omitempty"` + Host string `json:"host"` + Port int `json:"port"` + ManagementURL string `json:"managementUrl,omitempty"` + DownloadURL string `json:"downloadUrl,omitempty"` +} +type RegistrationStatus struct { + State string `json:"state"` + Error string `json:"error,omitempty"` + CheckedAt time.Time `json:"checkedAt"` +} + +func ValidateRegistration(r Registration) error { + for _, raw := range []string{r.ManagementURL, r.DownloadURL, r.HTTPURL} { + if raw == "" { + continue + } + u, err := url.Parse(raw) + if err != nil || (u.Scheme != "http" && u.Scheme != "https") || u.Hostname() == "" || u.User != nil || u.RawQuery != "" || u.Fragment != "" { + return errors.New("invalid advertised URL") + } + } + if !regexp.MustCompile(`^[a-z0-9][a-z0-9-]{0,62}$`).MatchString(r.NodeID) { + return errors.New("invalid immutable node ID") + } + if !r.Enabled { + return nil + } + host, port, err := net.SplitHostPort(r.Address) + p, e := strconv.Atoi(port) + if err != nil || e != nil || host == "" || p < 1 || p > 65535 || strings.ContainsAny(host, "/\\ \t\r\n") { + return errors.New("scheduler address must be host:port") + } + if r.Host == "" || strings.ContainsAny(r.Host, "/\\ \t\r\n") || r.Port < 1 || r.Port > 65535 { + return errors.New("invalid advertised host/port") + } + return nil +} +func (c *Config) registrationPath() string { + return filepath.Join(c.Repos(), ".scheduler-registration.json") +} +func (c *Config) Registration() Registration { + c.registrationOnce.Do(func() { + r := Registration{Enabled: c.Scheduler.OriginMode == "cluster" || c.Scheduler.Mode == "cluster", NodeID: c.Scheduler.Discovery.InstanceId, Address: c.Scheduler.Addr, HTTPURL: c.Scheduler.HTTPURL, Host: c.Scheduler.Discovery.Host, Port: c.Scheduler.Discovery.Port, ManagementURL: c.Upload.AdvertiseURL} + b, err := os.ReadFile(c.registrationPath()) + if err == nil { + err = json.Unmarshal(b, &r) + if err == nil { + err = ValidateRegistration(r) + } + c.registrationPersisted = err == nil + } + if err != nil && !os.IsNotExist(err) { + c.registrationErr = err + r.Enabled = false + } + c.registration = &r + }) + c.mu.RLock() + defer c.mu.RUnlock() + return *c.registration +} +func (c *Config) SaveRegistration(r Registration) error { + if err := ValidateRegistration(r); err != nil { + return err + } + c.Registration() + c.mu.Lock() + defer c.mu.Unlock() + if c.registrationErr != nil { + return fmt.Errorf("cannot replace unreadable registration: %w", c.registrationErr) + } + if c.registration.NodeID != "" && c.registration.NodeID != r.NodeID { + return errors.New("node ID is already bound and cannot change") + } + if *c.registration == r && c.registrationPersisted { + return nil + } + b, err := json.Marshal(r) + if err != nil { + return err + } + f, err := os.CreateTemp(c.Repos(), ".scheduler-registration-*") + if err != nil { + return err + } + defer os.Remove(f.Name()) + if err = f.Chmod(0600); err == nil { + _, err = f.Write(b) + } + if err == nil { + err = f.Sync() + } + closeErr := f.Close() + if err == nil { + err = closeErr + } + if err != nil { + return err + } + if err = os.Rename(f.Name(), c.registrationPath()); err != nil { + return err + } + c.registration = &r + c.registrationPersisted = true + c.registrationStatus = RegistrationStatus{State: "connecting"} + c.Scheduler.Mode = "standalone" + return nil +} +func (c *Config) SetRegistrationStatus(r Registration, state string, err error) { + c.mu.Lock() + defer c.mu.Unlock() + if c.registration == nil || *c.registration != r { + return + } + s := RegistrationStatus{State: state, CheckedAt: time.Now().UTC()} + if err != nil { + s.Error = err.Error() + } + c.registrationStatus = s + if state == "connected" { + c.Scheduler.Mode = "cluster" + } else { + c.Scheduler.Mode = "standalone" + } +} +func (c *Config) RegistrationState() RegistrationStatus { + r := c.Registration() + c.mu.RLock() + defer c.mu.RUnlock() + if c.registrationErr != nil { + return RegistrationStatus{State: "error", Error: c.registrationErr.Error()} + } + if !r.Enabled { + return RegistrationStatus{State: "disabled"} + } + s := c.registrationStatus + if s.State == "" { + s.State = "connecting" + } + return s +} +func (c *Config) SetSchedulerID(id int32) { c.mu.Lock(); defer c.mu.Unlock(); c.Id = id } +func (c *Config) SchedulerID() int32 { c.mu.RLock(); defer c.mu.RUnlock(); return c.Id } + +// Explicit URLs support proxies and published container ports. The default uses +// the advertised node host, never the upload listen address (often 0.0.0.0). +func (c *Config) RegistrationEndpoints() (string, string) { + r := c.Registration() + management, download := r.ManagementURL, r.DownloadURL + if management == "" { + management = c.Upload.AdvertiseURL + } + if management == "" && r.Host != "" { + port := c.Upload.Port + if port == 0 { + port = 8091 + } + management = "http://" + net.JoinHostPort(r.Host, strconv.Itoa(port)) + } + if download == "" && r.Host != "" { + download = "http://" + net.JoinHostPort(r.Host, strconv.Itoa(r.Port)) + } + return management, download +} diff --git a/pkg/config/registration_endpoints_test.go b/pkg/config/registration_endpoints_test.go new file mode 100644 index 0000000..896e29d --- /dev/null +++ b/pkg/config/registration_endpoints_test.go @@ -0,0 +1,26 @@ +package config + +import "testing" + +func TestRegistrationAdvertisesManagementEndpoint(t *testing.T) { + c := &Config{} + c.Server.Repos = t.TempDir() + c.Upload.Port = 18091 + c.Scheduler.Discovery.Host = "2001:db8::1" + c.Scheduler.Discovery.Port = 8090 + management, download := c.RegistrationEndpoints() + if management != "http://[2001:db8::1]:18091" || download != "http://[2001:db8::1]:8090" { + t.Fatalf("%s %s", management, download) + } + r := c.Registration() + r.NodeID = "node-one" + r.ManagementURL = "https://proxy.example/manage" + r.DownloadURL = "https://proxy.example/download" + if err := c.SaveRegistration(r); err != nil { + t.Fatal(err) + } + management, download = c.RegistrationEndpoints() + if management != r.ManagementURL || download != r.DownloadURL { + t.Fatal("explicit published endpoints lost") + } +} diff --git a/pkg/config/registration_test.go b/pkg/config/registration_test.go new file mode 100644 index 0000000..0e4e2cc --- /dev/null +++ b/pkg/config/registration_test.go @@ -0,0 +1,63 @@ +package config + +import ( + "os" + "path/filepath" + "testing" +) + +func TestRegistrationPersistenceAndIdentity(t *testing.T) { + c := &Config{} + c.Server.Repos = t.TempDir() + r := Registration{Enabled: true, NodeID: "node-a", Address: "scheduler:19091", Host: "speed", Port: 8090} + if err := c.SaveRegistration(r); err != nil { + t.Fatal(err) + } + next := &Config{} + next.Server.Repos = c.Server.Repos + if next.Registration() != r { + t.Fatal("lost registration on restart") + } + r.NodeID = "node-b" + if next.SaveRegistration(r) == nil { + t.Fatal("changed bound ID") + } + r.NodeID = "node-a" + r.Enabled = false + if err := next.SaveRegistration(r); err != nil { + t.Fatal(err) + } + if next.GetOriginSchedulerModel() != "standalone" { + t.Fatal("disable not applied") + } +} +func TestRegistrationYAMLBindingAndCorruptFile(t *testing.T) { + c := &Config{} + c.Server.Repos = t.TempDir() + c.Scheduler.Discovery.InstanceId = "legacy" + if c.SaveRegistration(Registration{NodeID: "other"}) == nil { + t.Fatal("overrode YAML binding") + } + if err := c.SaveRegistration(Registration{NodeID: "legacy"}); err != nil { + t.Fatal(err) + } + if _, err := os.Stat(filepath.Join(c.Server.Repos, ".scheduler-registration.json")); err != nil { + t.Fatal(err) + } + dir := t.TempDir() + if err := os.WriteFile(filepath.Join(dir, ".scheduler-registration.json"), []byte("bad"), 0600); err != nil { + t.Fatal(err) + } + bad := &Config{} + bad.Server.Repos = dir + if bad.SaveRegistration(Registration{NodeID: "a"}) == nil { + t.Fatal("overwrote corrupt binding") + } +} +func TestRegistrationValidation(t *testing.T) { + for _, r := range []Registration{{NodeID: "../x"}, {Enabled: true, NodeID: "n", Address: "http://scheduler:90", Host: "speed", Port: 8090}, {Enabled: true, NodeID: "n", Address: "host:99999", Host: "speed", Port: 8090}} { + if ValidateRegistration(r) == nil { + t.Fatalf("accepted %+v", r) + } + } +} diff --git a/pkg/dependency/alerts_test.go b/pkg/dependency/alerts_test.go new file mode 100644 index 0000000..117fa3f --- /dev/null +++ b/pkg/dependency/alerts_test.go @@ -0,0 +1,53 @@ +package dependency + +import ( + "bytes" + "gopkg.in/yaml.v3" + "os" + "testing" + "text/template" +) + +// Validate the shipped YAML and notification templates locally. A deployment +// must additionally run promtool against its installed Prometheus version. +func TestAlertRuleFileAndTemplates(t *testing.T) { + contents, err := os.ReadFile("../../config/prometheus/dependency_alerts.yml") + if err != nil { + t.Fatal(err) + } + var rules struct { + Groups []struct { + Name string `yaml:"name"` + Rules []struct { + Alert string `yaml:"alert"` + Expr string `yaml:"expr"` + For string `yaml:"for"` + Labels map[string]string `yaml:"labels"` + Annotations map[string]string `yaml:"annotations"` + } `yaml:"rules"` + } `yaml:"groups"` + } + decoder := yaml.NewDecoder(bytes.NewReader(contents)) + decoder.KnownFields(true) + if err := decoder.Decode(&rules); err != nil { + t.Fatal(err) + } + if len(rules.Groups) != 1 || len(rules.Groups[0].Rules) != 2 { + t.Fatal("missing dependency rules") + } + for _, rule := range rules.Groups[0].Rules { + if rule.For == "" || rule.Labels["severity"] != "warning" { + t.Fatal("missing persistence or severity") + } + for _, value := range rule.Annotations { + tmpl, err := template.New(rule.Alert).Parse("{{$labels := .}}" + value) + if err != nil { + t.Fatal(err) + } + var out bytes.Buffer + if err := tmpl.Execute(&out, map[string]string{"dependency": "metadata_storage_read", "impact": "local_metadata_lookup", "instance": "test-node"}); err != nil { + t.Fatal(err) + } + } + } +} diff --git a/pkg/dependency/impact.go b/pkg/dependency/impact.go new file mode 100644 index 0000000..f48fa38 --- /dev/null +++ b/pkg/dependency/impact.go @@ -0,0 +1,80 @@ +package dependency + +import "time" + +type Assessment string + +const ( + ImpactUnknown Assessment = "unknown" + ImpactHealthy Assessment = "healthy_evidence" + ImpactPartial Assessment = "partial_evidence" + ImpactBusinessFailure Assessment = "business_failure" + ImpactBusinessFailureProbeHealthy Assessment = "business_failure_probe_healthy" + ImpactProbeFailure Assessment = "probe_failure" + ImpactCorrelatedFailure Assessment = "correlated_failure" + ImpactUncorrelatedFailures Assessment = "uncorrelated_failures" + ImpactUnresolved Assessment = "unresolved_evidence" +) + +var assessments = [...]Assessment{ImpactUnknown, ImpactHealthy, ImpactPartial, ImpactBusinessFailure, ImpactBusinessFailureProbeHealthy, ImpactProbeFailure, ImpactCorrelatedFailure, ImpactUncorrelatedFailures, ImpactUnresolved} + +// Impact is evidence for one storage operation, never a whole-node or mount +// diagnosis. Original observations are retained for explanation, not mutated. +type Impact struct { + Assessment Assessment + Business, Probe Snapshot + ProbeEnabled, ProbeBlocked bool +} + +// AssessStorage combines passive business evidence and an independent probe. +// The result is deliberately not an HTTP, readiness, retry or restart decision. +func AssessStorage(business, probe Snapshot, enabled, blocked bool, now time.Time) Impact { + fresh := func(s Snapshot) Snapshot { + if s.LastObservation.IsZero() || s.LastObservation.After(now) || now.Sub(s.LastObservation) > Freshness { + s.State = Unknown + } + return s + } + business = fresh(business) + if enabled { + probe = fresh(probe) + } else { + probe = Snapshot{} + blocked = false + } + r := Impact{Business: business, Probe: probe, ProbeEnabled: enabled, ProbeBlocked: blocked} + bad := func(s Snapshot) bool { return s.State == Suspect || s.State == Degraded } + bBad, pBad := bad(business), bad(probe) || blocked + bGood, pGood := business.State == ObservedHealthy && !business.Unresolved, probe.State == ObservedHealthy && !probe.Unresolved + switch { + case bBad && pBad: + probeFailure := probe.LastFailure + if blocked { + probeFailure = now + } // Still in flight is current evidence. + delta := business.LastFailure.Sub(probeFailure) + if delta < 0 { + delta = -delta + } + r.Assessment = ImpactUncorrelatedFailures + if !business.LastFailure.IsZero() && !probeFailure.IsZero() && delta <= EvidenceWindow { + r.Assessment = ImpactCorrelatedFailure + } + case bBad: + r.Assessment = ImpactBusinessFailure + if pGood && !business.LastFailure.IsZero() && !probe.LastObservation.Before(business.LastFailure) { + r.Assessment = ImpactBusinessFailureProbeHealthy + } + case pBad: + r.Assessment = ImpactProbeFailure + case business.Unresolved || probe.Unresolved: + r.Assessment = ImpactUnresolved + case bGood && pGood: + r.Assessment = ImpactHealthy + case bGood || pGood: + r.Assessment = ImpactPartial + default: + r.Assessment = ImpactUnknown + } + return r +} diff --git a/pkg/dependency/impact_metrics.go b/pkg/dependency/impact_metrics.go new file mode 100644 index 0000000..3b5b8ec --- /dev/null +++ b/pkg/dependency/impact_metrics.go @@ -0,0 +1,46 @@ +package dependency + +import ( + "github.com/prometheus/client_golang/prometheus" + "time" +) + +// ProbeEvidence must only read memory. A nil source means probes are disabled. +type ProbeEvidence func(ID, time.Time) (Snapshot, bool) + +type impactCollector struct { + business *Monitor + probe ProbeEvidence + impact, enabled *prometheus.Desc +} + +func ImpactCollector(business *Monitor, probe ProbeEvidence) prometheus.Collector { + return &impactCollector{business: business, probe: probe, + impact: prometheus.NewDesc("dingospeed_storage_impact", "Storage evidence assessment, one-hot. Not a whole-mount diagnosis or request policy.", []string{"operation", "assessment"}, nil), + enabled: prometheus.NewDesc("dingospeed_storage_impact_probe_enabled", "Whether active storage probe evidence is configured for this assessment.", nil, nil), + } +} +func (c *impactCollector) Describe(ch chan<- *prometheus.Desc) { ch <- c.impact; ch <- c.enabled } +func (c *impactCollector) Collect(ch chan<- prometheus.Metric) { + now := time.Now() + on := 0. + if c.probe != nil { + on = 1 + } + ch <- prometheus.MustNewConstMetric(c.enabled, prometheus.GaugeValue, on) + for id, operation := range []string{"read", "write"} { + var probe Snapshot + blocked := false + if c.probe != nil { + probe, blocked = c.probe(ID(id), now) + } + r := AssessStorage(c.business.Snapshot(ID(id), now), probe, c.probe != nil, blocked, now) + for _, assessment := range assessments { + value := 0. + if assessment == r.Assessment { + value = 1 + } + ch <- prometheus.MustNewConstMetric(c.impact, prometheus.GaugeValue, value, operation, string(assessment)) + } + } +} diff --git a/pkg/dependency/impact_test.go b/pkg/dependency/impact_test.go new file mode 100644 index 0000000..7a037af --- /dev/null +++ b/pkg/dependency/impact_test.go @@ -0,0 +1,142 @@ +package dependency + +import ( + "sync" + "testing" + "time" + + "github.com/prometheus/client_golang/prometheus" +) + +func TestStorageImpactMatrix(t *testing.T) { + now := time.Unix(1000, 0) + healthy := Snapshot{State: ObservedHealthy, LastObservation: now} + failure := Snapshot{State: Suspect, LastObservation: now, LastFailure: now, Unresolved: true} + recovering := failure + recovering.State = Recovering + stale := failure + stale.LastObservation = now.Add(-3 * time.Minute) + stale.LastFailure = stale.LastObservation + earlierHealthy := healthy + earlierHealthy.LastObservation = now.Add(-time.Second) + earlierFailure := failure + earlierFailure.LastObservation = now.Add(-40 * time.Second) + earlierFailure.LastFailure = earlierFailure.LastObservation + cases := []struct { + name string + business, probe Snapshot + enabled, blocked bool + want Assessment + }{ + {"nothing observed", Snapshot{}, Snapshot{}, false, false, ImpactUnknown}, + {"both healthy", healthy, healthy, true, false, ImpactHealthy}, + {"disabled is not healthy", healthy, healthy, false, false, ImpactPartial}, + {"no business traffic", Snapshot{}, healthy, true, false, ImpactPartial}, + {"business failed probe disabled", failure, failure, false, true, ImpactBusinessFailure}, + {"business failed probe unknown", failure, Snapshot{}, true, false, ImpactBusinessFailure}, + {"successful probe after failure", failure, healthy, true, false, ImpactBusinessFailureProbeHealthy}, + {"successful probe before failure", failure, earlierHealthy, true, false, ImpactBusinessFailure}, + {"probe failure only", healthy, failure, true, false, ImpactProbeFailure}, + {"temporally aligned failures", failure, failure, true, false, ImpactCorrelatedFailure}, + {"separate failure times", failure, earlierFailure, true, false, ImpactUncorrelatedFailures}, + {"stale is unresolved", stale, healthy, true, false, ImpactUnresolved}, + {"recovering is not recovered", recovering, healthy, true, false, ImpactUnresolved}, + {"expired but still blocked", healthy, stale, true, true, ImpactProbeFailure}, + {"new business error during old blockage", failure, stale, true, true, ImpactCorrelatedFailure}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + got := AssessStorage(tc.business, tc.probe, tc.enabled, tc.blocked, now) + if got.Assessment != tc.want { + t.Fatalf("got %+v want %s", got, tc.want) + } + }) + } +} + +func TestStorageImpactRecoveryAndExpiry(t *testing.T) { + business, probe := &Monitor{}, &Monitor{} + base := time.Unix(1000, 0) + for _, sec := range []int{0, 5, 10} { + at := base.Add(time.Duration(sec) * time.Second) + business.Observe(MetadataRead, false, at) + probe.Observe(MetadataRead, false, at) + } + assess := func(sec int) Assessment { + now := base.Add(time.Duration(sec) * time.Second) + return AssessStorage(business.Snapshot(MetadataRead, now), probe.Snapshot(MetadataRead, now), true, false, now).Assessment + } + if assess(10) != ImpactCorrelatedFailure { + t.Fatal("paired failures lost") + } + before := business.Snapshot(MetadataRead, base.Add(10*time.Second)) + for _, sec := range []int{15, 20, 25} { + probe.Observe(MetadataRead, true, base.Add(time.Duration(sec)*time.Second)) + } + if assess(25) != ImpactBusinessFailureProbeHealthy { + t.Fatal("probe recovery cleared business failure") + } + if business.Snapshot(MetadataRead, base.Add(25*time.Second)) != before { + t.Fatal("assessment modified source") + } + if assess(180) != ImpactUnresolved { + t.Fatal("no traffic treated as recovery") + } + if s := business.Snapshot(MetadataWrite, base); s.State != Unknown { + t.Fatal("read affected write") + } + if s := business.Snapshot(Scheduler, base); s.State != Unknown { + t.Fatal("storage affected scheduler") + } +} + +func TestImpactCollectorBoundedReadOnlyAndConcurrent(t *testing.T) { + business, probe := &Monitor{}, &Monitor{} + registry := prometheus.NewPedanticRegistry() + registry.MustRegister(ImpactCollector(business, func(id ID, now time.Time) (Snapshot, bool) { return probe.Snapshot(id, now), false })) + var wg sync.WaitGroup + for i := 0; i < 3; i++ { + wg.Add(1) + go func() { + defer wg.Done() + for j := 0; j < 100; j++ { + business.Observe(MetadataRead, j%2 == 0, time.Now()) + probe.Observe(MetadataWrite, j%2 == 0, time.Now()) + if _, err := registry.Gather(); err != nil { + t.Error(err) + } + } + }() + } + wg.Wait() + now := time.Now() + before := business.Snapshot(MetadataRead, now) + metrics, err := registry.Gather() + if err != nil || len(metrics) != 2 { + t.Fatalf("metrics: %v %d", err, len(metrics)) + } + for _, family := range metrics { + if family.GetName() != "dingospeed_storage_impact" { + continue + } + if len(family.Metric) != 18 { + t.Fatal("unbounded assessment labels") + } + sums := map[string]float64{} + for _, m := range family.Metric { + operation := "" + for _, l := range m.Label { + if l.GetName() == "operation" { + operation = l.GetValue() + } + } + sums[operation] += m.Gauge.GetValue() + } + if sums["read"] != 1 || sums["write"] != 1 { + t.Fatalf("not one-hot: %v", sums) + } + } + if business.Snapshot(MetadataRead, now) != before { + t.Fatal("scrape changed evidence") + } +} diff --git a/pkg/dependency/metrics.go b/pkg/dependency/metrics.go new file mode 100644 index 0000000..d5f3389 --- /dev/null +++ b/pkg/dependency/metrics.go @@ -0,0 +1,43 @@ +package dependency + +import ( + "github.com/prometheus/client_golang/prometheus" + "time" +) + +type collector struct { + monitor *Monitor + state, unresolved, last *prometheus.Desc +} + +func newCollector(m *Monitor) *collector { + labels := []string{"dependency", "impact"} + return &collector{m, + prometheus.NewDesc("dingospeed_dependency_state", "Observed capability state: 0 unknown, 1 observed healthy, 2 suspect, 3 degraded, 4 recovering. Not a node readiness decision.", labels, nil), + prometheus.NewDesc("dingospeed_dependency_unresolved", "Observed failure without sufficient positive recovery evidence.", labels, nil), + prometheus.NewDesc("dingospeed_dependency_last_observation_seconds", "Unix timestamp of last completed instrumented operation; zero means unobserved.", labels, nil), + } +} +func init() { prometheus.MustRegister(newCollector(Default)) } +func (c *collector) Describe(ch chan<- *prometheus.Desc) { + ch <- c.state + ch <- c.unresolved + ch <- c.last +} +func (c *collector) Collect(ch chan<- prometheus.Metric) { + now := time.Now() + for id, labels := range [][2]string{{"metadata_storage_read", "local_metadata_lookup"}, {"metadata_storage_write", "metadata_cache_persistence"}, {"scheduler_rpc", "cluster_coordination"}} { + s := c.monitor.Snapshot(ID(id), now) + unresolved := 0. + timestamp := 0. + if s.Unresolved { + unresolved = 1 + } + if !s.LastObservation.IsZero() { + timestamp = float64(s.LastObservation.Unix()) + } + ch <- prometheus.MustNewConstMetric(c.state, prometheus.GaugeValue, float64(s.State), labels[0], labels[1]) + ch <- prometheus.MustNewConstMetric(c.unresolved, prometheus.GaugeValue, unresolved, labels[0], labels[1]) + ch <- prometheus.MustNewConstMetric(c.last, prometheus.GaugeValue, timestamp, labels[0], labels[1]) + } +} diff --git a/pkg/dependency/promtool_test.go b/pkg/dependency/promtool_test.go new file mode 100644 index 0000000..c7ac54b --- /dev/null +++ b/pkg/dependency/promtool_test.go @@ -0,0 +1,108 @@ +package dependency + +import ( + "fmt" + "os" + "os/exec" + "path/filepath" + "strings" + "testing" + "time" + + "gopkg.in/yaml.v3" +) + +// PROMTOOL points to a separately installed, trusted promtool executable. +// Samples come from Monitor itself, so this tests the observation-to-alert chain. +func TestPrometheusAlertTimelines(t *testing.T) { + tool := os.Getenv("PROMTOOL") + if tool == "" { + t.Skip("set PROMTOOL to run the actual Prometheus rule engine") + } + rules, err := filepath.Abs("../../config/prometheus/dependency_alerts.yml") + if err != nil { + t.Fatal(err) + } + if out, err := exec.Command(tool, "check", "rules", rules).CombinedOutput(); err != nil { + t.Fatalf("check rules: %v\n%s", err, out) + } + type checkpoint struct { + sec int + degraded, unresolved bool + } + cases := []struct { + name string + id ID + events map[int]bool + checks []checkpoint + }{ + {"brief_failure_verified_recovery", MetadataRead, map[int]bool{0: false, 5: true, 10: true, 15: true}, []checkpoint{{10, false, false}, {15, false, false}, {180, false, false}}}, + {"failure_then_no_traffic", MetadataRead, map[int]bool{0: false, 5: false, 10: false}, []checkpoint{{35, false, false}, {40, true, false}, {115, true, false}, {120, true, true}, {135, false, true}, {180, false, true}}}, + {"write_failure_read_and_scheduler_healthy", MetadataWrite, map[int]bool{0: false, 5: false, 10: false, 140: true, 145: true, 150: true}, []checkpoint{{40, true, false}, {120, true, true}, {140, false, true}, {145, false, true}, {150, false, false}}}, + {"scheduler_single_failure_no_traffic", Scheduler, map[int]bool{0: false}, []checkpoint{{115, false, false}, {120, false, true}, {125, false, true}, {180, false, true}}}, + {"interrupted_recovery", Scheduler, map[int]bool{0: false, 5: false, 10: false, 45: true, 50: true, 55: false, 60: false, 65: false}, []checkpoint{{40, true, false}, {50, false, false}, {90, false, false}, {95, true, false}, {120, true, true}}}, + } + deps := []string{"metadata_storage_read", "metadata_storage_write", "scheduler_rpc"} + impacts := []string{"local_metadata_lookup", "metadata_cache_persistence", "cluster_coordination"} + var groups []any + for _, tc := range cases { + m := &Monitor{} + var states, unresolved [dependencyCount][]string + base := time.Unix(1000, 0) + for sec := 0; sec <= 180; sec += 5 { + now := base.Add(time.Duration(sec) * time.Second) + for id := ID(0); id < dependencyCount; id++ { + if id != tc.id { + m.Observe(id, true, now) + } else if success, ok := tc.events[sec]; ok { + m.Observe(id, success, now) + } + s := m.Snapshot(id, now) + states[id] = append(states[id], fmt.Sprint(int(s.State))) + v := "0" + if s.Unresolved { + v = "1" + } + unresolved[id] = append(unresolved[id], v) + } + } + var series, checks []any + for id := ID(0); id < dependencyCount; id++ { + labels := fmt.Sprintf(`{job="dingospeed",instance="node-a",dependency="%s",impact="%s"}`, deps[id], impacts[id]) + series = append(series, map[string]any{"series": "dingospeed_dependency_state" + labels, "values": strings.Join(states[id], " ")}, map[string]any{"series": "dingospeed_dependency_unresolved" + labels, "values": strings.Join(unresolved[id], " ")}) + other := strings.Replace(labels, "node-a", "node-b", 1) + series = append(series, map[string]any{"series": "dingospeed_dependency_state" + other, "values": "1x36"}, map[string]any{"series": "dingospeed_dependency_unresolved" + other, "values": "0x36"}) + } + for _, check := range tc.checks { + for i, firing := range []bool{check.degraded, check.unresolved} { + name := "DingoSpeedDependencyDegraded" + summary := "Dependency capability affected: " + deps[tc.id] + description := fmt.Sprintf("Instance node-a has sustained failed operations affecting %s. This does not establish whole-node or mount failure; do not automatically restart.", impacts[tc.id]) + if i == 1 { + name = "DingoSpeedDependencyUnresolved" + summary = "Dependency failure lacks recovery evidence: " + deps[tc.id] + description = fmt.Sprintf("Instance node-a has an unresolved observation affecting %s. Repeated failures or absence of further traffic may explain this; inspect evidence before changing request policy.", impacts[tc.id]) + } + expected := []any{} + if firing { + expected = append(expected, map[string]any{"exp_labels": map[string]string{"job": "dingospeed", "instance": "node-a", "dependency": deps[tc.id], "impact": impacts[tc.id], "severity": "warning"}, "exp_annotations": map[string]string{"summary": summary, "description": description}}) + } + checks = append(checks, map[string]any{"eval_time": fmt.Sprintf("%ds", check.sec), "alertname": name, "exp_alerts": expected}) + } + } + groups = append(groups, map[string]any{"name": tc.name, "interval": "5s", "input_series": series, "alert_rule_test": checks}) + } + data, err := yaml.Marshal(map[string]any{"rule_files": []string{rules}, "evaluation_interval": "5s", "tests": groups}) + if err != nil { + t.Fatal(err) + } + fixture := filepath.Join(t.TempDir(), "dependency_timelines.yml") + if err := os.WriteFile(fixture, data, 0600); err != nil { + t.Fatal(err) + } + if out, err := exec.Command(tool, "test", "rules", fixture).CombinedOutput(); err != nil { + t.Fatalf("timeline test: %v\n%s", err, out) + } else { + t.Log(string(out)) + } +} diff --git a/pkg/dependency/state.go b/pkg/dependency/state.go new file mode 100644 index 0000000..f24046d --- /dev/null +++ b/pkg/dependency/state.go @@ -0,0 +1,131 @@ +// Package dependency observes existing operations without controlling requests. +package dependency + +import ( + "sync" + "time" +) + +type ID int + +const ( + MetadataRead ID = iota + MetadataWrite + Scheduler + dependencyCount +) + +type State int + +const ( + Unknown State = iota + ObservedHealthy + Suspect + Degraded + Recovering +) + +// These conservative defaults require evidence spread over time, not a burst +// of failed suboperations from a single request. No traffic means unknown. +const ( + EvidenceWindow = 30 * time.Second + MinEvidenceSpan = 10 * time.Second + Freshness = 2 * time.Minute +) + +type evidence struct { + // At most 31 distinct seconds intersect an inclusive 30-second window. + times [31]time.Time + samples int +} + +func (e *evidence) add(now time.Time) { + start := 0 + for start < e.samples && now.Sub(e.times[start]) > EvidenceWindow { + start++ + } + if start > 0 { + e.samples = copy(e.times[:], e.times[start:e.samples]) + } + if e.samples > 0 && now.Unix() == e.times[e.samples-1].Unix() { + return + } + e.times[e.samples] = now + e.samples++ +} +func (e evidence) sufficient() bool { + return e.samples >= 3 && e.times[e.samples-1].Sub(e.times[0]) >= MinEvidenceSpan +} + +type Snapshot struct { + State State + LastObservation time.Time + LastFailure time.Time + Unresolved bool +} +type entry struct { + Snapshot + failures, successes evidence +} +type Monitor struct { + mu sync.Mutex + entries [dependencyCount]entry +} + +var Default = &Monitor{} + +// Observe records a completed operation. Missing files and invalid content must +// be filtered by the caller. Success is positive evidence for this capability +// only; it does not prove every path, mount or other capability is healthy. +func (m *Monitor) Observe(id ID, success bool, now time.Time) { + if id < 0 || id >= dependencyCount { + return + } + m.mu.Lock() + defer m.mu.Unlock() + e := &m.entries[id] + if !e.LastObservation.IsZero() && now.Before(e.LastObservation) { + return + } + e.LastObservation = now + if !success { + e.LastFailure = now + e.Unresolved = true + e.successes = evidence{} + e.failures.add(now) + if e.State != Degraded { + e.State = Suspect + } + if e.failures.sufficient() { + e.State = Degraded + } + return + } + if !e.Unresolved { + e.State = ObservedHealthy + return + } + e.successes.add(now) + e.State = Recovering + if e.successes.sufficient() { + e.State = ObservedHealthy + e.Unresolved = false + e.failures = evidence{} + e.successes = evidence{} + } +} + +// Snapshot does not probe dependencies or modify state. Stale evidence never +// resolves an incident: Unresolved stays true until positive recovery evidence. +func (m *Monitor) Snapshot(id ID, now time.Time) Snapshot { + if id < 0 || id >= dependencyCount { + return Snapshot{State: Unknown} + } + m.mu.Lock() + defer m.mu.Unlock() + s := m.entries[id].Snapshot + if s.LastObservation.IsZero() || now.Sub(s.LastObservation) > Freshness { + s.State = Unknown + } + return s +} diff --git a/pkg/dependency/state_test.go b/pkg/dependency/state_test.go new file mode 100644 index 0000000..ca117f2 --- /dev/null +++ b/pkg/dependency/state_test.go @@ -0,0 +1,132 @@ +package dependency + +import ( + "github.com/prometheus/client_golang/prometheus" + "sync" + "testing" + "time" +) + +func TestScopeAndRecovery(t *testing.T) { + m := &Monitor{} + now := time.Unix(1000, 0) + if m.Snapshot(MetadataRead, now).State != Unknown { + t.Fatal("unobserved is not healthy") + } + m.Observe(MetadataRead, false, now) + for i := 0; i < 100; i++ { + m.Observe(MetadataRead, false, now) + } + if m.Snapshot(MetadataRead, now).State != Suspect { + t.Fatal("burst promoted to dependency failure") + } + m.Observe(MetadataRead, false, now.Add(5*time.Second)) + m.Observe(MetadataRead, false, now.Add(10*time.Second)) + if m.Snapshot(MetadataRead, now.Add(10*time.Second)).State != Degraded { + t.Fatal("sustained failure not identified") + } + if m.Snapshot(MetadataWrite, now).State != Unknown || m.Snapshot(Scheduler, now).State != Unknown { + t.Fatal("failure scope spread") + } + stale := m.Snapshot(MetadataRead, now.Add(3*time.Minute)) + if stale.State != Unknown || !stale.Unresolved { + t.Fatal("silence treated as recovery") + } + for _, sec := range []int{180, 185} { + m.Observe(MetadataRead, true, now.Add(time.Duration(sec)*time.Second)) + } + if s := m.Snapshot(MetadataRead, now.Add(185*time.Second)); s.State != Recovering || !s.Unresolved { + t.Fatal("premature recovery") + } + m.Observe(MetadataRead, true, now.Add(190*time.Second)) + if s := m.Snapshot(MetadataRead, now.Add(190*time.Second)); s.State != ObservedHealthy || s.Unresolved { + t.Fatal("positive recovery not accepted") + } +} + +func TestEvidenceSlidingWindow(t *testing.T) { + for _, success := range []bool{false, true} { + m := &Monitor{} + base := time.Unix(1000, 0) + if success { + m.Observe(MetadataRead, false, base.Add(-time.Second)) + } + for _, sec := range []int{0, 25, 31, 35} { + m.Observe(MetadataRead, success, base.Add(time.Duration(sec)*time.Second)) + } + want := Degraded + if success { + want = ObservedHealthy + } + if got := m.Snapshot(MetadataRead, base.Add(35*time.Second)); got.State != want { + t.Fatalf("success=%v: recent samples at 25,31,35 were lost: %+v", success, got) + } + } + var e evidence + base := time.Unix(1000, 0) + for i := 0; i < 10000; i++ { + e.add(base.Add(time.Duration(i) * 100 * time.Millisecond)) + if e.samples > 31 || !e.times[e.samples-1].After(base.Add(-time.Second)) { + t.Fatal("window is not bounded") + } + } +} + +func TestWindowsAndInterruptedRecovery(t *testing.T) { + m := &Monitor{} + now := time.Unix(1000, 0) + for _, sec := range []int{0, 40, 80} { + m.Observe(Scheduler, false, now.Add(time.Duration(sec)*time.Second)) + } + if m.Snapshot(Scheduler, now.Add(80*time.Second)).State != Suspect { + t.Fatal("sparse errors aggregated outside window") + } + m.Observe(Scheduler, true, now.Add(81*time.Second)) + m.Observe(Scheduler, true, now.Add(86*time.Second)) + m.Observe(Scheduler, false, now.Add(90*time.Second)) + m.Observe(Scheduler, true, now.Add(91*time.Second)) + if s := m.Snapshot(Scheduler, now.Add(91*time.Second)); !s.Unresolved || s.State != Recovering { + t.Fatal("interrupted recovery accepted") + } + m.Observe(Scheduler, true, now) // Out of order evidence cannot resolve incident. + if !m.Snapshot(Scheduler, now.Add(91*time.Second)).Unresolved { + t.Fatal("out of order observation changed state") + } +} + +func TestCollectorConcurrentAndReadOnly(t *testing.T) { + m := &Monitor{} + registry := prometheus.NewPedanticRegistry() + registry.MustRegister(newCollector(m)) + var wg sync.WaitGroup + for i := 0; i < 4; i++ { + wg.Add(1) + go func() { + defer wg.Done() + for j := 0; j < 100; j++ { + m.Observe(MetadataWrite, j%2 == 0, time.Now()) + if _, err := registry.Gather(); err != nil { + t.Error(err) + } + } + }() + } + wg.Wait() + before := m.Snapshot(MetadataWrite, time.Now()) + families, err := registry.Gather() + if err != nil { + t.Fatal(err) + } + if len(families) != 3 { + t.Fatalf("unexpected metrics count %d", len(families)) + } + for _, f := range families { + if len(f.Metric) != 3 { + t.Fatal("unbounded labels") + } + } + after := m.Snapshot(MetadataWrite, time.Now()) + if before != after { + t.Fatal("scrape mutated state") + } +} diff --git a/pkg/error/error.go b/pkg/error/error.go index 3dd6d09..388c9e4 100644 --- a/pkg/error/error.go +++ b/pkg/error/error.go @@ -46,3 +46,7 @@ func NewAppendCode(code int, msg string) Error { func Wrap(msg string, err error) Error { return Error{msg: msg, err: err} } + +func WrapCode(code int, msg string, err error) Error { + return Error{statusCode: code, msg: msg, err: err} +} diff --git a/pkg/hfprojection/index.go b/pkg/hfprojection/index.go new file mode 100644 index 0000000..e559270 --- /dev/null +++ b/pkg/hfprojection/index.go @@ -0,0 +1,124 @@ +package hfprojection + +import ( + "os" + "sync" +) + +// Index is the read-only projection used by repository catalogue queries. A +// refresh scans the existing HF cache once and atomically replaces one repo +// type; manifest reads then reuse that snapshot without walking the cache. +type Index struct { + mu sync.RWMutex + refreshMu sync.Mutex + scopes map[string]indexedScope +} + +type indexedScope struct { + catalog Catalog + manifests map[string]Manifest + errors map[string]error +} + +func NewIndex() *Index { return &Index{scopes: make(map[string]indexedScope)} } + +var DefaultIndex = NewIndex() + +func scopeKey(root, repoType string) string { return root + "\x00" + repoType } +func manifestKey(repo, revision string) string { return repo + "\x00" + revision } + +// Refresh rebuilds only from local cache files. It never calls HF and never +// writes the provider cache or its database. +func (i *Index) Refresh(root, repoType string) (Catalog, error) { + i.refreshMu.Lock() + defer i.refreshMu.Unlock() + return i.refresh(root, repoType) +} + +func (i *Index) refresh(root, repoType string) (Catalog, error) { + reader := Reader{Root: root} + catalog, err := reader.Catalog(repoType) + if err != nil { + return Catalog{}, err + } + scope := indexedScope{catalog: catalog, manifests: make(map[string]Manifest), errors: make(map[string]error)} + for _, repo := range catalog.Repos { + if repo.Error != "" { + continue + } + for _, revision := range repo.Revisions { + manifest, readErr := reader.Manifest(repoType, repo.Repo, revision.Name) + key := manifestKey(repo.Repo, revision.Name) + if readErr != nil { + scope.errors[key] = readErr + continue + } + scope.manifests[key] = manifest + commitKey := manifestKey(repo.Repo, revision.Commit) + if _, seen := scope.manifests[commitKey]; seen { + continue + } + commitManifest := manifest + commitManifest.Revision = revision.Commit + scope.manifests[commitKey] = commitManifest + } + } + i.mu.Lock() + i.scopes[scopeKey(root, repoType)] = scope + i.mu.Unlock() + return catalog, nil +} + +func (i *Index) Manifest(root, repoType, repo, revision string) (Manifest, error) { + key := scopeKey(root, repoType) + i.mu.RLock() + _, exists := i.scopes[key] + i.mu.RUnlock() + if !exists { + i.refreshMu.Lock() + i.mu.RLock() + _, exists = i.scopes[key] + i.mu.RUnlock() + if !exists { + _, err := i.refresh(root, repoType) + if err != nil { + i.refreshMu.Unlock() + return Manifest{}, err + } + } + i.refreshMu.Unlock() + } + i.mu.RLock() + scope := i.scopes[key] + manifest, ok := scope.manifests[manifestKey(repo, revision)] + readErr := scope.errors[manifestKey(repo, revision)] + i.mu.RUnlock() + if readErr != nil { + return Manifest{}, readErr + } + if !ok { + return Manifest{}, os.ErrNotExist + } + return manifest, nil +} + +// Invalidate makes the next catalogue or manifest read rebuild local facts. +// It is called after a cache task changes HF bytes. +func (i *Index) Invalidate(root, repoType string) { + if i == nil { + return + } + i.refreshMu.Lock() + defer i.refreshMu.Unlock() + i.mu.Lock() + defer i.mu.Unlock() + if repoType == "" { + for key := range i.scopes { + if len(key) > len(root) && key[:len(root)+1] == root+"\x00" { + delete(i.scopes, key) + } + } + return + } + delete(i.scopes, scopeKey(root, repoType)) +} diff --git a/pkg/hfprojection/mount_sources.go b/pkg/hfprojection/mount_sources.go new file mode 100644 index 0000000..472a0bb --- /dev/null +++ b/pkg/hfprojection/mount_sources.go @@ -0,0 +1,28 @@ +package hfprojection + +import "fmt" + +// MountSources resolves one observed manifest, without re-reading a mutable +// branch for each file. The caller must validate and pin every payload before +// publishing the mount. Paths are local administrative data, not an HTTP API. +func (r Reader) MountSources(typ, repo, revision string) (Manifest, map[string]string, error) { + m, err := r.Manifest(typ, repo, revision) + if err != nil { + return m, nil, err + } + if !m.ManifestComplete { + return m, nil, fmt.Errorf("complete cached repository manifest required") + } + paths := make(map[string]string, len(m.Files)) + for _, f := range m.Files { + if f.localPath == "" { + return m, nil, fmt.Errorf("missing content: %s", f.Path) + } + p, err := r.contentPath(m.RepoKey, f.localPath) + if err != nil { + return m, nil, fmt.Errorf("%s: %w", f.Path, err) + } + paths[f.Path] = p + } + return m, paths, nil +} diff --git a/pkg/hfprojection/mount_sources_test.go b/pkg/hfprojection/mount_sources_test.go new file mode 100644 index 0000000..3680c58 --- /dev/null +++ b/pkg/hfprojection/mount_sources_test.go @@ -0,0 +1,50 @@ +package hfprojection + +import ( + "os" + "path/filepath" + "testing" +) + +func TestMountSourcesPinsOneManifest(t *testing.T) { + r, _ := fixture(t) + blob(t, filepath.Join(r.Root, "files", "models", "Qwen", "demo", "blobs", "missing"), 9, 7) + m, paths, err := r.MountSources("models", "Qwen/demo", "main") + if err != nil { + t.Fatal(err) + } + if m.Commit != "commit1" || len(paths) != 3 { + t.Fatalf("unexpected sources: %+v %v", m, paths) + } + if paths["weights.bin"] != filepath.Join(r.Root, "files", "models", "Qwen", "demo", "blobs", "weight") { + t.Fatal(paths) + } +} + +func TestMountSourcesRequiresCompleteMetadata(t *testing.T) { + r, api := fixture(t) + if err := os.Chmod(filepath.Join(api, "revision", "main", "meta_get.json"), 0600); err != nil { + t.Fatal(err) + } + response(t, filepath.Join(api, "revision", "main", "meta_get.json"), map[string]any{"sha": "commit1"}) + if _, _, err := r.MountSources("models", "Qwen/demo", "main"); err == nil { + t.Fatal("accepted incomplete manifest") + } +} + +func TestMountSourcesRejectsEscapingLink(t *testing.T) { + r, _ := fixture(t) + blob(t, filepath.Join(r.Root, "files", "models", "Qwen", "demo", "blobs", "missing"), 9, 7) + p := filepath.Join(r.Root, "files", "models", "Qwen", "demo", "blobs", "weight") + if err := os.Remove(p); err != nil { + t.Fatal(err) + } + outside := filepath.Join(t.TempDir(), "other") + blob(t, outside, 8, 3) + if err := os.Symlink(outside, p); err != nil { + t.Skip(err) + } + if _, _, err := r.MountSources("models", "Qwen/demo", "main"); err == nil { + t.Fatal("accepted escaping blob") + } +} diff --git a/pkg/hfprojection/projection.go b/pkg/hfprojection/projection.go new file mode 100644 index 0000000..eb373d5 --- /dev/null +++ b/pkg/hfprojection/projection.go @@ -0,0 +1,633 @@ +// Package hfprojection reads existing HF cache facts without registration, +// network access, database access, or opening any file for writing. +package hfprojection + +import ( + "encoding/binary" + "encoding/hex" + "encoding/json" + "errors" + "fmt" + "io" + "io/fs" + "math" + "os" + "path/filepath" + "sort" + "strings" + + "dingospeed/pkg/repository" +) + +type Revision struct { + Name string `json:"name"` + Commit string `json:"commit"` +} +type Repo struct { + repository.RepoKey + Revisions []Revision `json:"revisions"` + Access string `json:"access"` + Error string `json:"error,omitempty"` +} +type Catalog struct { + Namespace string `json:"namespace"` + RepoType string `json:"repoType"` + Repos []Repo `json:"repos"` +} +type File struct { + Path string `json:"path"` + OID string `json:"oid"` + Size int64 `json:"size"` + CachedBytes int64 `json:"cachedBytes"` + CacheStatus string `json:"cacheStatus"` + localPath string +} +type Manifest struct { + repository.RepoKey + Revision string `json:"revision"` + Commit string `json:"commit"` + MetadataAvailable bool `json:"metadataAvailable"` + ManifestComplete bool `json:"manifestComplete"` + Access string `json:"access"` + Files []File `json:"files"` +} + +// Reader has only a filesystem root. No downloader or database can be injected. +type Reader struct { + Root string + Namespace string +} + +func (r Reader) namespace() string { + if r.Namespace != "" { + return r.Namespace + } + return repository.HuggingFace +} + +func (r Reader) key(typ, repo string) (repository.RepoKey, error) { + k := repository.RepoKey{Namespace: r.namespace(), RepoType: typ, Repo: repo} + if err := k.Validate(); err != nil { + return k, err + } + p := strings.Split(repo, "/") + if len(p) > 2 || (r.namespace() == repository.HuggingFace && (strings.EqualFold(p[0], repository.Local) || strings.EqualFold(p[0], repository.ModelScope))) { + return k, fmt.Errorf("not an upstream HF repository") + } + return k, nil +} +func (r Reader) safe(p string) error { return repository.SafePath(r.Root, p) } +func (r Reader) dirs(p string) ([]os.DirEntry, error) { + if err := r.safe(p); err != nil { + return nil, err + } + v, err := os.ReadDir(p) + if errors.Is(err, os.ErrNotExist) { + return nil, nil + } + return v, err +} + +// Catalog discovers legacy layouts, including repositories without a marker. +// It never traverses hosted/modelscope roots or guesses identity from blobs. +func (r Reader) Catalog(typ string) (Catalog, error) { + out := Catalog{Namespace: r.namespace(), RepoType: typ, Repos: []Repo{}} + if _, err := r.key(typ, "placeholder"); err != nil { + return out, err + } + if err := r.safe(r.Root); err != nil { + return out, err + } + if _, err := os.Stat(r.Root); err != nil { + return out, err + } + names := map[string]bool{} + if r.namespace() == repository.ModelScope { + descriptors, err := repository.List(r.Root) + if err != nil { + return out, err + } + for _, d := range descriptors { + if d.Namespace == r.namespace() && d.RepoType == typ { + names[d.Repo] = true + } + } + } else { + for _, area := range []string{"api", "files"} { + base := filepath.Join(r.Root, area, typ) + owners, err := r.dirs(base) + if err != nil { + return out, err + } + for _, owner := range owners { + if !owner.IsDir() || strings.EqualFold(owner.Name(), repository.Local) || strings.EqualFold(owner.Name(), repository.ModelScope) { + continue + } + children, err := r.dirs(filepath.Join(base, owner.Name())) + if err != nil { + return out, err + } + nested := map[string]bool{} + for _, child := range children { + if !child.IsDir() { + continue + } + entries, err := r.dirs(filepath.Join(base, owner.Name(), child.Name())) + if err != nil { + return out, err + } + for _, e := range entries { + if e.IsDir() && cacheArea(e.Name(), area) { + names[owner.Name()+"/"+child.Name()] = true + nested[child.Name()] = true + } + } + } + // A real two-segment repository can itself be named revision, + // paths-info or resolve. Prefer its own layout over treating its + // name as the enclosing single-segment repository's cache area. + for _, child := range children { + if child.IsDir() && cacheArea(child.Name(), area) && !nested[child.Name()] { + names[owner.Name()] = true + } + } + } + } + } + for name := range names { + k, err := r.key(typ, name) + if err != nil { + return out, err + } + repo := Repo{RepoKey: k, Revisions: []Revision{}, Access: "unknown"} + repo.Revisions, err = r.revisions(k) + if err == nil { + repo.Access, err = r.access(k, repo.Revisions) + } + if err != nil { + repo.Error = "cached repository metadata is unreadable or invalid" + repo.Revisions = []Revision{} + } + out.Repos = append(out.Repos, repo) + } + sort.Slice(out.Repos, func(i, j int) bool { return out.Repos[i].Repo < out.Repos[j].Repo }) + return out, nil +} +func cacheArea(name, area string) bool { + return area == "api" && (name == "revision" || name == "paths-info") || area == "files" && name == "resolve" +} + +type metadata struct { + SHA string `json:"sha"` + Private *bool `json:"private"` + Gated json.RawMessage `json:"gated"` + Siblings []struct { + Path string `json:"rfilename"` + Size *int64 `json:"size"` + OID string `json:"blobId"` + LFS *struct { + OID string `json:"sha256"` + Size int64 `json:"size"` + } `json:"lfs"` + } `json:"siblings"` +} + +func metadataAccess(m metadata) string { + if m.Private != nil && *m.Private { + return "restricted" + } + gated := strings.TrimSpace(string(m.Gated)) + if gated != "" && gated != "false" && gated != "null" { + return "restricted" + } + if m.Private != nil && !*m.Private && gated == "false" { + return "public" + } + return "unknown" +} + +// No cached permission fact is inferred from successful download or a repo +// name. A missing permission field is unknown, not permission to publish it. +func (r Reader) access(k repository.RepoKey, revisions []Revision) (string, error) { + access := "public" + // Hidden fixed-commit metadata can still contain a restrictive permission + // fact. Deduplicating the displayed branch list must not drop that fact. + seen := map[string]bool{} + all := append([]Revision(nil), revisions...) + for _, rev := range all { + seen[rev.Name] = true + } + entries, err := r.dirs(filepath.Join(k.APIRoot(r.Root), "revision")) + if err != nil { + return "unknown", err + } + for _, entry := range entries { + if entry.IsDir() && !seen[entry.Name()] { + all = append(all, Revision{Name: entry.Name()}) + } + } + if len(all) == 0 { + return "unknown", nil + } + for _, rev := range all { + var m metadata + err := r.decode(filepath.Join(k.Revision(r.Root, rev.Name), "meta_get.json"), &m) + if errors.Is(err, os.ErrNotExist) { + access = "unknown" + continue + } + if err != nil { + return "unknown", err + } + switch metadataAccess(m) { + case "restricted": + return "restricted", nil + case "unknown": + access = "unknown" + } + } + return access, nil +} + +func (r Reader) decode(p string, out any) error { + if err := r.safe(p); err != nil { + return err + } + f, err := os.Open(p) + if err != nil { + return err + } + defer f.Close() + // Cache responses are bounded; malformed/truncated data must fail closed. + b, err := io.ReadAll(io.LimitReader(f, 32*1024*1024+1)) + if err != nil { + return err + } + if len(b) > 32*1024*1024 { + return fmt.Errorf("cached metadata exceeds projection limit") + } + var wrapper struct { + Status int `json:"status_code"` + Content string `json:"content"` + } + if err = json.Unmarshal(b, &wrapper); err != nil { + return fmt.Errorf("invalid cached metadata: %w", err) + } + if wrapper.Status != 200 { + return fmt.Errorf("cached metadata is not successful (HTTP %d)", wrapper.Status) + } + b, err = hex.DecodeString(wrapper.Content) + if err != nil { + return fmt.Errorf("invalid cached metadata encoding: %w", err) + } + if err = json.Unmarshal(b, out); err != nil { + return fmt.Errorf("invalid cached response body: %w", err) + } + return nil +} +func (r Reader) revisions(k repository.RepoKey) ([]Revision, error) { + byName := map[string]string{} + entries, err := r.dirs(filepath.Join(k.APIRoot(r.Root), "revision")) + if err != nil { + return nil, err + } + for _, e := range entries { + if !e.IsDir() { + continue + } + if err := repository.Segment(e.Name()); err != nil { + return nil, err + } + var meta metadata + err := r.decode(filepath.Join(k.Revision(r.Root, e.Name()), "meta_get.json"), &meta) + if errors.Is(err, os.ErrNotExist) { + continue + } + if err != nil { + return nil, err + } + if err := repository.Segment(meta.SHA); err != nil { + return nil, fmt.Errorf("invalid cached commit") + } + byName[e.Name()] = meta.SHA + } + if k.Namespace != repository.ModelScope { + for _, base := range []string{filepath.Join(k.APIRoot(r.Root), "paths-info"), filepath.Join(k.FilesRoot(r.Root), "resolve")} { + entries, err := r.dirs(base) + if err != nil { + return nil, err + } + for _, e := range entries { + if e.IsDir() { + if err := repository.Segment(e.Name()); err != nil { + return nil, err + } + if _, ok := byName[e.Name()]; !ok { + byName[e.Name()] = e.Name() + } + } + } + } + } + out := []Revision{} + pointed := map[string]bool{} + for name, commit := range byName { + if name != commit { + pointed[commit] = true + } + } + for name, commit := range byName { + if name == commit && pointed[commit] { + continue + } + out = append(out, Revision{name, commit}) + } + sort.Slice(out, func(i, j int) bool { return out[i].Name < out[j].Name }) + return out, nil +} + +func (r Reader) Manifest(typ, repo, revision string) (Manifest, error) { + k, err := r.key(typ, repo) + out := Manifest{RepoKey: k, Revision: revision, Files: []File{}, Access: "unknown"} + if err != nil { + return out, err + } + if err = repository.Segment(revision); err != nil { + return out, err + } + revisions, err := r.revisions(k) + if err != nil { + return out, err + } + for _, rev := range revisions { + if rev.Name == revision { + out.Commit = rev.Commit + break + } + } + // Fixed commits remain directly readable even when hidden from the branch + // list to avoid counting the same cached revision twice. + if out.Commit == "" { + for _, rev := range revisions { + if rev.Commit == revision { + out.Commit = revision + break + } + } + } + if out.Commit == "" { + return out, os.ErrNotExist + } + out.Access, err = r.access(k, revisions) + if err != nil { + return out, err + } + files := map[string]File{} + var meta metadata + err = r.decode(filepath.Join(k.Revision(r.Root, revision), "meta_get.json"), &meta) + if errors.Is(err, os.ErrNotExist) && revision != out.Commit { + err = r.decode(filepath.Join(k.Revision(r.Root, out.Commit), "meta_get.json"), &meta) + } + if err != nil && !errors.Is(err, os.ErrNotExist) { + return out, err + } + if err == nil { + if meta.SHA != out.Commit { + return out, fmt.Errorf("cached revision changed; refresh projection") + } + out.MetadataAvailable = true + out.ManifestComplete = meta.Siblings != nil + for _, s := range meta.Siblings { + if err := repository.Relative(s.Path); err != nil { + return out, err + } + v := File{Path: s.Path, OID: s.OID, Size: -1, CacheStatus: "unknown"} + if s.Size != nil { + v.Size = *s.Size + } + if s.LFS != nil { + v.OID = s.LFS.OID + v.Size = s.LFS.Size + } + files[s.Path] = v + } + } + infoRoot := filepath.Join(k.APIRoot(r.Root), "paths-info", out.Commit) + err = r.walk(infoRoot, func(p string, e fs.DirEntry) error { + if e.Name() != "paths-info_post.json" { + return nil + } + rel, _ := filepath.Rel(infoRoot, filepath.Dir(p)) + path := filepath.ToSlash(rel) + if err := repository.Relative(path); err != nil { + return err + } + var items []struct { + OID string `json:"oid"` + Size int64 `json:"size"` + LFS struct { + OID string `json:"oid"` + } `json:"lfs"` + } + if err := r.decode(p, &items); err != nil { + return err + } + if len(items) != 1 { + return fmt.Errorf("invalid cached file metadata") + } + v := files[path] + v.Path = path + v.Size = items[0].Size + v.OID = items[0].OID + if items[0].LFS.OID != "" { + v.OID = items[0].LFS.OID + } + files[path] = v + return nil + }) + if err != nil { + return out, err + } + resolveRoot := filepath.Join(k.FilesRoot(r.Root), "resolve", out.Commit) + err = r.walk(resolveRoot, func(p string, e fs.DirEntry) error { + rel, _ := filepath.Rel(resolveRoot, p) + path := filepath.ToSlash(rel) + if err := repository.Relative(path); err != nil { + return err + } + v, ok := files[path] + if !ok { + v = File{Path: path, Size: -1} + } + v.localPath = p + files[path] = v + return nil + }) + if err != nil { + return out, err + } + for _, v := range files { + if v.OID != "" { + if err := repository.Segment(v.OID); err != nil { + return out, fmt.Errorf("invalid cached object identity") + } + v.localPath = k.Blob(r.Root, v.OID) + } + // Catalogue manifests describe the cached upstream snapshot. Payload + // completeness belongs to transfer/file-serving paths and is deliberately + // not inspected while listing repositories or files. + v.CacheStatus = "unknown" + out.Files = append(out.Files, v) + } + sort.Slice(out.Files, func(i, j int) bool { return out.Files[i].Path < out.Files[j].Path }) + return out, nil +} + +func (r Reader) walk(root string, visit func(string, fs.DirEntry) error) error { + if err := r.safe(root); err != nil { + return err + } + err := filepath.WalkDir(root, func(p string, e fs.DirEntry, err error) error { + if err != nil { + return err + } + if e.IsDir() { + return nil + } + return visit(p, e) + }) + if errors.Is(err, os.ErrNotExist) { + if _, rootErr := os.Stat(root); errors.Is(rootErr, os.ErrNotExist) { + return nil + } + } + return err +} + +// Only resolve symlinks into this repository's blob directory. Never follow +// arbitrary links out of the cache or into a different repository. +func (r Reader) contentPath(k repository.RepoKey, p string) (string, error) { + if err := r.safe(filepath.Dir(p)); err != nil { + return "", err + } + info, err := os.Lstat(p) + if err != nil { + return "", err + } + if info.Mode()&os.ModeSymlink != 0 { + target, err := filepath.EvalSymlinks(p) + if err != nil { + return "", err + } + abs, _ := filepath.Abs(filepath.Join(k.FilesRoot(r.Root), "blobs")) + target, _ = filepath.Abs(target) + if filepath.Dir(target) != abs { + return "", fmt.Errorf("cached resolve link leaves repository blobs") + } + p = target + } + if err := r.safe(p); err != nil { + return "", err + } + return p, nil +} + +// inspect validates a bounded OLAH v8 header and actual bytes for every cached +// block. Corrupt OLAH is unknown, never treated as a complete plain file. +func inspect(f *os.File, expected int64) (status string, size, cached, offset int64, err error) { + info, err := f.Stat() + if err != nil { + return "unknown", 0, 0, 0, err + } + if !info.Mode().IsRegular() { + return "unknown", 0, 0, 0, fmt.Errorf("cache content is not a regular file") + } + var raw [36]byte + n, readErr := f.ReadAt(raw[:], 0) + if n < 4 || string(raw[:4]) != "OLAH" { + if expected < 0 || expected != info.Size() { + return "unknown", info.Size(), 0, 0, nil + } + return "complete", info.Size(), info.Size(), 0, nil + } + if readErr != nil { + return "unknown", expected, 0, 0, nil + } + version := binary.LittleEndian.Uint64(raw[4:12]) + block := binary.LittleEndian.Uint64(raw[12:20]) + length := binary.LittleEndian.Uint64(raw[20:28]) + bits := binary.LittleEndian.Uint64(raw[28:36]) + if version != 8 || block == 0 || block > math.MaxInt64 || length > math.MaxInt64 || bits > 8*1024*1024 || bits == 0 { + return "unknown", expected, 0, 0, nil + } + count := length / block + if length%block != 0 { + count++ + } + if count > bits { + return "unknown", expected, 0, 0, nil + } + mask := make([]byte, (bits+7)/8) + if _, err := f.ReadAt(mask, 36); err != nil { + return "unknown", expected, 0, 0, nil + } + size = int64(length) + offset = 36 + int64(len(mask)) + if expected >= 0 && expected != size { + return "unknown", size, 0, offset, nil + } + complete := true + for i := uint64(0); i < count; i++ { + start := i * block + end := start + min(block, length-start) + if mask[i/8]&(1<<(i%8)) != 0 && info.Size()-offset >= int64(end) { + cached += int64(end - start) + } else { + complete = false + } + } + if complete { + return "complete", size, cached, offset, nil + } + return "partial", size, cached, offset, nil +} + +var ErrIncomplete = errors.New("file is not completely cached") + +// OpenFile returns only an already complete local payload, pinned to an open +// read-only descriptor. Callers must close it; no cache manager is involved. +func (r Reader) OpenFile(typ, repo, revision, path string) (*os.File, *io.SectionReader, error) { + if err := repository.Relative(path); err != nil { + return nil, nil, err + } + m, err := r.Manifest(typ, repo, revision) + if err != nil { + return nil, nil, err + } + for _, v := range m.Files { + if v.Path == path { + if v.localPath == "" { + return nil, nil, os.ErrNotExist + } + p, err := r.contentPath(m.RepoKey, v.localPath) + if err != nil { + return nil, nil, err + } + f, err := os.Open(p) + if err != nil { + return nil, nil, err + } + state, size, _, offset, err := inspect(f, v.Size) + if err != nil { + f.Close() + return nil, nil, err + } + if state != "complete" { + f.Close() + return nil, nil, ErrIncomplete + } + return f, io.NewSectionReader(f, offset, size), nil + } + } + return nil, nil, os.ErrNotExist +} diff --git a/pkg/hfprojection/projection_test.go b/pkg/hfprojection/projection_test.go new file mode 100644 index 0000000..0b0977f --- /dev/null +++ b/pkg/hfprojection/projection_test.go @@ -0,0 +1,347 @@ +package hfprojection + +import ( + "crypto/sha256" + "encoding/binary" + "encoding/hex" + "encoding/json" + "errors" + "io" + "os" + "path/filepath" + "reflect" + "testing" +) + +func write(t *testing.T, p string, b []byte) { + t.Helper() + if err := os.MkdirAll(filepath.Dir(p), 0755); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(p, b, 0444); err != nil { + t.Fatal(err) + } +} +func response(t *testing.T, p string, v any) { + t.Helper() + body, err := json.Marshal(v) + if err != nil { + t.Fatal(err) + } + b, _ := json.Marshal(map[string]any{"status_code": 200, "content": hex.EncodeToString(body)}) + write(t, p, b) +} +func blob(t *testing.T, p string, size uint64, mask byte) { + t.Helper() + b := make([]byte, 37+size) + copy(b, "OLAH") + binary.LittleEndian.PutUint64(b[4:], 8) + binary.LittleEndian.PutUint64(b[12:], 4) + binary.LittleEndian.PutUint64(b[20:], size) + binary.LittleEndian.PutUint64(b[28:], 8) + b[36] = mask + copy(b[37:], "abcdefgh") + write(t, p, b) +} +func fixture(t *testing.T) (Reader, string) { + t.Helper() + root := t.TempDir() + api := filepath.Join(root, "api", "models", "Qwen", "demo") + response(t, filepath.Join(api, "revision", "main", "meta_get.json"), map[string]any{"sha": "commit1", "siblings": []any{map[string]any{"rfilename": "README.md"}, map[string]any{"rfilename": "weights.bin"}, map[string]any{"rfilename": "missing.bin"}}}) + for _, item := range []struct { + path, oid string + size int + }{{"README.md", "readme", 8}, {"weights.bin", "weight", 8}, {"missing.bin", "missing", 9}} { + response(t, filepath.Join(api, "paths-info", "commit1", item.path, "paths-info_post.json"), []any{map[string]any{"path": item.path, "oid": item.oid, "size": item.size}}) + } + blob(t, filepath.Join(root, "files", "models", "Qwen", "demo", "blobs", "readme"), 8, 3) + blob(t, filepath.Join(root, "files", "models", "Qwen", "demo", "blobs", "weight"), 8, 1) + return Reader{Root: root}, api +} + +type stamp struct { + Sum [32]byte + Size int64 + Time int64 + Mode os.FileMode +} + +func snapshot(t *testing.T, root string) map[string]stamp { + t.Helper() + out := map[string]stamp{} + err := filepath.Walk(root, func(p string, i os.FileInfo, err error) error { + if err != nil { + return err + } + b := []byte{} + if !i.IsDir() { + b, err = os.ReadFile(p) + if err != nil { + return err + } + } + out[p] = stamp{sha256.Sum256(b), i.Size(), i.ModTime().UnixNano(), i.Mode()} + return nil + }) + if err != nil { + t.Fatal(err) + } + return out +} +func TestLegacyProjectionReadsWithoutChangingAnyCacheEntry(t *testing.T) { + r, api := fixture(t) + // A colliding company model must never be discovered as a HF owner. + response(t, filepath.Join(r.Root, "api", "models", "dingo-local", "datacanvas", "demo", "revision", "main", "meta_get.json"), map[string]any{"sha": "hosted", "siblings": []any{}}) + before := snapshot(t, r.Root) + for range 2 { + cat, err := r.Catalog("models") + if err != nil { + t.Fatal(err) + } + if len(cat.Repos) != 1 || cat.Repos[0].Repo != "Qwen/demo" || cat.Repos[0].Error != "" { + t.Fatalf("catalog=%+v", cat) + } + m, err := r.Manifest("models", "Qwen/demo", "main") + if err != nil { + t.Fatal(err) + } + if !m.ManifestComplete || !m.MetadataAvailable || m.Commit != "commit1" || len(m.Files) != 3 { + t.Fatalf("manifest=%+v", m) + } + byPath := map[string]File{} + for _, f := range m.Files { + byPath[f.Path] = f + } + if byPath["README.md"].CacheStatus != "unknown" || byPath["weights.bin"].CacheStatus != "unknown" || byPath["missing.bin"].CacheStatus != "unknown" { + t.Fatalf("files=%+v", m.Files) + } + f, section, err := r.OpenFile("models", "Qwen/demo", "main", "README.md") + if err != nil { + t.Fatal(err) + } + b, err := io.ReadAll(section) + f.Close() + if err != nil || string(b) != "abcdefgh" { + t.Fatalf("payload=%q err=%v", b, err) + } + if _, _, err := r.OpenFile("models", "Qwen/demo", "main", "weights.bin"); !errors.Is(err, ErrIncomplete) { + t.Fatalf("partial returned %v", err) + } + if _, _, err := r.OpenFile("models", "Qwen/demo", "main", "missing.bin"); !errors.Is(err, os.ErrNotExist) { + t.Fatalf("missing returned %v", err) + } + } + if !reflect.DeepEqual(before, snapshot(t, r.Root)) { + t.Fatal("projection changed cache entries") + } + if _, err := os.Stat(filepath.Join(api, "repository.json")); !errors.Is(err, os.ErrNotExist) { + t.Fatal("projection registered repository") + } +} + +func TestIndexReusesSnapshotRefreshesExplicitlyAndKeepsCommitLookup(t *testing.T) { + r, api := fixture(t) + index := NewIndex() + if _, err := index.Refresh(r.Root, "models"); err != nil { + t.Fatal(err) + } + for _, revision := range []string{"main", "commit1"} { + manifest, err := index.Manifest(r.Root, "models", "Qwen/demo", revision) + if err != nil || manifest.Commit != "commit1" || len(manifest.Files) != 3 { + t.Fatalf("revision %s manifest=%+v err=%v", revision, manifest, err) + } + } + if err := os.Remove(filepath.Join(api, "revision", "main", "meta_get.json")); err != nil { + t.Fatal(err) + } + if _, err := index.Manifest(r.Root, "models", "Qwen/demo", "main"); err != nil { + t.Fatalf("indexed manifest unexpectedly rescanned local cache: %v", err) + } + index.Invalidate(r.Root, "models") + if _, err := index.Manifest(r.Root, "models", "Qwen/demo", "main"); !errors.Is(err, os.ErrNotExist) { + t.Fatalf("invalidated index did not rebuild local facts: %v", err) + } +} + +func TestMetadataMissingKeepsObservedFilesAndCorruptionIsExplicit(t *testing.T) { + r, api := fixture(t) + if err := os.Remove(filepath.Join(api, "revision", "main", "meta_get.json")); err != nil { + t.Fatal(err) + } + m, err := r.Manifest("models", "Qwen/demo", "commit1") + if err != nil || m.MetadataAvailable || m.ManifestComplete || len(m.Files) != 3 { + t.Fatalf("manifest=%+v err=%v", m, err) + } + if _, err = r.Manifest("models", "Qwen/demo", "nonexistent"); !errors.Is(err, os.ErrNotExist) { + t.Fatalf("missing revision err=%v", err) + } + write(t, filepath.Join(api, "revision", "main", "meta_get.json"), []byte("broken")) + cat, err := r.Catalog("models") + if err != nil || len(cat.Repos) != 1 || cat.Repos[0].Error == "" { + t.Fatalf("catalog=%+v err=%v", cat, err) + } + if _, err = r.Manifest("models", "Qwen/demo", "main"); err == nil { + t.Fatal("corrupt metadata silently accepted") + } +} +func TestCorruptOrTruncatedBlobNeverComplete(t *testing.T) { + for _, mode := range []string{"truncated-mask", "truncated-body", "huge-mask", "invalid-version"} { + t.Run(mode, func(t *testing.T) { + p := filepath.Join(t.TempDir(), "blob") + blob(t, p, 8, 3) + b, _ := os.ReadFile(p) + switch mode { + case "truncated-mask": + b = b[:36] + case "truncated-body": + b = b[:39] + case "huge-mask": + binary.LittleEndian.PutUint64(b[28:], ^uint64(0)) + case "invalid-version": + binary.LittleEndian.PutUint64(b[4:], 99) + } + os.Chmod(p, 0644) + if err := os.WriteFile(p, b, 0444); err != nil { + t.Fatal(err) + } + f, err := os.Open(p) + if err != nil { + t.Fatal(err) + } + defer f.Close() + status, _, _, _, err := inspect(f, 8) + if err != nil || status == "complete" { + t.Fatalf("status=%s err=%v", status, err) + } + }) + } +} + +func TestUpstreamSiblingsWithoutObjectDetailsStayVisible(t *testing.T) { + r, api := fixture(t) + p := filepath.Join(api, "revision", "main", "meta_get.json") + if err := os.Chmod(p, 0644); err != nil { + t.Fatal(err) + } + response(t, p, map[string]any{"sha": "commit1", "siblings": []any{map[string]any{"rfilename": "not-downloaded.bin"}}}) + m, err := r.Manifest("models", "Qwen/demo", "main") + if err != nil { + t.Fatal(err) + } + for _, f := range m.Files { + if f.Path == "not-downloaded.bin" { + if f.OID != "" || f.Size != -1 || f.CacheStatus != "unknown" { + t.Fatalf("unknown upstream file = %+v", f) + } + return + } + } + t.Fatal("upstream sibling was dropped") +} + +func TestCachedVisibilityAndBranchCommitDeduplication(t *testing.T) { + for _, tc := range []struct { + name string + private any + gated any + want string + }{ + {"public", false, false, "public"}, {"private", true, false, "restricted"}, {"gated", false, "auto", "restricted"}, {"unknown", nil, nil, "unknown"}, + } { + t.Run(tc.name, func(t *testing.T) { + r, api := fixture(t) + p := filepath.Join(api, "revision", "main", "meta_get.json") + os.Chmod(p, 0644) + response(t, p, map[string]any{"sha": "commit1", "private": tc.private, "gated": tc.gated, "siblings": []any{}}) + cat, err := r.Catalog("models") + if err != nil { + t.Fatal(err) + } + if len(cat.Repos) != 1 || cat.Repos[0].Access != tc.want || len(cat.Repos[0].Revisions) != 1 || cat.Repos[0].Revisions[0].Name != "main" { + t.Fatalf("catalog=%+v", cat) + } + m, err := r.Manifest("models", "Qwen/demo", "commit1") + if err != nil || m.Access != tc.want || len(m.Files) != 3 { + t.Fatalf("manifest=%+v err=%v", m, err) + } + }) + } +} + +func TestLayoutNamesAreValidRepoNamesAndBlobsAloneAreNotDiscovery(t *testing.T) { + r := Reader{Root: t.TempDir()} + for _, name := range []string{"revision", "paths-info", "resolve", "blobs"} { + response(t, filepath.Join(r.Root, "api", "models", "acme", name, "revision", "main", "meta_get.json"), map[string]any{"sha": "commit1", "siblings": []any{}}) + } + blob(t, filepath.Join(r.Root, "files", "models", "temporary", "orphan", "blobs", "unreferenced"), 8, 3) + cat, err := r.Catalog("models") + if err != nil { + t.Fatal(err) + } + if len(cat.Repos) != 4 { + t.Fatalf("catalog=%+v", cat) + } + for _, repo := range cat.Repos { + if repo.Repo == "acme" || repo.Repo == "temporary/orphan" { + t.Fatalf("misidentified %s", repo.Repo) + } + } +} + +func TestHiddenCommitPermissionCannotBeDiscardedByBranchDeduplication(t *testing.T) { + r, api := fixture(t) + p := filepath.Join(api, "revision", "main", "meta_get.json") + os.Chmod(p, 0644) + response(t, p, map[string]any{"sha": "commit1", "private": false, "gated": false, "siblings": []any{}}) + response(t, filepath.Join(api, "revision", "commit1", "meta_get.json"), map[string]any{"sha": "commit1", "private": true, "gated": false, "siblings": []any{}}) + cat, err := r.Catalog("models") + if err != nil { + t.Fatal(err) + } + if len(cat.Repos[0].Revisions) != 1 || cat.Repos[0].Access != "restricted" { + t.Fatalf("catalog=%+v", cat) + } +} +func TestProjectionRejectsUnsafeIdentitiesAndEscapingLinks(t *testing.T) { + r, _ := fixture(t) + for _, repo := range []string{"../demo", "dingo-local/alice", "modelscope/Qwen", "a/b/c"} { + if _, err := r.Manifest("models", repo, "main"); err == nil { + t.Fatalf("accepted %s", repo) + } + } + outside := filepath.Join(t.TempDir(), "secret") + write(t, outside, []byte("private")) + p := filepath.Join(r.Root, "files", "models", "Qwen", "demo", "resolve", "commit1", "escape.txt") + os.MkdirAll(filepath.Dir(p), 0755) + if err := os.Symlink(outside, p); err != nil { + t.Skipf("symlinks unavailable: %v", err) + } + manifest, err := r.Manifest("models", "Qwen/demo", "main") + if err != nil { + t.Fatal(err) + } + for _, file := range manifest.Files { + if file.CacheStatus != "unknown" { + t.Fatalf("listing claimed payload availability: %+v", file) + } + } + if f, _, err := r.OpenFile("models", "Qwen/demo", "main", "escape.txt"); err == nil { + f.Close() + t.Fatal("opened external symlink") + } + // A valid resolve link into this repository remains readable. No metadata + // OID is supplied, so this exercises the resolve link rather than blob lookup. + response(t, filepath.Join(r.Root, "api", "models", "Qwen", "demo", "paths-info", "commit1", "linked.txt", "paths-info_post.json"), []any{map[string]any{"size": 8}}) + if err := os.Symlink(filepath.Join(r.Root, "files", "models", "Qwen", "demo", "blobs", "readme"), filepath.Join(filepath.Dir(p), "linked.txt")); err != nil { + t.Fatal(err) + } + f, content, err := r.OpenFile("models", "Qwen/demo", "main", "linked.txt") + if err != nil { + t.Fatal(err) + } + defer f.Close() + b, err := io.ReadAll(content) + if err != nil || string(b) != "abcdefgh" { + t.Fatalf("valid link content=%q err=%v", b, err) + } +} diff --git a/pkg/inventory/protocol.go b/pkg/inventory/protocol.go new file mode 100644 index 0000000..fb550ab --- /dev/null +++ b/pkg/inventory/protocol.go @@ -0,0 +1,53 @@ +// Package inventory owns only uploaded inventory, never remote cache progress. +package inventory + +import ( + "crypto/sha256" + "encoding/hex" + "encoding/json" +) + +type Key struct { + Namespace string `json:"namespace"` + RepoType string `json:"repoType"` + Repo string `json:"repo"` +} + +func (k Key) ID() string { return k.Namespace + "\x00" + k.RepoType + "\x00" + k.Repo } + +type File struct { + Key + Path string `json:"path"` + SHA256 string `json:"sha256"` + Size int64 `json:"size"` +} +type Report struct { + Version int `json:"version"` + InstanceID string `json:"instanceId"` + Epoch string `json:"epoch"` + Sequence uint64 `json:"sequence"` + Baseline bool `json:"baseline"` + Key + Deleted bool `json:"deleted"` + Files []File `json:"files"` +} + +func (r Report) Digest() string { + b, _ := json.Marshal(r) + h := sha256.Sum256(b) + return hex.EncodeToString(h[:]) +} + +type Ack struct { + Epoch string `json:"epoch"` + Sequence uint64 `json:"sequence"` + Digest string `json:"digest"` + Status string `json:"status"` +} +type Session struct { + InstanceID string `json:"instanceId"` + Epoch string `json:"epoch"` + PendingEpoch string `json:"pendingEpoch"` + Status string `json:"status"` + Error string `json:"error,omitempty"` +} diff --git a/pkg/inventory/store.go b/pkg/inventory/store.go new file mode 100644 index 0000000..2c78a8a --- /dev/null +++ b/pkg/inventory/store.go @@ -0,0 +1,305 @@ +package inventory + +import ( + "bytes" + "encoding/json" + "errors" + "fmt" + "os" + "path/filepath" + "runtime" + "sync" + "time" +) + +// Barrier excludes effective metadata mutations while a snapshot is fixed. +// Network calls never hold it. Durable state lives outside every repository. +var Barrier sync.RWMutex +var storeMu sync.Mutex +var Wake = make(chan struct{}, 1) +var ErrDamaged = errors.New("inventory state damaged") + +func Signal() { + select { + case Wake <- struct{}{}: + default: + } +} + +type Operation struct { + Kind string `json:"kind"` + Data json.RawMessage `json:"data"` +} +type Entry struct { + Key + Sequence uint64 `json:"sequence"` + Deleted bool `json:"deleted"` + Operation *Operation `json:"operation,omitempty"` + Report *Report `json:"report,omitempty"` + Due time.Time `json:"due"` + Attempts int `json:"attempts"` + Error string `json:"error,omitempty"` +} +type State struct { + BuildAttempts int `json:"buildAttempts,omitempty"` + BuildBlocked bool `json:"buildBlocked,omitempty"` + Version int `json:"version"` + NodeID string `json:"nodeId"` + Epoch string `json:"epoch"` + Sequence uint64 `json:"sequence"` + Entries map[string]*Entry `json:"entries"` + ReconcileEpoch string `json:"reconcileEpoch,omitempty"` + Baseline *Report `json:"baseline,omitempty"` + Cutoff uint64 `json:"cutoff"` + RetryAt time.Time `json:"retryAt"` + Attempts int `json:"attempts"` + Error string `json:"error,omitempty"` +} + +func statePath(root string) string { return filepath.Join(root, ".upload-inventory", "outbox.json") } +func load(root string) (*State, error) { + b, e := os.ReadFile(statePath(root)) + if os.IsNotExist(e) { + return &State{Version: 2, Entries: map[string]*Entry{}}, nil + } + if e != nil { + return nil, e + } + var s State + if e = json.Unmarshal(b, &s); e != nil { + return nil, fmt.Errorf("%w; explicit recovery required: %v", ErrDamaged, e) + } + if s.Version != 2 || s.Entries == nil { + return nil, ErrDamaged + } + return &s, nil +} +func Read(root string) (*State, error) { storeMu.Lock(); defer storeMu.Unlock(); return load(root) } +func Update(root string, fn func(*State) error) error { + storeMu.Lock() + defer storeMu.Unlock() + s, e := load(root) + if e != nil { + return e + } + before, e := json.Marshal(s) + if e != nil { + return e + } + if e = fn(s); e != nil { + return e + } + b, e := json.Marshal(s) + if e != nil { + return e + } + if bytes.Equal(before, b) { + return nil + } + return AtomicWrite(statePath(root), b) +} +func AtomicWrite(path string, b []byte) error { + dir := filepath.Dir(path) + if e := os.MkdirAll(dir, 0700); e != nil { + return e + } + f, e := os.CreateTemp(dir, ".inventory-*") + if e != nil { + return e + } + name := f.Name() + defer os.Remove(name) + if _, e = f.Write(b); e == nil { + e = f.Sync() + } + ce := f.Close() + if e != nil { + return e + } + if ce != nil { + return ce + } + if e = os.Rename(name, path); e != nil { + return e + } + return SyncParents(filepath.Dir(dir), path) +} + +// SyncParents makes newly created directories and renamed/deleted entries +// durable before their write-ahead operation is acknowledged locally. +func SyncParents(root, path string) error { + if runtime.GOOS == "windows" { + return nil + } // directory fsync is unavailable here + root, err := filepath.Abs(root) + if err != nil { + return err + } + dir, err := filepath.Abs(filepath.Dir(path)) + if err != nil { + return err + } + for { + rel, err := filepath.Rel(root, dir) + if err != nil || rel == ".." || len(rel) > 3 && rel[:3] == ".."+string(filepath.Separator) { + return errors.New("inventory sync outside repository root") + } + f, err := os.Open(dir) + if err == nil { + err = f.Sync() + _ = f.Close() + if err != nil { + return err + } + } else if !os.IsNotExist(err) { + return err + } + if dir == root { + return nil + } + dir = filepath.Dir(dir) + } +} + +// Run journals before the first effective change. A failed completion write leaves +// a replayable operation, not an unrecorded successful local mutation. +func Run(root string, k Key, kind string, data any, fn func() error) error { + Barrier.RLock() + defer Barrier.RUnlock() + b, e := json.Marshal(data) + if e != nil { + return e + } + e = Update(root, func(s *State) error { + old := s.Entries[k.ID()] + if old != nil && old.Operation != nil { + return errors.New("repository operation awaits recovery") + } + if s.Sequence == ^uint64(0) { + return errors.New("inventory sequence exhausted") + } + s.Sequence++ + if old == nil { + old = &Entry{Key: k} + s.Entries[k.ID()] = old + } + old.Sequence = s.Sequence + old.Operation = &Operation{Kind: kind, Data: b} + // The first dirty deadline is not pushed back by continuous mutations. + if old.Due.IsZero() { + old.Due = time.Now().Add(300 * time.Millisecond) + } + return nil + }) + if e != nil { + return e + } + defer Signal() + if e = fn(); e != nil { + return e + } + // The business effect is already committed. A failed journal cleanup is retried. + _ = Finish(root, k, kind) + return nil +} +func Finish(root string, k Key, kind string) error { + return Update(root, func(s *State) error { + e := s.Entries[k.ID()] + if e == nil { + return errors.New("missing operation record") + } + e.Operation = nil + e.Deleted = kind == "delete-repository" + e.Error = "" + return nil + }) +} +func RequestReconcile(root, epoch string) error { + if epoch == "" { + return errors.New("missing reconciliation epoch") + } + err := Update(root, func(s *State) error { + if s.Epoch == epoch { + return nil + } + if s.ReconcileEpoch != epoch { + s.ReconcileEpoch = epoch + s.Baseline = nil + s.BuildAttempts = 0 + s.BuildBlocked = false + s.Attempts = 0 + s.RetryAt = time.Time{} + s.Error = "" + } + return nil + }) + if errors.Is(err, ErrDamaged) { + // Only an explicit request can replace corrupt state. Preserve its bytes + // for investigation instead of silently treating corruption as an empty queue. + storeMu.Lock() + old, readErr := os.ReadFile(statePath(root)) + if readErr == nil { + backup := statePath(root) + ".damaged-" + time.Now().UTC().Format("20060102T150405.000000000") + readErr = AtomicWrite(backup, old) + } + if readErr == nil { + b, _ := json.Marshal(&State{Version: 2, Entries: map[string]*Entry{}, ReconcileEpoch: epoch}) + readErr = AtomicWrite(statePath(root), b) + } + storeMu.Unlock() + err = readErr + } + if err == nil { + Signal() + } + return err +} + +// Confirm consumes exactly the report acknowledged. New mutations survive. +func Confirm(root string, p *Report, a Ack) error { + if a.Epoch != p.Epoch || a.Sequence != p.Sequence || a.Digest != p.Digest() || (a.Status != "accepted" && a.Status != "confirmed" && a.Status != "obsolete") { + return errors.New("inventory acknowledgement mismatch") + } + return Update(root, func(s *State) error { + if p.Baseline { + if s.Baseline == nil || s.Baseline.Digest() != p.Digest() { + return errors.New("baseline changed") + } + // Epoch fencing makes reset safe. Rebase post-snapshot changes only. + cutoff := s.Cutoff + s.Sequence -= cutoff + for id, e := range s.Entries { + if e.Sequence <= cutoff { + delete(s.Entries, id) + } else { + e.Sequence -= cutoff + e.Report = nil + e.Attempts = 0 + e.Due = time.Now() + } + } + s.Epoch = p.Epoch + s.BuildAttempts = 0 + s.BuildBlocked = false + s.Baseline = nil + s.ReconcileEpoch = "" + s.Cutoff = 0 + s.Attempts = 0 + s.RetryAt = time.Time{} + s.Error = "" + return nil + } + e := s.Entries[p.Key.ID()] + if e == nil || e.Report == nil || e.Report.Digest() != p.Digest() { + return errors.New("pending report changed") + } + if e.Sequence == p.Sequence && e.Operation == nil { + delete(s.Entries, p.Key.ID()) + } else { + e.Report = nil + e.Attempts = 0 + e.Due = time.Now() + } + return nil + }) +} diff --git a/pkg/inventory/store_test.go b/pkg/inventory/store_test.go new file mode 100644 index 0000000..3430e26 --- /dev/null +++ b/pkg/inventory/store_test.go @@ -0,0 +1,93 @@ +package inventory + +import ( + "errors" + "testing" +) + +func TestDurableReportConfirmationAndResetKeepNewMutations(t *testing.T) { + root := t.TempDir() + k := Key{Namespace: "team", RepoType: "models", Repo: "full/name"} + change := func() { + t.Helper() + if err := Run(root, k, "metadata", map[string]string{"commit": "fixed"}, func() error { return nil }); err != nil { + t.Fatal(err) + } + } + change() + var first *Report + if err := Update(root, func(s *State) error { + s.Epoch = "old" + s.NodeID = "A" + first = &Report{Version: 2, InstanceID: "A", Epoch: "old", Sequence: s.Sequence, Key: k} + s.Entries[k.ID()].Report = first + return nil + }); err != nil { + t.Fatal(err) + } + change() + q, err := Read(root) + if err != nil { + t.Fatal(err) + } + if q.Entries[k.ID()].Report.Digest() != first.Digest() { + t.Fatal("mutation rewrote fixed report") + } + ack := func(p *Report) Ack { + return Ack{Epoch: p.Epoch, Sequence: p.Sequence, Digest: p.Digest(), Status: "accepted"} + } + if err = Confirm(root, first, ack(first)); err != nil { + t.Fatal(err) + } + q, _ = Read(root) + if q.Entries[k.ID()] == nil || q.Entries[k.ID()].Sequence != 2 { + t.Fatal("old ACK erased new mutation") + } + base := &Report{Version: 2, InstanceID: "A", Epoch: "new", Baseline: true} + if err = Update(root, func(s *State) error { s.Baseline = base; s.Cutoff = s.Sequence; s.ReconcileEpoch = "new"; return nil }); err != nil { + t.Fatal(err) + } + change() + if err = Confirm(root, base, ack(base)); err != nil { + t.Fatal(err) + } + q, _ = Read(root) + if q.Epoch != "new" || q.Sequence != 1 || q.Entries[k.ID()].Sequence != 1 || q.Entries[k.ID()].Report != nil { + t.Fatalf("post-baseline mutation lost: %+v", q) + } + // A reset with no concurrent changes leaves zero sequence and an empty queue. + base.Epoch = "next" + if err = Update(root, func(s *State) error { s.Baseline = base; s.Cutoff = s.Sequence; return nil }); err != nil { + t.Fatal(err) + } + if err = Confirm(root, base, ack(base)); err != nil { + t.Fatal(err) + } + q, _ = Read(root) + if q.Sequence != 0 || len(q.Entries) != 0 { + t.Fatal("manual reconciliation did not clear covered queue") + } +} + +func TestFailedOperationRemainsRecoverableOutsideRepository(t *testing.T) { + root := t.TempDir() + k := Key{Namespace: "team", RepoType: "models", Repo: "gone"} + err := Run(root, k, "delete-repository", k, func() error { return errors.New("crash after filesystem change") }) + if err == nil { + t.Fatal("expected failure") + } + q, err := Read(root) + if err != nil { + t.Fatal(err) + } + if q.Entries[k.ID()].Operation == nil { + t.Fatal("lost recovery intent") + } + if err = Finish(root, k, "delete-repository"); err != nil { + t.Fatal(err) + } + q, _ = Read(root) + if !q.Entries[k.ID()].Deleted || q.Entries[k.ID()].Operation != nil { + t.Fatal("deletion tombstone missing") + } +} diff --git a/pkg/middleware/route_params.go b/pkg/middleware/route_params.go new file mode 100644 index 0000000..043a743 --- /dev/null +++ b/pkg/middleware/route_params.go @@ -0,0 +1,25 @@ +package middleware + +import ( + "github.com/labstack/echo/v4" + "net/url" +) + +// Echo routes against RawPath when it is present. Only in that case are the +// route values still escaped; URL.Path and query values are already decoded. +func DecodeRouteParams(next echo.HandlerFunc) echo.HandlerFunc { + return func(c echo.Context) error { + if c.Request().URL.RawPath != "" { + values := append([]string(nil), c.ParamValues()...) + for i, value := range values { + decoded, err := url.PathUnescape(value) + if err != nil { + return echo.NewHTTPError(400, "invalid route encoding") + } + values[i] = decoded + } + c.SetParamValues(values...) + } + return next(c) + } +} diff --git a/pkg/middleware/upload_guard.go b/pkg/middleware/upload_guard.go index f67aa14..05704c1 100644 --- a/pkg/middleware/upload_guard.go +++ b/pkg/middleware/upload_guard.go @@ -21,23 +21,8 @@ import ( "go.uber.org/zap" ) -// UploadGuardMiddleware 挡住浏览器发起的跨站请求。 -// -// 上传口没有任何身份校验，安全性完全建立在“只有本机的 ingest agent 会调它” -// 这个假设上。但绑在 127.0.0.1 挡不住浏览器——用户的浏览器就在回环上。 -// -// 具体的攻击路径：上传口的写接口全部读裸 body、不校验 Content-Type， -// 于是一个 text/plain 的 POST 就能带着 JSON 打进来。而 text/plain 的 POST -// 属于 CORS 简单请求，不触发预检，浏览器会直接送出去。用户访问任意一个恶意 -// 页面，那个页面就能 POST /api/cache/orphans/delete 把缓存删掉。 -// -// 这里用 Origin 头来区分：浏览器发起的跨站请求一定带 Origin（简单请求也带）， -// 而 ingest agent、curl 这类服务端调用方不会带。因此拒掉“带 Origin 的写请求” -// 既能挡住浏览器，又不需要任何调用方改代码。Sec-Fetch-Site 是同一判断的补强， -// 现代浏览器都会发，且是禁止修改的头。 -// -// 这不能替代真正的鉴权，只是把浏览器这条路堵死。任何能直接发 HTTP 的进程 -// 仍然可以访问上传口——那个问题要靠给上传口加回身份校验来解决。 +// UploadGuardMiddleware 使用 Origin / Sec-Fetch-Site 拒绝浏览器跨站写入。 +// 浏览器简单请求也带 Origin，因此这里不依赖 CORS 预检或 Content-Type。 func UploadGuardMiddleware() echo.MiddlewareFunc { return func(next echo.HandlerFunc) echo.HandlerFunc { return func(c echo.Context) error { diff --git a/pkg/proto/manager.proto b/pkg/proto/manager.proto index 2813c76..286db6c 100644 --- a/pkg/proto/manager.proto +++ b/pkg/proto/manager.proto @@ -5,8 +5,11 @@ import "google/protobuf/empty.proto"; option go_package = ".;manager"; +// Existing RPC org/repo keep legacy remote coordinates; hosted org is dingo-local/namespace. + // 用户服务定义 service Manager { + rpc IngestRepository (IngestRepositoryRequest) returns (IngestRepositoryResponse); // 启动注册方法 rpc Register (RegisterRequest) returns (RegisterResponse) {}; // 心跳方法 @@ -34,6 +37,8 @@ message RegisterRequest { string host = 2; int32 port = 3; bool online = 4; + string managementUrl = 5; + string downloadUrl = 6; } // 注册响应 @@ -47,6 +52,32 @@ message HeartbeatRequest { int32 id = 1; string instanceId = 2; bool online = 3; + NodeHealthSnapshot health = 4; +} + +// In-memory observations only; absence means the client does not report health. +message NodeHealthSnapshot { + uint32 version = 1; + string process_id = 2; + int64 started_at = 3; + int64 collected_at = 4; + uint32 heartbeat_period_seconds = 5; + repeated CapabilityObservation capabilities = 6; + repeated StorageErrorCount errors = 7; +} + +message CapabilityObservation { + string capability = 1; + int32 state = 2; // 0 unknown, 1 observed healthy, 2 suspect, 3 degraded, 4 recovering + bool unresolved = 3; + int64 last_observation = 4; + int64 last_failure = 5; +} + +message StorageErrorCount { + string operation = 1; + string kind = 2; + uint64 count = 3; // Cumulative within process_id, not an active-failure count. } message SchedulerFileRequest { @@ -135,4 +166,15 @@ message UpdateRepositoryMountStatusReq{ int64 id = 1; int32 status = 2; string errorMsg = 3; -} \ No newline at end of file +} + +message IngestRepositoryRequest { + string instance_id = 1; + string namespace = 2; + string repo_type = 3; + string repo = 4; + string revision = 5; + string commit = 6; + bool online = 7; +} +message IngestRepositoryResponse { int64 repository_id = 1; string commit = 2; int64 file_count = 3; int64 used_storage = 4; } diff --git a/pkg/proto/manager/manager.pb.go b/pkg/proto/manager/manager.pb.go index 2d47159..c094ab3 100644 --- a/pkg/proto/manager/manager.pb.go +++ b/pkg/proto/manager/manager.pb.go @@ -1,7 +1,7 @@ // Code generated by protoc-gen-go. DO NOT EDIT. // versions: -// protoc-gen-go v1.36.5 -// protoc v5.29.3 +// protoc-gen-go v1.36.6 +// protoc v7.35.1 // source: manager.proto package manager @@ -29,6 +29,8 @@ type RegisterRequest struct { Host string `protobuf:"bytes,2,opt,name=host,proto3" json:"host,omitempty"` Port int32 `protobuf:"varint,3,opt,name=port,proto3" json:"port,omitempty"` Online bool `protobuf:"varint,4,opt,name=online,proto3" json:"online,omitempty"` + ManagementUrl string `protobuf:"bytes,5,opt,name=managementUrl,proto3" json:"managementUrl,omitempty"` + DownloadUrl string `protobuf:"bytes,6,opt,name=downloadUrl,proto3" json:"downloadUrl,omitempty"` unknownFields protoimpl.UnknownFields sizeCache protoimpl.SizeCache } @@ -91,6 +93,20 @@ func (x *RegisterRequest) GetOnline() bool { return false } +func (x *RegisterRequest) GetManagementUrl() string { + if x != nil { + return x.ManagementUrl + } + return "" +} + +func (x *RegisterRequest) GetDownloadUrl() string { + if x != nil { + return x.DownloadUrl + } + return "" +} + // 注册响应 type RegisterResponse struct { state protoimpl.MessageState `protogen:"open.v1"` @@ -150,6 +166,7 @@ type HeartbeatRequest struct { Id int32 `protobuf:"varint,1,opt,name=id,proto3" json:"id,omitempty"` InstanceId string `protobuf:"bytes,2,opt,name=instanceId,proto3" json:"instanceId,omitempty"` Online bool `protobuf:"varint,3,opt,name=online,proto3" json:"online,omitempty"` + Health *NodeHealthSnapshot `protobuf:"bytes,4,opt,name=health,proto3" json:"health,omitempty"` unknownFields protoimpl.UnknownFields sizeCache protoimpl.SizeCache } @@ -205,6 +222,242 @@ func (x *HeartbeatRequest) GetOnline() bool { return false } +func (x *HeartbeatRequest) GetHealth() *NodeHealthSnapshot { + if x != nil { + return x.Health + } + return nil +} + +// In-memory observations only; absence means the client does not report health. +type NodeHealthSnapshot struct { + state protoimpl.MessageState `protogen:"open.v1"` + Version uint32 `protobuf:"varint,1,opt,name=version,proto3" json:"version,omitempty"` + ProcessId string `protobuf:"bytes,2,opt,name=process_id,json=processId,proto3" json:"process_id,omitempty"` + StartedAt int64 `protobuf:"varint,3,opt,name=started_at,json=startedAt,proto3" json:"started_at,omitempty"` + CollectedAt int64 `protobuf:"varint,4,opt,name=collected_at,json=collectedAt,proto3" json:"collected_at,omitempty"` + HeartbeatPeriodSeconds uint32 `protobuf:"varint,5,opt,name=heartbeat_period_seconds,json=heartbeatPeriodSeconds,proto3" json:"heartbeat_period_seconds,omitempty"` + Capabilities []*CapabilityObservation `protobuf:"bytes,6,rep,name=capabilities,proto3" json:"capabilities,omitempty"` + Errors []*StorageErrorCount `protobuf:"bytes,7,rep,name=errors,proto3" json:"errors,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *NodeHealthSnapshot) Reset() { + *x = NodeHealthSnapshot{} + mi := &file_manager_proto_msgTypes[3] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *NodeHealthSnapshot) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*NodeHealthSnapshot) ProtoMessage() {} + +func (x *NodeHealthSnapshot) ProtoReflect() protoreflect.Message { + mi := &file_manager_proto_msgTypes[3] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use NodeHealthSnapshot.ProtoReflect.Descriptor instead. +func (*NodeHealthSnapshot) Descriptor() ([]byte, []int) { + return file_manager_proto_rawDescGZIP(), []int{3} +} + +func (x *NodeHealthSnapshot) GetVersion() uint32 { + if x != nil { + return x.Version + } + return 0 +} + +func (x *NodeHealthSnapshot) GetProcessId() string { + if x != nil { + return x.ProcessId + } + return "" +} + +func (x *NodeHealthSnapshot) GetStartedAt() int64 { + if x != nil { + return x.StartedAt + } + return 0 +} + +func (x *NodeHealthSnapshot) GetCollectedAt() int64 { + if x != nil { + return x.CollectedAt + } + return 0 +} + +func (x *NodeHealthSnapshot) GetHeartbeatPeriodSeconds() uint32 { + if x != nil { + return x.HeartbeatPeriodSeconds + } + return 0 +} + +func (x *NodeHealthSnapshot) GetCapabilities() []*CapabilityObservation { + if x != nil { + return x.Capabilities + } + return nil +} + +func (x *NodeHealthSnapshot) GetErrors() []*StorageErrorCount { + if x != nil { + return x.Errors + } + return nil +} + +type CapabilityObservation struct { + state protoimpl.MessageState `protogen:"open.v1"` + Capability string `protobuf:"bytes,1,opt,name=capability,proto3" json:"capability,omitempty"` + State int32 `protobuf:"varint,2,opt,name=state,proto3" json:"state,omitempty"` // 0 unknown, 1 observed healthy, 2 suspect, 3 degraded, 4 recovering + Unresolved bool `protobuf:"varint,3,opt,name=unresolved,proto3" json:"unresolved,omitempty"` + LastObservation int64 `protobuf:"varint,4,opt,name=last_observation,json=lastObservation,proto3" json:"last_observation,omitempty"` + LastFailure int64 `protobuf:"varint,5,opt,name=last_failure,json=lastFailure,proto3" json:"last_failure,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *CapabilityObservation) Reset() { + *x = CapabilityObservation{} + mi := &file_manager_proto_msgTypes[4] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *CapabilityObservation) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*CapabilityObservation) ProtoMessage() {} + +func (x *CapabilityObservation) ProtoReflect() protoreflect.Message { + mi := &file_manager_proto_msgTypes[4] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use CapabilityObservation.ProtoReflect.Descriptor instead. +func (*CapabilityObservation) Descriptor() ([]byte, []int) { + return file_manager_proto_rawDescGZIP(), []int{4} +} + +func (x *CapabilityObservation) GetCapability() string { + if x != nil { + return x.Capability + } + return "" +} + +func (x *CapabilityObservation) GetState() int32 { + if x != nil { + return x.State + } + return 0 +} + +func (x *CapabilityObservation) GetUnresolved() bool { + if x != nil { + return x.Unresolved + } + return false +} + +func (x *CapabilityObservation) GetLastObservation() int64 { + if x != nil { + return x.LastObservation + } + return 0 +} + +func (x *CapabilityObservation) GetLastFailure() int64 { + if x != nil { + return x.LastFailure + } + return 0 +} + +type StorageErrorCount struct { + state protoimpl.MessageState `protogen:"open.v1"` + Operation string `protobuf:"bytes,1,opt,name=operation,proto3" json:"operation,omitempty"` + Kind string `protobuf:"bytes,2,opt,name=kind,proto3" json:"kind,omitempty"` + Count uint64 `protobuf:"varint,3,opt,name=count,proto3" json:"count,omitempty"` // Cumulative within process_id, not an active-failure count. + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *StorageErrorCount) Reset() { + *x = StorageErrorCount{} + mi := &file_manager_proto_msgTypes[5] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *StorageErrorCount) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*StorageErrorCount) ProtoMessage() {} + +func (x *StorageErrorCount) ProtoReflect() protoreflect.Message { + mi := &file_manager_proto_msgTypes[5] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use StorageErrorCount.ProtoReflect.Descriptor instead. +func (*StorageErrorCount) Descriptor() ([]byte, []int) { + return file_manager_proto_rawDescGZIP(), []int{5} +} + +func (x *StorageErrorCount) GetOperation() string { + if x != nil { + return x.Operation + } + return "" +} + +func (x *StorageErrorCount) GetKind() string { + if x != nil { + return x.Kind + } + return "" +} + +func (x *StorageErrorCount) GetCount() uint64 { + if x != nil { + return x.Count + } + return 0 +} + type SchedulerFileRequest struct { state protoimpl.MessageState `protogen:"open.v1"` DataType string `protobuf:"bytes,1,opt,name=dataType,proto3" json:"dataType,omitempty"` @@ -222,7 +475,7 @@ type SchedulerFileRequest struct { func (x *SchedulerFileRequest) Reset() { *x = SchedulerFileRequest{} - mi := &file_manager_proto_msgTypes[3] + mi := &file_manager_proto_msgTypes[6] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -234,7 +487,7 @@ func (x *SchedulerFileRequest) String() string { func (*SchedulerFileRequest) ProtoMessage() {} func (x *SchedulerFileRequest) ProtoReflect() protoreflect.Message { - mi := &file_manager_proto_msgTypes[3] + mi := &file_manager_proto_msgTypes[6] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -247,7 +500,7 @@ func (x *SchedulerFileRequest) ProtoReflect() protoreflect.Message { // Deprecated: Use SchedulerFileRequest.ProtoReflect.Descriptor instead. func (*SchedulerFileRequest) Descriptor() ([]byte, []int) { - return file_manager_proto_rawDescGZIP(), []int{3} + return file_manager_proto_rawDescGZIP(), []int{6} } func (x *SchedulerFileRequest) GetDataType() string { @@ -322,7 +575,7 @@ type SyncFileProcessReq struct { func (x *SyncFileProcessReq) Reset() { *x = SyncFileProcessReq{} - mi := &file_manager_proto_msgTypes[4] + mi := &file_manager_proto_msgTypes[7] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -334,7 +587,7 @@ func (x *SyncFileProcessReq) String() string { func (*SyncFileProcessReq) ProtoMessage() {} func (x *SyncFileProcessReq) ProtoReflect() protoreflect.Message { - mi := &file_manager_proto_msgTypes[4] + mi := &file_manager_proto_msgTypes[7] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -347,7 +600,7 @@ func (x *SyncFileProcessReq) ProtoReflect() protoreflect.Message { // Deprecated: Use SyncFileProcessReq.ProtoReflect.Descriptor instead. func (*SyncFileProcessReq) Descriptor() ([]byte, []int) { - return file_manager_proto_rawDescGZIP(), []int{4} + return file_manager_proto_rawDescGZIP(), []int{7} } func (x *SyncFileProcessReq) GetFileProcessEntries() []*FileProcessEntry { @@ -376,7 +629,7 @@ type FileProcessEntry struct { func (x *FileProcessEntry) Reset() { *x = FileProcessEntry{} - mi := &file_manager_proto_msgTypes[5] + mi := &file_manager_proto_msgTypes[8] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -388,7 +641,7 @@ func (x *FileProcessEntry) String() string { func (*FileProcessEntry) ProtoMessage() {} func (x *FileProcessEntry) ProtoReflect() protoreflect.Message { - mi := &file_manager_proto_msgTypes[5] + mi := &file_manager_proto_msgTypes[8] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -401,7 +654,7 @@ func (x *FileProcessEntry) ProtoReflect() protoreflect.Message { // Deprecated: Use FileProcessEntry.ProtoReflect.Descriptor instead. func (*FileProcessEntry) Descriptor() ([]byte, []int) { - return file_manager_proto_rawDescGZIP(), []int{5} + return file_manager_proto_rawDescGZIP(), []int{8} } func (x *FileProcessEntry) GetDataType() string { @@ -496,7 +749,7 @@ type SchedulerFileResponse struct { func (x *SchedulerFileResponse) Reset() { *x = SchedulerFileResponse{} - mi := &file_manager_proto_msgTypes[6] + mi := &file_manager_proto_msgTypes[9] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -508,7 +761,7 @@ func (x *SchedulerFileResponse) String() string { func (*SchedulerFileResponse) ProtoMessage() {} func (x *SchedulerFileResponse) ProtoReflect() protoreflect.Message { - mi := &file_manager_proto_msgTypes[6] + mi := &file_manager_proto_msgTypes[9] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -521,7 +774,7 @@ func (x *SchedulerFileResponse) ProtoReflect() protoreflect.Message { // Deprecated: Use SchedulerFileResponse.ProtoReflect.Descriptor instead. func (*SchedulerFileResponse) Descriptor() ([]byte, []int) { - return file_manager_proto_rawDescGZIP(), []int{6} + return file_manager_proto_rawDescGZIP(), []int{9} } func (x *SchedulerFileResponse) GetSchedulerType() int32 { @@ -578,7 +831,7 @@ type FileProcessRequest struct { func (x *FileProcessRequest) Reset() { *x = FileProcessRequest{} - mi := &file_manager_proto_msgTypes[7] + mi := &file_manager_proto_msgTypes[10] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -590,7 +843,7 @@ func (x *FileProcessRequest) String() string { func (*FileProcessRequest) ProtoMessage() {} func (x *FileProcessRequest) ProtoReflect() protoreflect.Message { - mi := &file_manager_proto_msgTypes[7] + mi := &file_manager_proto_msgTypes[10] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -603,7 +856,7 @@ func (x *FileProcessRequest) ProtoReflect() protoreflect.Message { // Deprecated: Use FileProcessRequest.ProtoReflect.Descriptor instead. func (*FileProcessRequest) Descriptor() ([]byte, []int) { - return file_manager_proto_rawDescGZIP(), []int{7} + return file_manager_proto_rawDescGZIP(), []int{10} } func (x *FileProcessRequest) GetProcessId() int64 { @@ -648,7 +901,7 @@ type DeleteByEtagsAndFieldsRequest struct { func (x *DeleteByEtagsAndFieldsRequest) Reset() { *x = DeleteByEtagsAndFieldsRequest{} - mi := &file_manager_proto_msgTypes[8] + mi := &file_manager_proto_msgTypes[11] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -660,7 +913,7 @@ func (x *DeleteByEtagsAndFieldsRequest) String() string { func (*DeleteByEtagsAndFieldsRequest) ProtoMessage() {} func (x *DeleteByEtagsAndFieldsRequest) ProtoReflect() protoreflect.Message { - mi := &file_manager_proto_msgTypes[8] + mi := &file_manager_proto_msgTypes[11] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -673,7 +926,7 @@ func (x *DeleteByEtagsAndFieldsRequest) ProtoReflect() protoreflect.Message { // Deprecated: Use DeleteByEtagsAndFieldsRequest.ProtoReflect.Descriptor instead. func (*DeleteByEtagsAndFieldsRequest) Descriptor() ([]byte, []int) { - return file_manager_proto_rawDescGZIP(), []int{8} + return file_manager_proto_rawDescGZIP(), []int{11} } func (x *DeleteByEtagsAndFieldsRequest) GetEtag() string { @@ -734,7 +987,7 @@ type CreateCacheJobReq struct { func (x *CreateCacheJobReq) Reset() { *x = CreateCacheJobReq{} - mi := &file_manager_proto_msgTypes[9] + mi := &file_manager_proto_msgTypes[12] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -746,7 +999,7 @@ func (x *CreateCacheJobReq) String() string { func (*CreateCacheJobReq) ProtoMessage() {} func (x *CreateCacheJobReq) ProtoReflect() protoreflect.Message { - mi := &file_manager_proto_msgTypes[9] + mi := &file_manager_proto_msgTypes[12] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -759,7 +1012,7 @@ func (x *CreateCacheJobReq) ProtoReflect() protoreflect.Message { // Deprecated: Use CreateCacheJobReq.ProtoReflect.Descriptor instead. func (*CreateCacheJobReq) Descriptor() ([]byte, []int) { - return file_manager_proto_rawDescGZIP(), []int{9} + return file_manager_proto_rawDescGZIP(), []int{12} } func (x *CreateCacheJobReq) GetType() int32 { @@ -828,7 +1081,7 @@ type CreateCacheJobResp struct { func (x *CreateCacheJobResp) Reset() { *x = CreateCacheJobResp{} - mi := &file_manager_proto_msgTypes[10] + mi := &file_manager_proto_msgTypes[13] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -840,7 +1093,7 @@ func (x *CreateCacheJobResp) String() string { func (*CreateCacheJobResp) ProtoMessage() {} func (x *CreateCacheJobResp) ProtoReflect() protoreflect.Message { - mi := &file_manager_proto_msgTypes[10] + mi := &file_manager_proto_msgTypes[13] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -853,7 +1106,7 @@ func (x *CreateCacheJobResp) ProtoReflect() protoreflect.Message { // Deprecated: Use CreateCacheJobResp.ProtoReflect.Descriptor instead. func (*CreateCacheJobResp) Descriptor() ([]byte, []int) { - return file_manager_proto_rawDescGZIP(), []int{10} + return file_manager_proto_rawDescGZIP(), []int{13} } func (x *CreateCacheJobResp) GetId() int64 { @@ -878,7 +1131,7 @@ type UpdateCacheJobStatusReq struct { func (x *UpdateCacheJobStatusReq) Reset() { *x = UpdateCacheJobStatusReq{} - mi := &file_manager_proto_msgTypes[11] + mi := &file_manager_proto_msgTypes[14] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -890,7 +1143,7 @@ func (x *UpdateCacheJobStatusReq) String() string { func (*UpdateCacheJobStatusReq) ProtoMessage() {} func (x *UpdateCacheJobStatusReq) ProtoReflect() protoreflect.Message { - mi := &file_manager_proto_msgTypes[11] + mi := &file_manager_proto_msgTypes[14] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -903,7 +1156,7 @@ func (x *UpdateCacheJobStatusReq) ProtoReflect() protoreflect.Message { // Deprecated: Use UpdateCacheJobStatusReq.ProtoReflect.Descriptor instead. func (*UpdateCacheJobStatusReq) Descriptor() ([]byte, []int) { - return file_manager_proto_rawDescGZIP(), []int{11} + return file_manager_proto_rawDescGZIP(), []int{14} } func (x *UpdateCacheJobStatusReq) GetId() int64 { @@ -966,7 +1219,7 @@ type UpdateRepositoryMountStatusReq struct { func (x *UpdateRepositoryMountStatusReq) Reset() { *x = UpdateRepositoryMountStatusReq{} - mi := &file_manager_proto_msgTypes[12] + mi := &file_manager_proto_msgTypes[15] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -978,7 +1231,7 @@ func (x *UpdateRepositoryMountStatusReq) String() string { func (*UpdateRepositoryMountStatusReq) ProtoMessage() {} func (x *UpdateRepositoryMountStatusReq) ProtoReflect() protoreflect.Message { - mi := &file_manager_proto_msgTypes[12] + mi := &file_manager_proto_msgTypes[15] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -991,7 +1244,7 @@ func (x *UpdateRepositoryMountStatusReq) ProtoReflect() protoreflect.Message { // Deprecated: Use UpdateRepositoryMountStatusReq.ProtoReflect.Descriptor instead. func (*UpdateRepositoryMountStatusReq) Descriptor() ([]byte, []int) { - return file_manager_proto_rawDescGZIP(), []int{12} + return file_manager_proto_rawDescGZIP(), []int{15} } func (x *UpdateRepositoryMountStatusReq) GetId() int64 { @@ -1015,182 +1268,317 @@ func (x *UpdateRepositoryMountStatusReq) GetErrorMsg() string { return "" } +type IngestRepositoryRequest struct { + state protoimpl.MessageState `protogen:"open.v1"` + InstanceId string `protobuf:"bytes,1,opt,name=instance_id,json=instanceId,proto3" json:"instance_id,omitempty"` + Namespace string `protobuf:"bytes,2,opt,name=namespace,proto3" json:"namespace,omitempty"` + RepoType string `protobuf:"bytes,3,opt,name=repo_type,json=repoType,proto3" json:"repo_type,omitempty"` + Repo string `protobuf:"bytes,4,opt,name=repo,proto3" json:"repo,omitempty"` + Revision string `protobuf:"bytes,5,opt,name=revision,proto3" json:"revision,omitempty"` + Commit string `protobuf:"bytes,6,opt,name=commit,proto3" json:"commit,omitempty"` + Online bool `protobuf:"varint,7,opt,name=online,proto3" json:"online,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *IngestRepositoryRequest) Reset() { + *x = IngestRepositoryRequest{} + mi := &file_manager_proto_msgTypes[16] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *IngestRepositoryRequest) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*IngestRepositoryRequest) ProtoMessage() {} + +func (x *IngestRepositoryRequest) ProtoReflect() protoreflect.Message { + mi := &file_manager_proto_msgTypes[16] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use IngestRepositoryRequest.ProtoReflect.Descriptor instead. +func (*IngestRepositoryRequest) Descriptor() ([]byte, []int) { + return file_manager_proto_rawDescGZIP(), []int{16} +} + +func (x *IngestRepositoryRequest) GetInstanceId() string { + if x != nil { + return x.InstanceId + } + return "" +} + +func (x *IngestRepositoryRequest) GetNamespace() string { + if x != nil { + return x.Namespace + } + return "" +} + +func (x *IngestRepositoryRequest) GetRepoType() string { + if x != nil { + return x.RepoType + } + return "" +} + +func (x *IngestRepositoryRequest) GetRepo() string { + if x != nil { + return x.Repo + } + return "" +} + +func (x *IngestRepositoryRequest) GetRevision() string { + if x != nil { + return x.Revision + } + return "" +} + +func (x *IngestRepositoryRequest) GetCommit() string { + if x != nil { + return x.Commit + } + return "" +} + +func (x *IngestRepositoryRequest) GetOnline() bool { + if x != nil { + return x.Online + } + return false +} + +type IngestRepositoryResponse struct { + state protoimpl.MessageState `protogen:"open.v1"` + RepositoryId int64 `protobuf:"varint,1,opt,name=repository_id,json=repositoryId,proto3" json:"repository_id,omitempty"` + Commit string `protobuf:"bytes,2,opt,name=commit,proto3" json:"commit,omitempty"` + FileCount int64 `protobuf:"varint,3,opt,name=file_count,json=fileCount,proto3" json:"file_count,omitempty"` + UsedStorage int64 `protobuf:"varint,4,opt,name=used_storage,json=usedStorage,proto3" json:"used_storage,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *IngestRepositoryResponse) Reset() { + *x = IngestRepositoryResponse{} + mi := &file_manager_proto_msgTypes[17] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *IngestRepositoryResponse) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*IngestRepositoryResponse) ProtoMessage() {} + +func (x *IngestRepositoryResponse) ProtoReflect() protoreflect.Message { + mi := &file_manager_proto_msgTypes[17] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use IngestRepositoryResponse.ProtoReflect.Descriptor instead. +func (*IngestRepositoryResponse) Descriptor() ([]byte, []int) { + return file_manager_proto_rawDescGZIP(), []int{17} +} + +func (x *IngestRepositoryResponse) GetRepositoryId() int64 { + if x != nil { + return x.RepositoryId + } + return 0 +} + +func (x *IngestRepositoryResponse) GetCommit() string { + if x != nil { + return x.Commit + } + return "" +} + +func (x *IngestRepositoryResponse) GetFileCount() int64 { + if x != nil { + return x.FileCount + } + return 0 +} + +func (x *IngestRepositoryResponse) GetUsedStorage() int64 { + if x != nil { + return x.UsedStorage + } + return 0 +} + var File_manager_proto protoreflect.FileDescriptor -var file_manager_proto_rawDesc = string([]byte{ - 0x0a, 0x0d, 0x6d, 0x61, 0x6e, 0x61, 0x67, 0x65, 0x72, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x12, - 0x07, 0x6d, 0x61, 0x6e, 0x61, 0x67, 0x65, 0x72, 0x1a, 0x1b, 0x67, 0x6f, 0x6f, 0x67, 0x6c, 0x65, - 0x2f, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x62, 0x75, 0x66, 0x2f, 0x65, 0x6d, 0x70, 0x74, 0x79, 0x2e, - 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x22, 0x71, 0x0a, 0x0f, 0x52, 0x65, 0x67, 0x69, 0x73, 0x74, 0x65, - 0x72, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x12, 0x1e, 0x0a, 0x0a, 0x69, 0x6e, 0x73, 0x74, - 0x61, 0x6e, 0x63, 0x65, 0x49, 0x64, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x0a, 0x69, 0x6e, - 0x73, 0x74, 0x61, 0x6e, 0x63, 0x65, 0x49, 0x64, 0x12, 0x12, 0x0a, 0x04, 0x68, 0x6f, 0x73, 0x74, - 0x18, 0x02, 0x20, 0x01, 0x28, 0x09, 0x52, 0x04, 0x68, 0x6f, 0x73, 0x74, 0x12, 0x12, 0x0a, 0x04, - 0x70, 0x6f, 0x72, 0x74, 0x18, 0x03, 0x20, 0x01, 0x28, 0x05, 0x52, 0x04, 0x70, 0x6f, 0x72, 0x74, - 0x12, 0x16, 0x0a, 0x06, 0x6f, 0x6e, 0x6c, 0x69, 0x6e, 0x65, 0x18, 0x04, 0x20, 0x01, 0x28, 0x08, - 0x52, 0x06, 0x6f, 0x6e, 0x6c, 0x69, 0x6e, 0x65, 0x22, 0x3c, 0x0a, 0x10, 0x52, 0x65, 0x67, 0x69, - 0x73, 0x74, 0x65, 0x72, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x0e, 0x0a, 0x02, - 0x69, 0x64, 0x18, 0x01, 0x20, 0x01, 0x28, 0x05, 0x52, 0x02, 0x69, 0x64, 0x12, 0x18, 0x0a, 0x07, - 0x73, 0x75, 0x63, 0x63, 0x65, 0x73, 0x73, 0x18, 0x02, 0x20, 0x01, 0x28, 0x08, 0x52, 0x07, 0x73, - 0x75, 0x63, 0x63, 0x65, 0x73, 0x73, 0x22, 0x5a, 0x0a, 0x10, 0x48, 0x65, 0x61, 0x72, 0x74, 0x62, - 0x65, 0x61, 0x74, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x12, 0x0e, 0x0a, 0x02, 0x69, 0x64, - 0x18, 0x01, 0x20, 0x01, 0x28, 0x05, 0x52, 0x02, 0x69, 0x64, 0x12, 0x1e, 0x0a, 0x0a, 0x69, 0x6e, - 0x73, 0x74, 0x61, 0x6e, 0x63, 0x65, 0x49, 0x64, 0x18, 0x02, 0x20, 0x01, 0x28, 0x09, 0x52, 0x0a, - 0x69, 0x6e, 0x73, 0x74, 0x61, 0x6e, 0x63, 0x65, 0x49, 0x64, 0x12, 0x16, 0x0a, 0x06, 0x6f, 0x6e, - 0x6c, 0x69, 0x6e, 0x65, 0x18, 0x03, 0x20, 0x01, 0x28, 0x08, 0x52, 0x06, 0x6f, 0x6e, 0x6c, 0x69, - 0x6e, 0x65, 0x22, 0xf0, 0x01, 0x0a, 0x14, 0x53, 0x63, 0x68, 0x65, 0x64, 0x75, 0x6c, 0x65, 0x72, - 0x46, 0x69, 0x6c, 0x65, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x12, 0x1a, 0x0a, 0x08, 0x64, - 0x61, 0x74, 0x61, 0x54, 0x79, 0x70, 0x65, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x08, 0x64, - 0x61, 0x74, 0x61, 0x54, 0x79, 0x70, 0x65, 0x12, 0x10, 0x0a, 0x03, 0x6f, 0x72, 0x67, 0x18, 0x02, - 0x20, 0x01, 0x28, 0x09, 0x52, 0x03, 0x6f, 0x72, 0x67, 0x12, 0x12, 0x0a, 0x04, 0x72, 0x65, 0x70, - 0x6f, 0x18, 0x03, 0x20, 0x01, 0x28, 0x09, 0x52, 0x04, 0x72, 0x65, 0x70, 0x6f, 0x12, 0x12, 0x0a, - 0x04, 0x6e, 0x61, 0x6d, 0x65, 0x18, 0x04, 0x20, 0x01, 0x28, 0x09, 0x52, 0x04, 0x6e, 0x61, 0x6d, - 0x65, 0x12, 0x12, 0x0a, 0x04, 0x65, 0x74, 0x61, 0x67, 0x18, 0x05, 0x20, 0x01, 0x28, 0x09, 0x52, - 0x04, 0x65, 0x74, 0x61, 0x67, 0x12, 0x1e, 0x0a, 0x0a, 0x69, 0x6e, 0x73, 0x74, 0x61, 0x6e, 0x63, - 0x65, 0x49, 0x64, 0x18, 0x06, 0x20, 0x01, 0x28, 0x09, 0x52, 0x0a, 0x69, 0x6e, 0x73, 0x74, 0x61, - 0x6e, 0x63, 0x65, 0x49, 0x64, 0x12, 0x1a, 0x0a, 0x08, 0x73, 0x74, 0x61, 0x72, 0x74, 0x50, 0x6f, - 0x73, 0x18, 0x07, 0x20, 0x01, 0x28, 0x03, 0x52, 0x08, 0x73, 0x74, 0x61, 0x72, 0x74, 0x50, 0x6f, - 0x73, 0x12, 0x16, 0x0a, 0x06, 0x65, 0x6e, 0x64, 0x50, 0x6f, 0x73, 0x18, 0x08, 0x20, 0x01, 0x28, - 0x03, 0x52, 0x06, 0x65, 0x6e, 0x64, 0x50, 0x6f, 0x73, 0x12, 0x1a, 0x0a, 0x08, 0x66, 0x69, 0x6c, - 0x65, 0x53, 0x69, 0x7a, 0x65, 0x18, 0x09, 0x20, 0x01, 0x28, 0x03, 0x52, 0x08, 0x66, 0x69, 0x6c, - 0x65, 0x53, 0x69, 0x7a, 0x65, 0x22, 0x5f, 0x0a, 0x12, 0x53, 0x79, 0x6e, 0x63, 0x46, 0x69, 0x6c, - 0x65, 0x50, 0x72, 0x6f, 0x63, 0x65, 0x73, 0x73, 0x52, 0x65, 0x71, 0x12, 0x49, 0x0a, 0x12, 0x66, - 0x69, 0x6c, 0x65, 0x50, 0x72, 0x6f, 0x63, 0x65, 0x73, 0x73, 0x45, 0x6e, 0x74, 0x72, 0x69, 0x65, - 0x73, 0x18, 0x01, 0x20, 0x03, 0x28, 0x0b, 0x32, 0x19, 0x2e, 0x6d, 0x61, 0x6e, 0x61, 0x67, 0x65, - 0x72, 0x2e, 0x46, 0x69, 0x6c, 0x65, 0x50, 0x72, 0x6f, 0x63, 0x65, 0x73, 0x73, 0x45, 0x6e, 0x74, - 0x72, 0x79, 0x52, 0x12, 0x66, 0x69, 0x6c, 0x65, 0x50, 0x72, 0x6f, 0x63, 0x65, 0x73, 0x73, 0x45, - 0x6e, 0x74, 0x72, 0x69, 0x65, 0x73, 0x22, 0xa2, 0x02, 0x0a, 0x10, 0x46, 0x69, 0x6c, 0x65, 0x50, - 0x72, 0x6f, 0x63, 0x65, 0x73, 0x73, 0x45, 0x6e, 0x74, 0x72, 0x79, 0x12, 0x1a, 0x0a, 0x08, 0x64, - 0x61, 0x74, 0x61, 0x54, 0x79, 0x70, 0x65, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x08, 0x64, - 0x61, 0x74, 0x61, 0x54, 0x79, 0x70, 0x65, 0x12, 0x10, 0x0a, 0x03, 0x6f, 0x72, 0x67, 0x18, 0x02, - 0x20, 0x01, 0x28, 0x09, 0x52, 0x03, 0x6f, 0x72, 0x67, 0x12, 0x12, 0x0a, 0x04, 0x72, 0x65, 0x70, - 0x6f, 0x18, 0x03, 0x20, 0x01, 0x28, 0x09, 0x52, 0x04, 0x72, 0x65, 0x70, 0x6f, 0x12, 0x12, 0x0a, - 0x04, 0x6e, 0x61, 0x6d, 0x65, 0x18, 0x04, 0x20, 0x01, 0x28, 0x09, 0x52, 0x04, 0x6e, 0x61, 0x6d, - 0x65, 0x12, 0x12, 0x0a, 0x04, 0x65, 0x74, 0x61, 0x67, 0x18, 0x05, 0x20, 0x01, 0x28, 0x09, 0x52, - 0x04, 0x65, 0x74, 0x61, 0x67, 0x12, 0x1e, 0x0a, 0x0a, 0x69, 0x6e, 0x73, 0x74, 0x61, 0x6e, 0x63, - 0x65, 0x49, 0x64, 0x18, 0x06, 0x20, 0x01, 0x28, 0x09, 0x52, 0x0a, 0x69, 0x6e, 0x73, 0x74, 0x61, - 0x6e, 0x63, 0x65, 0x49, 0x64, 0x12, 0x1a, 0x0a, 0x08, 0x73, 0x74, 0x61, 0x72, 0x74, 0x50, 0x6f, - 0x73, 0x18, 0x07, 0x20, 0x01, 0x28, 0x03, 0x52, 0x08, 0x73, 0x74, 0x61, 0x72, 0x74, 0x50, 0x6f, - 0x73, 0x12, 0x16, 0x0a, 0x06, 0x65, 0x6e, 0x64, 0x50, 0x6f, 0x73, 0x18, 0x08, 0x20, 0x01, 0x28, - 0x03, 0x52, 0x06, 0x65, 0x6e, 0x64, 0x50, 0x6f, 0x73, 0x12, 0x1a, 0x0a, 0x08, 0x66, 0x69, 0x6c, - 0x65, 0x53, 0x69, 0x7a, 0x65, 0x18, 0x09, 0x20, 0x01, 0x28, 0x03, 0x52, 0x08, 0x66, 0x69, 0x6c, - 0x65, 0x53, 0x69, 0x7a, 0x65, 0x12, 0x16, 0x0a, 0x06, 0x73, 0x74, 0x61, 0x74, 0x75, 0x73, 0x18, - 0x0a, 0x20, 0x01, 0x28, 0x05, 0x52, 0x06, 0x73, 0x74, 0x61, 0x74, 0x75, 0x73, 0x12, 0x1c, 0x0a, - 0x09, 0x70, 0x72, 0x6f, 0x63, 0x65, 0x73, 0x73, 0x49, 0x64, 0x18, 0x0b, 0x20, 0x01, 0x28, 0x03, - 0x52, 0x09, 0x70, 0x72, 0x6f, 0x63, 0x65, 0x73, 0x73, 0x49, 0x64, 0x22, 0xcd, 0x01, 0x0a, 0x15, - 0x53, 0x63, 0x68, 0x65, 0x64, 0x75, 0x6c, 0x65, 0x72, 0x46, 0x69, 0x6c, 0x65, 0x52, 0x65, 0x73, - 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x24, 0x0a, 0x0d, 0x73, 0x63, 0x68, 0x65, 0x64, 0x75, 0x6c, - 0x65, 0x72, 0x54, 0x79, 0x70, 0x65, 0x18, 0x01, 0x20, 0x01, 0x28, 0x05, 0x52, 0x0d, 0x73, 0x63, - 0x68, 0x65, 0x64, 0x75, 0x6c, 0x65, 0x72, 0x54, 0x79, 0x70, 0x65, 0x12, 0x1c, 0x0a, 0x09, 0x70, - 0x72, 0x6f, 0x63, 0x65, 0x73, 0x73, 0x49, 0x64, 0x18, 0x02, 0x20, 0x01, 0x28, 0x03, 0x52, 0x09, - 0x70, 0x72, 0x6f, 0x63, 0x65, 0x73, 0x73, 0x49, 0x64, 0x12, 0x12, 0x0a, 0x04, 0x68, 0x6f, 0x73, - 0x74, 0x18, 0x03, 0x20, 0x01, 0x28, 0x09, 0x52, 0x04, 0x68, 0x6f, 0x73, 0x74, 0x12, 0x12, 0x0a, - 0x04, 0x70, 0x6f, 0x72, 0x74, 0x18, 0x04, 0x20, 0x01, 0x28, 0x05, 0x52, 0x04, 0x70, 0x6f, 0x72, - 0x74, 0x12, 0x2a, 0x0a, 0x10, 0x6d, 0x61, 0x73, 0x74, 0x65, 0x72, 0x49, 0x6e, 0x73, 0x74, 0x61, - 0x6e, 0x63, 0x65, 0x49, 0x64, 0x18, 0x05, 0x20, 0x01, 0x28, 0x09, 0x52, 0x10, 0x6d, 0x61, 0x73, - 0x74, 0x65, 0x72, 0x49, 0x6e, 0x73, 0x74, 0x61, 0x6e, 0x63, 0x65, 0x49, 0x64, 0x12, 0x1c, 0x0a, - 0x09, 0x6d, 0x61, 0x78, 0x4f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x18, 0x06, 0x20, 0x01, 0x28, 0x03, - 0x52, 0x09, 0x6d, 0x61, 0x78, 0x4f, 0x66, 0x66, 0x73, 0x65, 0x74, 0x22, 0x7a, 0x0a, 0x12, 0x46, - 0x69, 0x6c, 0x65, 0x50, 0x72, 0x6f, 0x63, 0x65, 0x73, 0x73, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, - 0x74, 0x12, 0x1c, 0x0a, 0x09, 0x70, 0x72, 0x6f, 0x63, 0x65, 0x73, 0x73, 0x49, 0x64, 0x18, 0x01, - 0x20, 0x01, 0x28, 0x03, 0x52, 0x09, 0x70, 0x72, 0x6f, 0x63, 0x65, 0x73, 0x73, 0x49, 0x64, 0x12, - 0x16, 0x0a, 0x06, 0x73, 0x74, 0x61, 0x50, 0x6f, 0x73, 0x18, 0x02, 0x20, 0x01, 0x28, 0x03, 0x52, - 0x06, 0x73, 0x74, 0x61, 0x50, 0x6f, 0x73, 0x12, 0x16, 0x0a, 0x06, 0x65, 0x6e, 0x64, 0x50, 0x6f, - 0x73, 0x18, 0x03, 0x20, 0x01, 0x28, 0x03, 0x52, 0x06, 0x65, 0x6e, 0x64, 0x50, 0x6f, 0x73, 0x12, - 0x16, 0x0a, 0x06, 0x73, 0x74, 0x61, 0x74, 0x75, 0x73, 0x18, 0x04, 0x20, 0x01, 0x28, 0x05, 0x52, - 0x06, 0x73, 0x74, 0x61, 0x74, 0x75, 0x73, 0x22, 0xa9, 0x01, 0x0a, 0x1d, 0x44, 0x65, 0x6c, 0x65, - 0x74, 0x65, 0x42, 0x79, 0x45, 0x74, 0x61, 0x67, 0x73, 0x41, 0x6e, 0x64, 0x46, 0x69, 0x65, 0x6c, - 0x64, 0x73, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x12, 0x12, 0x0a, 0x04, 0x65, 0x74, 0x61, - 0x67, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x04, 0x65, 0x74, 0x61, 0x67, 0x12, 0x1a, 0x0a, - 0x08, 0x64, 0x61, 0x74, 0x61, 0x74, 0x79, 0x70, 0x65, 0x18, 0x02, 0x20, 0x01, 0x28, 0x09, 0x52, - 0x08, 0x64, 0x61, 0x74, 0x61, 0x74, 0x79, 0x70, 0x65, 0x12, 0x10, 0x0a, 0x03, 0x6f, 0x72, 0x67, - 0x18, 0x03, 0x20, 0x01, 0x28, 0x09, 0x52, 0x03, 0x6f, 0x72, 0x67, 0x12, 0x12, 0x0a, 0x04, 0x72, - 0x65, 0x70, 0x6f, 0x18, 0x04, 0x20, 0x01, 0x28, 0x09, 0x52, 0x04, 0x72, 0x65, 0x70, 0x6f, 0x12, - 0x12, 0x0a, 0x04, 0x6e, 0x61, 0x6d, 0x65, 0x18, 0x05, 0x20, 0x01, 0x28, 0x09, 0x52, 0x04, 0x6e, - 0x61, 0x6d, 0x65, 0x12, 0x1e, 0x0a, 0x0a, 0x69, 0x6e, 0x73, 0x74, 0x61, 0x6e, 0x63, 0x65, 0x49, - 0x44, 0x18, 0x06, 0x20, 0x01, 0x28, 0x09, 0x52, 0x0a, 0x69, 0x6e, 0x73, 0x74, 0x61, 0x6e, 0x63, - 0x65, 0x49, 0x44, 0x22, 0xdb, 0x01, 0x0a, 0x11, 0x43, 0x72, 0x65, 0x61, 0x74, 0x65, 0x43, 0x61, - 0x63, 0x68, 0x65, 0x4a, 0x6f, 0x62, 0x52, 0x65, 0x71, 0x12, 0x12, 0x0a, 0x04, 0x74, 0x79, 0x70, - 0x65, 0x18, 0x01, 0x20, 0x01, 0x28, 0x05, 0x52, 0x04, 0x74, 0x79, 0x70, 0x65, 0x12, 0x1e, 0x0a, - 0x0a, 0x69, 0x6e, 0x73, 0x74, 0x61, 0x6e, 0x63, 0x65, 0x49, 0x64, 0x18, 0x02, 0x20, 0x01, 0x28, - 0x09, 0x52, 0x0a, 0x69, 0x6e, 0x73, 0x74, 0x61, 0x6e, 0x63, 0x65, 0x49, 0x64, 0x12, 0x1a, 0x0a, - 0x08, 0x64, 0x61, 0x74, 0x61, 0x74, 0x79, 0x70, 0x65, 0x18, 0x03, 0x20, 0x01, 0x28, 0x09, 0x52, - 0x08, 0x64, 0x61, 0x74, 0x61, 0x74, 0x79, 0x70, 0x65, 0x12, 0x10, 0x0a, 0x03, 0x6f, 0x72, 0x67, - 0x18, 0x04, 0x20, 0x01, 0x28, 0x09, 0x52, 0x03, 0x6f, 0x72, 0x67, 0x12, 0x12, 0x0a, 0x04, 0x72, - 0x65, 0x70, 0x6f, 0x18, 0x05, 0x20, 0x01, 0x28, 0x09, 0x52, 0x04, 0x72, 0x65, 0x70, 0x6f, 0x12, - 0x20, 0x0a, 0x0b, 0x75, 0x73, 0x65, 0x64, 0x53, 0x74, 0x6f, 0x72, 0x61, 0x67, 0x65, 0x18, 0x06, - 0x20, 0x01, 0x28, 0x03, 0x52, 0x0b, 0x75, 0x73, 0x65, 0x64, 0x53, 0x74, 0x6f, 0x72, 0x61, 0x67, - 0x65, 0x12, 0x16, 0x0a, 0x06, 0x63, 0x6f, 0x6d, 0x6d, 0x69, 0x74, 0x18, 0x07, 0x20, 0x01, 0x28, - 0x09, 0x52, 0x06, 0x63, 0x6f, 0x6d, 0x6d, 0x69, 0x74, 0x12, 0x16, 0x0a, 0x06, 0x73, 0x74, 0x61, - 0x74, 0x75, 0x73, 0x18, 0x08, 0x20, 0x01, 0x28, 0x05, 0x52, 0x06, 0x73, 0x74, 0x61, 0x74, 0x75, - 0x73, 0x22, 0x24, 0x0a, 0x12, 0x43, 0x72, 0x65, 0x61, 0x74, 0x65, 0x43, 0x61, 0x63, 0x68, 0x65, - 0x4a, 0x6f, 0x62, 0x52, 0x65, 0x73, 0x70, 0x12, 0x0e, 0x0a, 0x02, 0x69, 0x64, 0x18, 0x01, 0x20, - 0x01, 0x28, 0x03, 0x52, 0x02, 0x69, 0x64, 0x22, 0xbd, 0x01, 0x0a, 0x17, 0x55, 0x70, 0x64, 0x61, - 0x74, 0x65, 0x43, 0x61, 0x63, 0x68, 0x65, 0x4a, 0x6f, 0x62, 0x53, 0x74, 0x61, 0x74, 0x75, 0x73, - 0x52, 0x65, 0x71, 0x12, 0x0e, 0x0a, 0x02, 0x69, 0x64, 0x18, 0x01, 0x20, 0x01, 0x28, 0x03, 0x52, - 0x02, 0x69, 0x64, 0x12, 0x16, 0x0a, 0x06, 0x73, 0x74, 0x61, 0x74, 0x75, 0x73, 0x18, 0x02, 0x20, - 0x01, 0x28, 0x05, 0x52, 0x06, 0x73, 0x74, 0x61, 0x74, 0x75, 0x73, 0x12, 0x1a, 0x0a, 0x08, 0x65, - 0x72, 0x72, 0x6f, 0x72, 0x4d, 0x73, 0x67, 0x18, 0x03, 0x20, 0x01, 0x28, 0x09, 0x52, 0x08, 0x65, - 0x72, 0x72, 0x6f, 0x72, 0x4d, 0x73, 0x67, 0x12, 0x1e, 0x0a, 0x0a, 0x69, 0x6e, 0x73, 0x74, 0x61, - 0x6e, 0x63, 0x65, 0x49, 0x64, 0x18, 0x04, 0x20, 0x01, 0x28, 0x09, 0x52, 0x0a, 0x69, 0x6e, 0x73, - 0x74, 0x61, 0x6e, 0x63, 0x65, 0x49, 0x64, 0x12, 0x10, 0x0a, 0x03, 0x6f, 0x72, 0x67, 0x18, 0x05, - 0x20, 0x01, 0x28, 0x09, 0x52, 0x03, 0x6f, 0x72, 0x67, 0x12, 0x12, 0x0a, 0x04, 0x72, 0x65, 0x70, - 0x6f, 0x18, 0x06, 0x20, 0x01, 0x28, 0x09, 0x52, 0x04, 0x72, 0x65, 0x70, 0x6f, 0x12, 0x18, 0x0a, - 0x07, 0x70, 0x72, 0x6f, 0x63, 0x65, 0x73, 0x73, 0x18, 0x07, 0x20, 0x01, 0x28, 0x02, 0x52, 0x07, - 0x70, 0x72, 0x6f, 0x63, 0x65, 0x73, 0x73, 0x22, 0x64, 0x0a, 0x1e, 0x55, 0x70, 0x64, 0x61, 0x74, - 0x65, 0x52, 0x65, 0x70, 0x6f, 0x73, 0x69, 0x74, 0x6f, 0x72, 0x79, 0x4d, 0x6f, 0x75, 0x6e, 0x74, - 0x53, 0x74, 0x61, 0x74, 0x75, 0x73, 0x52, 0x65, 0x71, 0x12, 0x0e, 0x0a, 0x02, 0x69, 0x64, 0x18, - 0x01, 0x20, 0x01, 0x28, 0x03, 0x52, 0x02, 0x69, 0x64, 0x12, 0x16, 0x0a, 0x06, 0x73, 0x74, 0x61, - 0x74, 0x75, 0x73, 0x18, 0x02, 0x20, 0x01, 0x28, 0x05, 0x52, 0x06, 0x73, 0x74, 0x61, 0x74, 0x75, - 0x73, 0x12, 0x1a, 0x0a, 0x08, 0x65, 0x72, 0x72, 0x6f, 0x72, 0x4d, 0x73, 0x67, 0x18, 0x03, 0x20, - 0x01, 0x28, 0x09, 0x52, 0x08, 0x65, 0x72, 0x72, 0x6f, 0x72, 0x4d, 0x73, 0x67, 0x32, 0xc5, 0x05, - 0x0a, 0x07, 0x4d, 0x61, 0x6e, 0x61, 0x67, 0x65, 0x72, 0x12, 0x41, 0x0a, 0x08, 0x52, 0x65, 0x67, - 0x69, 0x73, 0x74, 0x65, 0x72, 0x12, 0x18, 0x2e, 0x6d, 0x61, 0x6e, 0x61, 0x67, 0x65, 0x72, 0x2e, - 0x52, 0x65, 0x67, 0x69, 0x73, 0x74, 0x65, 0x72, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, - 0x19, 0x2e, 0x6d, 0x61, 0x6e, 0x61, 0x67, 0x65, 0x72, 0x2e, 0x52, 0x65, 0x67, 0x69, 0x73, 0x74, - 0x65, 0x72, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x12, 0x3e, 0x0a, 0x09, - 0x48, 0x65, 0x61, 0x72, 0x74, 0x62, 0x65, 0x61, 0x74, 0x12, 0x19, 0x2e, 0x6d, 0x61, 0x6e, 0x61, - 0x67, 0x65, 0x72, 0x2e, 0x48, 0x65, 0x61, 0x72, 0x74, 0x62, 0x65, 0x61, 0x74, 0x52, 0x65, 0x71, - 0x75, 0x65, 0x73, 0x74, 0x1a, 0x16, 0x2e, 0x67, 0x6f, 0x6f, 0x67, 0x6c, 0x65, 0x2e, 0x70, 0x72, - 0x6f, 0x74, 0x6f, 0x62, 0x75, 0x66, 0x2e, 0x45, 0x6d, 0x70, 0x74, 0x79, 0x12, 0x4e, 0x0a, 0x0d, - 0x53, 0x63, 0x68, 0x65, 0x64, 0x75, 0x6c, 0x65, 0x72, 0x46, 0x69, 0x6c, 0x65, 0x12, 0x1d, 0x2e, - 0x6d, 0x61, 0x6e, 0x61, 0x67, 0x65, 0x72, 0x2e, 0x53, 0x63, 0x68, 0x65, 0x64, 0x75, 0x6c, 0x65, - 0x72, 0x46, 0x69, 0x6c, 0x65, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x1e, 0x2e, 0x6d, - 0x61, 0x6e, 0x61, 0x67, 0x65, 0x72, 0x2e, 0x53, 0x63, 0x68, 0x65, 0x64, 0x75, 0x6c, 0x65, 0x72, - 0x46, 0x69, 0x6c, 0x65, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x48, 0x0a, 0x11, - 0x52, 0x65, 0x70, 0x6f, 0x72, 0x74, 0x46, 0x69, 0x6c, 0x65, 0x50, 0x72, 0x6f, 0x63, 0x65, 0x73, - 0x73, 0x12, 0x1b, 0x2e, 0x6d, 0x61, 0x6e, 0x61, 0x67, 0x65, 0x72, 0x2e, 0x46, 0x69, 0x6c, 0x65, - 0x50, 0x72, 0x6f, 0x63, 0x65, 0x73, 0x73, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, 0x1a, 0x16, - 0x2e, 0x67, 0x6f, 0x6f, 0x67, 0x6c, 0x65, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x62, 0x75, 0x66, - 0x2e, 0x45, 0x6d, 0x70, 0x74, 0x79, 0x12, 0x46, 0x0a, 0x0f, 0x53, 0x79, 0x6e, 0x63, 0x46, 0x69, - 0x6c, 0x65, 0x50, 0x72, 0x6f, 0x63, 0x65, 0x73, 0x73, 0x12, 0x1b, 0x2e, 0x6d, 0x61, 0x6e, 0x61, - 0x67, 0x65, 0x72, 0x2e, 0x53, 0x79, 0x6e, 0x63, 0x46, 0x69, 0x6c, 0x65, 0x50, 0x72, 0x6f, 0x63, - 0x65, 0x73, 0x73, 0x52, 0x65, 0x71, 0x1a, 0x16, 0x2e, 0x67, 0x6f, 0x6f, 0x67, 0x6c, 0x65, 0x2e, - 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x62, 0x75, 0x66, 0x2e, 0x45, 0x6d, 0x70, 0x74, 0x79, 0x12, 0x58, - 0x0a, 0x16, 0x44, 0x65, 0x6c, 0x65, 0x74, 0x65, 0x42, 0x79, 0x45, 0x74, 0x61, 0x67, 0x73, 0x41, - 0x6e, 0x64, 0x46, 0x69, 0x65, 0x6c, 0x64, 0x73, 0x12, 0x26, 0x2e, 0x6d, 0x61, 0x6e, 0x61, 0x67, - 0x65, 0x72, 0x2e, 0x44, 0x65, 0x6c, 0x65, 0x74, 0x65, 0x42, 0x79, 0x45, 0x74, 0x61, 0x67, 0x73, - 0x41, 0x6e, 0x64, 0x46, 0x69, 0x65, 0x6c, 0x64, 0x73, 0x52, 0x65, 0x71, 0x75, 0x65, 0x73, 0x74, - 0x1a, 0x16, 0x2e, 0x67, 0x6f, 0x6f, 0x67, 0x6c, 0x65, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x62, - 0x75, 0x66, 0x2e, 0x45, 0x6d, 0x70, 0x74, 0x79, 0x12, 0x49, 0x0a, 0x0e, 0x43, 0x72, 0x65, 0x61, - 0x74, 0x65, 0x43, 0x61, 0x63, 0x68, 0x65, 0x4a, 0x6f, 0x62, 0x12, 0x1a, 0x2e, 0x6d, 0x61, 0x6e, - 0x61, 0x67, 0x65, 0x72, 0x2e, 0x43, 0x72, 0x65, 0x61, 0x74, 0x65, 0x43, 0x61, 0x63, 0x68, 0x65, - 0x4a, 0x6f, 0x62, 0x52, 0x65, 0x71, 0x1a, 0x1b, 0x2e, 0x6d, 0x61, 0x6e, 0x61, 0x67, 0x65, 0x72, - 0x2e, 0x43, 0x72, 0x65, 0x61, 0x74, 0x65, 0x43, 0x61, 0x63, 0x68, 0x65, 0x4a, 0x6f, 0x62, 0x52, - 0x65, 0x73, 0x70, 0x12, 0x50, 0x0a, 0x14, 0x55, 0x70, 0x64, 0x61, 0x74, 0x65, 0x43, 0x61, 0x63, - 0x68, 0x65, 0x4a, 0x6f, 0x62, 0x53, 0x74, 0x61, 0x74, 0x75, 0x73, 0x12, 0x20, 0x2e, 0x6d, 0x61, - 0x6e, 0x61, 0x67, 0x65, 0x72, 0x2e, 0x55, 0x70, 0x64, 0x61, 0x74, 0x65, 0x43, 0x61, 0x63, 0x68, - 0x65, 0x4a, 0x6f, 0x62, 0x53, 0x74, 0x61, 0x74, 0x75, 0x73, 0x52, 0x65, 0x71, 0x1a, 0x16, 0x2e, - 0x67, 0x6f, 0x6f, 0x67, 0x6c, 0x65, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x62, 0x75, 0x66, 0x2e, - 0x45, 0x6d, 0x70, 0x74, 0x79, 0x12, 0x5e, 0x0a, 0x1b, 0x55, 0x70, 0x64, 0x61, 0x74, 0x65, 0x52, - 0x65, 0x70, 0x6f, 0x73, 0x69, 0x74, 0x6f, 0x72, 0x79, 0x4d, 0x6f, 0x75, 0x6e, 0x74, 0x53, 0x74, - 0x61, 0x74, 0x75, 0x73, 0x12, 0x27, 0x2e, 0x6d, 0x61, 0x6e, 0x61, 0x67, 0x65, 0x72, 0x2e, 0x55, - 0x70, 0x64, 0x61, 0x74, 0x65, 0x52, 0x65, 0x70, 0x6f, 0x73, 0x69, 0x74, 0x6f, 0x72, 0x79, 0x4d, - 0x6f, 0x75, 0x6e, 0x74, 0x53, 0x74, 0x61, 0x74, 0x75, 0x73, 0x52, 0x65, 0x71, 0x1a, 0x16, 0x2e, - 0x67, 0x6f, 0x6f, 0x67, 0x6c, 0x65, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x62, 0x75, 0x66, 0x2e, - 0x45, 0x6d, 0x70, 0x74, 0x79, 0x42, 0x0b, 0x5a, 0x09, 0x2e, 0x3b, 0x6d, 0x61, 0x6e, 0x61, 0x67, - 0x65, 0x72, 0x62, 0x06, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x33, -}) +const file_manager_proto_rawDesc = "" + + "\n" + + "\rmanager.proto\x12\amanager\x1a\x1bgoogle/protobuf/empty.proto\"\xb9\x01\n" + + "\x0fRegisterRequest\x12\x1e\n" + + "\n" + + "instanceId\x18\x01 \x01(\tR\n" + + "instanceId\x12\x12\n" + + "\x04host\x18\x02 \x01(\tR\x04host\x12\x12\n" + + "\x04port\x18\x03 \x01(\x05R\x04port\x12\x16\n" + + "\x06online\x18\x04 \x01(\bR\x06online\x12$\n" + + "\rmanagementUrl\x18\x05 \x01(\tR\rmanagementUrl\x12 \n" + + "\vdownloadUrl\x18\x06 \x01(\tR\vdownloadUrl\"<\n" + + "\x10RegisterResponse\x12\x0e\n" + + "\x02id\x18\x01 \x01(\x05R\x02id\x12\x18\n" + + "\asuccess\x18\x02 \x01(\bR\asuccess\"\x8f\x01\n" + + "\x10HeartbeatRequest\x12\x0e\n" + + "\x02id\x18\x01 \x01(\x05R\x02id\x12\x1e\n" + + "\n" + + "instanceId\x18\x02 \x01(\tR\n" + + "instanceId\x12\x16\n" + + "\x06online\x18\x03 \x01(\bR\x06online\x123\n" + + "\x06health\x18\x04 \x01(\v2\x1b.manager.NodeHealthSnapshotR\x06health\"\xc1\x02\n" + + "\x12NodeHealthSnapshot\x12\x18\n" + + "\aversion\x18\x01 \x01(\rR\aversion\x12\x1d\n" + + "\n" + + "process_id\x18\x02 \x01(\tR\tprocessId\x12\x1d\n" + + "\n" + + "started_at\x18\x03 \x01(\x03R\tstartedAt\x12!\n" + + "\fcollected_at\x18\x04 \x01(\x03R\vcollectedAt\x128\n" + + "\x18heartbeat_period_seconds\x18\x05 \x01(\rR\x16heartbeatPeriodSeconds\x12B\n" + + "\fcapabilities\x18\x06 \x03(\v2\x1e.manager.CapabilityObservationR\fcapabilities\x122\n" + + "\x06errors\x18\a \x03(\v2\x1a.manager.StorageErrorCountR\x06errors\"\xbb\x01\n" + + "\x15CapabilityObservation\x12\x1e\n" + + "\n" + + "capability\x18\x01 \x01(\tR\n" + + "capability\x12\x14\n" + + "\x05state\x18\x02 \x01(\x05R\x05state\x12\x1e\n" + + "\n" + + "unresolved\x18\x03 \x01(\bR\n" + + "unresolved\x12)\n" + + "\x10last_observation\x18\x04 \x01(\x03R\x0flastObservation\x12!\n" + + "\flast_failure\x18\x05 \x01(\x03R\vlastFailure\"[\n" + + "\x11StorageErrorCount\x12\x1c\n" + + "\toperation\x18\x01 \x01(\tR\toperation\x12\x12\n" + + "\x04kind\x18\x02 \x01(\tR\x04kind\x12\x14\n" + + "\x05count\x18\x03 \x01(\x04R\x05count\"\xf0\x01\n" + + "\x14SchedulerFileRequest\x12\x1a\n" + + "\bdataType\x18\x01 \x01(\tR\bdataType\x12\x10\n" + + "\x03org\x18\x02 \x01(\tR\x03org\x12\x12\n" + + "\x04repo\x18\x03 \x01(\tR\x04repo\x12\x12\n" + + "\x04name\x18\x04 \x01(\tR\x04name\x12\x12\n" + + "\x04etag\x18\x05 \x01(\tR\x04etag\x12\x1e\n" + + "\n" + + "instanceId\x18\x06 \x01(\tR\n" + + "instanceId\x12\x1a\n" + + "\bstartPos\x18\a \x01(\x03R\bstartPos\x12\x16\n" + + "\x06endPos\x18\b \x01(\x03R\x06endPos\x12\x1a\n" + + "\bfileSize\x18\t \x01(\x03R\bfileSize\"_\n" + + "\x12SyncFileProcessReq\x12I\n" + + "\x12fileProcessEntries\x18\x01 \x03(\v2\x19.manager.FileProcessEntryR\x12fileProcessEntries\"\xa2\x02\n" + + "\x10FileProcessEntry\x12\x1a\n" + + "\bdataType\x18\x01 \x01(\tR\bdataType\x12\x10\n" + + "\x03org\x18\x02 \x01(\tR\x03org\x12\x12\n" + + "\x04repo\x18\x03 \x01(\tR\x04repo\x12\x12\n" + + "\x04name\x18\x04 \x01(\tR\x04name\x12\x12\n" + + "\x04etag\x18\x05 \x01(\tR\x04etag\x12\x1e\n" + + "\n" + + "instanceId\x18\x06 \x01(\tR\n" + + "instanceId\x12\x1a\n" + + "\bstartPos\x18\a \x01(\x03R\bstartPos\x12\x16\n" + + "\x06endPos\x18\b \x01(\x03R\x06endPos\x12\x1a\n" + + "\bfileSize\x18\t \x01(\x03R\bfileSize\x12\x16\n" + + "\x06status\x18\n" + + " \x01(\x05R\x06status\x12\x1c\n" + + "\tprocessId\x18\v \x01(\x03R\tprocessId\"\xcd\x01\n" + + "\x15SchedulerFileResponse\x12$\n" + + "\rschedulerType\x18\x01 \x01(\x05R\rschedulerType\x12\x1c\n" + + "\tprocessId\x18\x02 \x01(\x03R\tprocessId\x12\x12\n" + + "\x04host\x18\x03 \x01(\tR\x04host\x12\x12\n" + + "\x04port\x18\x04 \x01(\x05R\x04port\x12*\n" + + "\x10masterInstanceId\x18\x05 \x01(\tR\x10masterInstanceId\x12\x1c\n" + + "\tmaxOffset\x18\x06 \x01(\x03R\tmaxOffset\"z\n" + + "\x12FileProcessRequest\x12\x1c\n" + + "\tprocessId\x18\x01 \x01(\x03R\tprocessId\x12\x16\n" + + "\x06staPos\x18\x02 \x01(\x03R\x06staPos\x12\x16\n" + + "\x06endPos\x18\x03 \x01(\x03R\x06endPos\x12\x16\n" + + "\x06status\x18\x04 \x01(\x05R\x06status\"\xa9\x01\n" + + "\x1dDeleteByEtagsAndFieldsRequest\x12\x12\n" + + "\x04etag\x18\x01 \x01(\tR\x04etag\x12\x1a\n" + + "\bdatatype\x18\x02 \x01(\tR\bdatatype\x12\x10\n" + + "\x03org\x18\x03 \x01(\tR\x03org\x12\x12\n" + + "\x04repo\x18\x04 \x01(\tR\x04repo\x12\x12\n" + + "\x04name\x18\x05 \x01(\tR\x04name\x12\x1e\n" + + "\n" + + "instanceID\x18\x06 \x01(\tR\n" + + "instanceID\"\xdb\x01\n" + + "\x11CreateCacheJobReq\x12\x12\n" + + "\x04type\x18\x01 \x01(\x05R\x04type\x12\x1e\n" + + "\n" + + "instanceId\x18\x02 \x01(\tR\n" + + "instanceId\x12\x1a\n" + + "\bdatatype\x18\x03 \x01(\tR\bdatatype\x12\x10\n" + + "\x03org\x18\x04 \x01(\tR\x03org\x12\x12\n" + + "\x04repo\x18\x05 \x01(\tR\x04repo\x12 \n" + + "\vusedStorage\x18\x06 \x01(\x03R\vusedStorage\x12\x16\n" + + "\x06commit\x18\a \x01(\tR\x06commit\x12\x16\n" + + "\x06status\x18\b \x01(\x05R\x06status\"$\n" + + "\x12CreateCacheJobResp\x12\x0e\n" + + "\x02id\x18\x01 \x01(\x03R\x02id\"\xbd\x01\n" + + "\x17UpdateCacheJobStatusReq\x12\x0e\n" + + "\x02id\x18\x01 \x01(\x03R\x02id\x12\x16\n" + + "\x06status\x18\x02 \x01(\x05R\x06status\x12\x1a\n" + + "\berrorMsg\x18\x03 \x01(\tR\berrorMsg\x12\x1e\n" + + "\n" + + "instanceId\x18\x04 \x01(\tR\n" + + "instanceId\x12\x10\n" + + "\x03org\x18\x05 \x01(\tR\x03org\x12\x12\n" + + "\x04repo\x18\x06 \x01(\tR\x04repo\x12\x18\n" + + "\aprocess\x18\a \x01(\x02R\aprocess\"d\n" + + "\x1eUpdateRepositoryMountStatusReq\x12\x0e\n" + + "\x02id\x18\x01 \x01(\x03R\x02id\x12\x16\n" + + "\x06status\x18\x02 \x01(\x05R\x06status\x12\x1a\n" + + "\berrorMsg\x18\x03 \x01(\tR\berrorMsg\"\xd5\x01\n" + + "\x17IngestRepositoryRequest\x12\x1f\n" + + "\vinstance_id\x18\x01 \x01(\tR\n" + + "instanceId\x12\x1c\n" + + "\tnamespace\x18\x02 \x01(\tR\tnamespace\x12\x1b\n" + + "\trepo_type\x18\x03 \x01(\tR\brepoType\x12\x12\n" + + "\x04repo\x18\x04 \x01(\tR\x04repo\x12\x1a\n" + + "\brevision\x18\x05 \x01(\tR\brevision\x12\x16\n" + + "\x06commit\x18\x06 \x01(\tR\x06commit\x12\x16\n" + + "\x06online\x18\a \x01(\bR\x06online\"\x99\x01\n" + + "\x18IngestRepositoryResponse\x12#\n" + + "\rrepository_id\x18\x01 \x01(\x03R\frepositoryId\x12\x16\n" + + "\x06commit\x18\x02 \x01(\tR\x06commit\x12\x1d\n" + + "\n" + + "file_count\x18\x03 \x01(\x03R\tfileCount\x12!\n" + + "\fused_storage\x18\x04 \x01(\x03R\vusedStorage2\x9e\x06\n" + + "\aManager\x12W\n" + + "\x10IngestRepository\x12 .manager.IngestRepositoryRequest\x1a!.manager.IngestRepositoryResponse\x12A\n" + + "\bRegister\x12\x18.manager.RegisterRequest\x1a\x19.manager.RegisterResponse\"\x00\x12>\n" + + "\tHeartbeat\x12\x19.manager.HeartbeatRequest\x1a\x16.google.protobuf.Empty\x12N\n" + + "\rSchedulerFile\x12\x1d.manager.SchedulerFileRequest\x1a\x1e.manager.SchedulerFileResponse\x12H\n" + + "\x11ReportFileProcess\x12\x1b.manager.FileProcessRequest\x1a\x16.google.protobuf.Empty\x12F\n" + + "\x0fSyncFileProcess\x12\x1b.manager.SyncFileProcessReq\x1a\x16.google.protobuf.Empty\x12X\n" + + "\x16DeleteByEtagsAndFields\x12&.manager.DeleteByEtagsAndFieldsRequest\x1a\x16.google.protobuf.Empty\x12I\n" + + "\x0eCreateCacheJob\x12\x1a.manager.CreateCacheJobReq\x1a\x1b.manager.CreateCacheJobResp\x12P\n" + + "\x14UpdateCacheJobStatus\x12 .manager.UpdateCacheJobStatusReq\x1a\x16.google.protobuf.Empty\x12^\n" + + "\x1bUpdateRepositoryMountStatus\x12'.manager.UpdateRepositoryMountStatusReq\x1a\x16.google.protobuf.EmptyB\vZ\t.;managerb\x06proto3" var ( file_manager_proto_rawDescOnce sync.Once @@ -1204,48 +1592,58 @@ func file_manager_proto_rawDescGZIP() []byte { return file_manager_proto_rawDescData } -var file_manager_proto_msgTypes = make([]protoimpl.MessageInfo, 13) +var file_manager_proto_msgTypes = make([]protoimpl.MessageInfo, 18) var file_manager_proto_goTypes = []any{ (*RegisterRequest)(nil), // 0: manager.RegisterRequest (*RegisterResponse)(nil), // 1: manager.RegisterResponse (*HeartbeatRequest)(nil), // 2: manager.HeartbeatRequest - (*SchedulerFileRequest)(nil), // 3: manager.SchedulerFileRequest - (*SyncFileProcessReq)(nil), // 4: manager.SyncFileProcessReq - (*FileProcessEntry)(nil), // 5: manager.FileProcessEntry - (*SchedulerFileResponse)(nil), // 6: manager.SchedulerFileResponse - (*FileProcessRequest)(nil), // 7: manager.FileProcessRequest - (*DeleteByEtagsAndFieldsRequest)(nil), // 8: manager.DeleteByEtagsAndFieldsRequest - (*CreateCacheJobReq)(nil), // 9: manager.CreateCacheJobReq - (*CreateCacheJobResp)(nil), // 10: manager.CreateCacheJobResp - (*UpdateCacheJobStatusReq)(nil), // 11: manager.UpdateCacheJobStatusReq - (*UpdateRepositoryMountStatusReq)(nil), // 12: manager.UpdateRepositoryMountStatusReq - (*emptypb.Empty)(nil), // 13: google.protobuf.Empty + (*NodeHealthSnapshot)(nil), // 3: manager.NodeHealthSnapshot + (*CapabilityObservation)(nil), // 4: manager.CapabilityObservation + (*StorageErrorCount)(nil), // 5: manager.StorageErrorCount + (*SchedulerFileRequest)(nil), // 6: manager.SchedulerFileRequest + (*SyncFileProcessReq)(nil), // 7: manager.SyncFileProcessReq + (*FileProcessEntry)(nil), // 8: manager.FileProcessEntry + (*SchedulerFileResponse)(nil), // 9: manager.SchedulerFileResponse + (*FileProcessRequest)(nil), // 10: manager.FileProcessRequest + (*DeleteByEtagsAndFieldsRequest)(nil), // 11: manager.DeleteByEtagsAndFieldsRequest + (*CreateCacheJobReq)(nil), // 12: manager.CreateCacheJobReq + (*CreateCacheJobResp)(nil), // 13: manager.CreateCacheJobResp + (*UpdateCacheJobStatusReq)(nil), // 14: manager.UpdateCacheJobStatusReq + (*UpdateRepositoryMountStatusReq)(nil), // 15: manager.UpdateRepositoryMountStatusReq + (*IngestRepositoryRequest)(nil), // 16: manager.IngestRepositoryRequest + (*IngestRepositoryResponse)(nil), // 17: manager.IngestRepositoryResponse + (*emptypb.Empty)(nil), // 18: google.protobuf.Empty } var file_manager_proto_depIdxs = []int32{ - 5, // 0: manager.SyncFileProcessReq.fileProcessEntries:type_name -> manager.FileProcessEntry - 0, // 1: manager.Manager.Register:input_type -> manager.RegisterRequest - 2, // 2: manager.Manager.Heartbeat:input_type -> manager.HeartbeatRequest - 3, // 3: manager.Manager.SchedulerFile:input_type -> manager.SchedulerFileRequest - 7, // 4: manager.Manager.ReportFileProcess:input_type -> manager.FileProcessRequest - 4, // 5: manager.Manager.SyncFileProcess:input_type -> manager.SyncFileProcessReq - 8, // 6: manager.Manager.DeleteByEtagsAndFields:input_type -> manager.DeleteByEtagsAndFieldsRequest - 9, // 7: manager.Manager.CreateCacheJob:input_type -> manager.CreateCacheJobReq - 11, // 8: manager.Manager.UpdateCacheJobStatus:input_type -> manager.UpdateCacheJobStatusReq - 12, // 9: manager.Manager.UpdateRepositoryMountStatus:input_type -> manager.UpdateRepositoryMountStatusReq - 1, // 10: manager.Manager.Register:output_type -> manager.RegisterResponse - 13, // 11: manager.Manager.Heartbeat:output_type -> google.protobuf.Empty - 6, // 12: manager.Manager.SchedulerFile:output_type -> manager.SchedulerFileResponse - 13, // 13: manager.Manager.ReportFileProcess:output_type -> google.protobuf.Empty - 13, // 14: manager.Manager.SyncFileProcess:output_type -> google.protobuf.Empty - 13, // 15: manager.Manager.DeleteByEtagsAndFields:output_type -> google.protobuf.Empty - 10, // 16: manager.Manager.CreateCacheJob:output_type -> manager.CreateCacheJobResp - 13, // 17: manager.Manager.UpdateCacheJobStatus:output_type -> google.protobuf.Empty - 13, // 18: manager.Manager.UpdateRepositoryMountStatus:output_type -> google.protobuf.Empty - 10, // [10:19] is the sub-list for method output_type - 1, // [1:10] is the sub-list for method input_type - 1, // [1:1] is the sub-list for extension type_name - 1, // [1:1] is the sub-list for extension extendee - 0, // [0:1] is the sub-list for field type_name + 3, // 0: manager.HeartbeatRequest.health:type_name -> manager.NodeHealthSnapshot + 4, // 1: manager.NodeHealthSnapshot.capabilities:type_name -> manager.CapabilityObservation + 5, // 2: manager.NodeHealthSnapshot.errors:type_name -> manager.StorageErrorCount + 8, // 3: manager.SyncFileProcessReq.fileProcessEntries:type_name -> manager.FileProcessEntry + 16, // 4: manager.Manager.IngestRepository:input_type -> manager.IngestRepositoryRequest + 0, // 5: manager.Manager.Register:input_type -> manager.RegisterRequest + 2, // 6: manager.Manager.Heartbeat:input_type -> manager.HeartbeatRequest + 6, // 7: manager.Manager.SchedulerFile:input_type -> manager.SchedulerFileRequest + 10, // 8: manager.Manager.ReportFileProcess:input_type -> manager.FileProcessRequest + 7, // 9: manager.Manager.SyncFileProcess:input_type -> manager.SyncFileProcessReq + 11, // 10: manager.Manager.DeleteByEtagsAndFields:input_type -> manager.DeleteByEtagsAndFieldsRequest + 12, // 11: manager.Manager.CreateCacheJob:input_type -> manager.CreateCacheJobReq + 14, // 12: manager.Manager.UpdateCacheJobStatus:input_type -> manager.UpdateCacheJobStatusReq + 15, // 13: manager.Manager.UpdateRepositoryMountStatus:input_type -> manager.UpdateRepositoryMountStatusReq + 17, // 14: manager.Manager.IngestRepository:output_type -> manager.IngestRepositoryResponse + 1, // 15: manager.Manager.Register:output_type -> manager.RegisterResponse + 18, // 16: manager.Manager.Heartbeat:output_type -> google.protobuf.Empty + 9, // 17: manager.Manager.SchedulerFile:output_type -> manager.SchedulerFileResponse + 18, // 18: manager.Manager.ReportFileProcess:output_type -> google.protobuf.Empty + 18, // 19: manager.Manager.SyncFileProcess:output_type -> google.protobuf.Empty + 18, // 20: manager.Manager.DeleteByEtagsAndFields:output_type -> google.protobuf.Empty + 13, // 21: manager.Manager.CreateCacheJob:output_type -> manager.CreateCacheJobResp + 18, // 22: manager.Manager.UpdateCacheJobStatus:output_type -> google.protobuf.Empty + 18, // 23: manager.Manager.UpdateRepositoryMountStatus:output_type -> google.protobuf.Empty + 14, // [14:24] is the sub-list for method output_type + 4, // [4:14] is the sub-list for method input_type + 4, // [4:4] is the sub-list for extension type_name + 4, // [4:4] is the sub-list for extension extendee + 0, // [0:4] is the sub-list for field type_name } func init() { file_manager_proto_init() } @@ -1259,7 +1657,7 @@ func file_manager_proto_init() { GoPackagePath: reflect.TypeOf(x{}).PkgPath(), RawDescriptor: unsafe.Slice(unsafe.StringData(file_manager_proto_rawDesc), len(file_manager_proto_rawDesc)), NumEnums: 0, - NumMessages: 13, + NumMessages: 18, NumExtensions: 0, NumServices: 1, }, diff --git a/pkg/proto/manager/manager_grpc.pb.go b/pkg/proto/manager/manager_grpc.pb.go index f7f5beb..c617498 100644 --- a/pkg/proto/manager/manager_grpc.pb.go +++ b/pkg/proto/manager/manager_grpc.pb.go @@ -1,7 +1,7 @@ // Code generated by protoc-gen-go-grpc. DO NOT EDIT. // versions: // - protoc-gen-go-grpc v1.5.1 -// - protoc v5.29.3 +// - protoc v7.35.1 // source: manager.proto package manager @@ -20,6 +20,7 @@ import ( const _ = grpc.SupportPackageIsVersion9 const ( + Manager_IngestRepository_FullMethodName = "/manager.Manager/IngestRepository" Manager_Register_FullMethodName = "/manager.Manager/Register" Manager_Heartbeat_FullMethodName = "/manager.Manager/Heartbeat" Manager_SchedulerFile_FullMethodName = "/manager.Manager/SchedulerFile" @@ -37,6 +38,7 @@ const ( // // 用户服务定义 type ManagerClient interface { + IngestRepository(ctx context.Context, in *IngestRepositoryRequest, opts ...grpc.CallOption) (*IngestRepositoryResponse, error) // 启动注册方法 Register(ctx context.Context, in *RegisterRequest, opts ...grpc.CallOption) (*RegisterResponse, error) // 心跳方法 @@ -63,6 +65,16 @@ func NewManagerClient(cc grpc.ClientConnInterface) ManagerClient { return &managerClient{cc} } +func (c *managerClient) IngestRepository(ctx context.Context, in *IngestRepositoryRequest, opts ...grpc.CallOption) (*IngestRepositoryResponse, error) { + cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) + out := new(IngestRepositoryResponse) + err := c.cc.Invoke(ctx, Manager_IngestRepository_FullMethodName, in, out, cOpts...) + if err != nil { + return nil, err + } + return out, nil +} + func (c *managerClient) Register(ctx context.Context, in *RegisterRequest, opts ...grpc.CallOption) (*RegisterResponse, error) { cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...) out := new(RegisterResponse) @@ -159,6 +171,7 @@ func (c *managerClient) UpdateRepositoryMountStatus(ctx context.Context, in *Upd // // 用户服务定义 type ManagerServer interface { + IngestRepository(context.Context, *IngestRepositoryRequest) (*IngestRepositoryResponse, error) // 启动注册方法 Register(context.Context, *RegisterRequest) (*RegisterResponse, error) // 心跳方法 @@ -185,6 +198,9 @@ type ManagerServer interface { // pointer dereference when methods are called. type UnimplementedManagerServer struct{} +func (UnimplementedManagerServer) IngestRepository(context.Context, *IngestRepositoryRequest) (*IngestRepositoryResponse, error) { + return nil, status.Errorf(codes.Unimplemented, "method IngestRepository not implemented") +} func (UnimplementedManagerServer) Register(context.Context, *RegisterRequest) (*RegisterResponse, error) { return nil, status.Errorf(codes.Unimplemented, "method Register not implemented") } @@ -233,6 +249,24 @@ func RegisterManagerServer(s grpc.ServiceRegistrar, srv ManagerServer) { s.RegisterService(&Manager_ServiceDesc, srv) } +func _Manager_IngestRepository_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) { + in := new(IngestRepositoryRequest) + if err := dec(in); err != nil { + return nil, err + } + if interceptor == nil { + return srv.(ManagerServer).IngestRepository(ctx, in) + } + info := &grpc.UnaryServerInfo{ + Server: srv, + FullMethod: Manager_IngestRepository_FullMethodName, + } + handler := func(ctx context.Context, req interface{}) (interface{}, error) { + return srv.(ManagerServer).IngestRepository(ctx, req.(*IngestRepositoryRequest)) + } + return interceptor(ctx, in, info, handler) +} + func _Manager_Register_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) { in := new(RegisterRequest) if err := dec(in); err != nil { @@ -402,6 +436,10 @@ var Manager_ServiceDesc = grpc.ServiceDesc{ ServiceName: "manager.Manager", HandlerType: (*ManagerServer)(nil), Methods: []grpc.MethodDesc{ + { + MethodName: "IngestRepository", + Handler: _Manager_IngestRepository_Handler, + }, { MethodName: "Register", Handler: _Manager_Register_Handler, diff --git a/pkg/repository/ambiguity_probe_test.go b/pkg/repository/ambiguity_probe_test.go new file mode 100644 index 0000000..68c11cd --- /dev/null +++ b/pkg/repository/ambiguity_probe_test.go @@ -0,0 +1,46 @@ +package repository + +import ( + "os" + "path/filepath" + "testing" +) + +// This probe documents the ambiguity that constrains the pruning rule. Two +// on-disk situations are byte-for-byte indistinguishable without extra +// information: +// +// A. orphan residue: api/models/alice/ghost/paths-info/... where the marker was +// deleted with its repository +// B. a legitimate repository whose own name begins with a reserved word, i.e. +// api/models/alice//paths-info/... where == "paths-info/repo" +// +// In (B) the directory named "paths-info" is a repository root's ancestor +// segment, not data, and pruning it would lose a real repository. The test +// asserts both trees can coexist and that only the marker-bearing one is found. +func TestPathsInfoAmbiguityIsResolvedByMarkerOnly(t *testing.T) { + root := t.TempDir() + // (B) a legitimate repository whose first segment is "paths-info". + legit := RepoKey{Namespace: "alice", RepoType: "models", Repo: "paths-info/repo"} + if err := Register(root, Hosted(legit)); err != nil { + t.Fatal(err) + } + // (A) orphan residue under a sibling namespace. + orphan := filepath.Join(root, "api", "models", "alice", "ghost") + if err := os.MkdirAll(filepath.Join(orphan, "paths-info", "commit", "w", "f.bin"), 0o755); err != nil { + t.Fatal(err) + } + + all, err := List(root) + if err != nil { + t.Fatal(err) + } + // Exactly the legitimate repository is found; the orphan contributes nothing + // and must not error either. + if len(all) != 1 { + t.Fatalf("expected only the marker-bearing repository, got %+v", all) + } + if all[0].RepoKey != legit { + t.Fatalf("wrong repository discovered: %+v", all[0]) + } +} \ No newline at end of file diff --git a/pkg/repository/conflict_scope_test.go b/pkg/repository/conflict_scope_test.go new file mode 100644 index 0000000..c354f90 --- /dev/null +++ b/pkg/repository/conflict_scope_test.go @@ -0,0 +1,148 @@ +package repository + +import ( + "fmt" + "math/rand" + "testing" +) + +// Register now scans only the repositories its conflict rules can act on. This +// asserts that scoping is verdict-preserving: for a given on-disk registry, the +// scoped view must yield the same decision and the same reason as the +// whole-registry view, across a wide spread of candidate descriptors. +func TestScopedConflictViewMatchesWholeRegistry(t *testing.T) { + seedRegistries := []struct { + name string + keys []Descriptor + }{ + {"hosted siblings", []Descriptor{ + Hosted(RepoKey{"alice", "models", "team/model-a"}), + Hosted(RepoKey{"alice", "models", "team/model-b"}), + Hosted(RepoKey{"bob", "models", "team/model-a"}), + }}, + {"namespace case", []Descriptor{ + Hosted(RepoKey{"Alice", "models", "a"}), + }}, + {"cross repo type", []Descriptor{ + Hosted(RepoKey{"alice", "models", "m"}), + Hosted(RepoKey{"alice", "datasets", "d"}), + }}, + {"remote providers", []Descriptor{ + Remote(RepoKey{HuggingFace, "models", "Qwen/Qwen3"}), + Remote(RepoKey{ModelScope, "models", "Qwen/Qwen3"}), + Hosted(RepoKey{"alice", "models", "local-one"}), + }}, + {"ancestor pairs", []Descriptor{ + Hosted(RepoKey{"alice", "models", "team"}), + Hosted(RepoKey{"alice", "models", "other/child"}), + }}, + } + + candidates := []Descriptor{ + Hosted(RepoKey{"alice", "models", "team/model-a"}), + Hosted(RepoKey{"alice", "models", "team/model-c"}), + Hosted(RepoKey{"Alice", "models", "newthing"}), + Hosted(RepoKey{"alice", "datasets", "fresh"}), + Hosted(RepoKey{"bob", "models", "team/model-a"}), + Hosted(RepoKey{"alice", "models", "team"}), + Hosted(RepoKey{"alice", "models", "team/model-a/deeper"}), + Hosted(RepoKey{"carol", "spaces", "demo"}), + Remote(RepoKey{HuggingFace, "models", "Qwen/Qwen3"}), + Remote(RepoKey{HuggingFace, "models", "Qwen/Qwen4"}), + Remote(RepoKey{ModelScope, "models", "Qwen/Qwen3"}), + Remote(RepoKey{ModelScope, "models", "Other/Thing"}), + Hosted(RepoKey{Local, "models", "plain"}), + Hosted(RepoKey{Local, "models", "team/nested"}), + } + + for _, registry := range seedRegistries { + for _, candidate := range candidates { + t.Run(fmt.Sprintf("%s/%s-%s", registry.name, candidate.Namespace, candidate.Repo), func(t *testing.T) { + root := t.TempDir() + for _, d := range registry.keys { + // Seeds may conflict with one another; only the resulting disk + // state matters here. + _ = Register(root, d) + } + whole, err := List(root) + if err != nil { + // A case-vs-directory mismatch is reported by List itself and is + // a pre-existing property of the registry, not of the scoped + // view; such a registry has no coherent verdict to compare. + t.Skipf("registry not enumerable as-is: %v", err) + } + scoped, err := listConflicts(root, candidate) + if err != nil { + // Same pre-existing enumeration property as above, reached + // through the scoped path. + t.Skipf("registry not enumerable as-is: %v", err) + } + wholeErr := registrationConflicts(root, candidate, whole) + scopedErr := registrationConflicts(root, candidate, scoped) + if (wholeErr == nil) != (scopedErr == nil) { + t.Fatalf("verdict drift for %+v:\n whole view: %v\n scoped view: %v", candidate.RepoKey, wholeErr, scopedErr) + } + if wholeErr != nil && scopedErr != nil && wholeErr.Error() != scopedErr.Error() { + t.Fatalf("reason drift for %+v: whole=%v scoped=%v", candidate.RepoKey, wholeErr, scopedErr) + } + // Every repository that actually conflicts with the candidate must + // be present in the scoped view. + for _, old := range whole { + if registrationConflicts(root, candidate, []Descriptor{old}) == nil { + continue + } + found := false + for _, s := range scoped { + if s.RepoKey == old.RepoKey { + found = true + break + } + } + if !found { + t.Fatalf("scoped view omitted conflicting repository %+v (scoped %+v)", old.RepoKey, scoped) + } + } + }) + } + } +} + +// Randomised differential test: many registries, many candidates, the scoped and +// whole-registry conflict views must never disagree. +func TestScopedConflictViewRandomised(t *testing.T) { + rng := rand.New(rand.NewSource(20260922)) + namespaces := []string{"alice", "bob", "carol", Local, HuggingFace, ModelScope} + types := []string{"models", "datasets", "spaces"} + repos := []string{"a", "b", "team/x", "team/y", "team", "deep/nest/z"} + + for iter := 0; iter < 150; iter++ { + root := t.TempDir() + for n := 0; n < rng.Intn(6); n++ { + k := RepoKey{namespaces[rng.Intn(len(namespaces))], types[rng.Intn(len(types))], repos[rng.Intn(len(repos))]} + d := Hosted(k) + if k.Namespace == HuggingFace || k.Namespace == ModelScope { + d = Remote(k) + } + _ = Register(root, d) + } + k := RepoKey{namespaces[rng.Intn(len(namespaces))], types[rng.Intn(len(types))], repos[rng.Intn(len(repos))]} + candidate := Hosted(k) + if k.Namespace == HuggingFace || k.Namespace == ModelScope { + candidate = Remote(k) + } + whole, err := List(root) + if err != nil { + continue // not enumerable as-is; covered by other tests + } + scoped, err := listConflicts(root, candidate) + if err != nil { + t.Fatal(err) + } + wholeErr := registrationConflicts(root, candidate, whole) + scopedErr := registrationConflicts(root, candidate, scoped) + if (wholeErr == nil) != (scopedErr == nil) { + t.Fatalf("iter %d: verdict drift for %+v whole=%v scoped=%v\n whole=%+v\n scoped=%+v", + iter, candidate.RepoKey, wholeErr, scopedErr, whole, scoped) + } + } +} \ No newline at end of file diff --git a/pkg/repository/layout_test.go b/pkg/repository/layout_test.go new file mode 100644 index 0000000..55b0d8b --- /dev/null +++ b/pkg/repository/layout_test.go @@ -0,0 +1,32 @@ +package repository + +import ( + "path/filepath" + "testing" +) + +func TestLegacyRemoteAndNamespacedUploadPhysicalPaths(t *testing.T) { + root := t.TempDir() + for _, tc := range []struct{ namespace, repo, files, api string }{ + {HuggingFace, "Qwen/demo", "files/models/Qwen/demo", "api/models/Qwen/demo"}, + {ModelScope, "Qwen/demo", "modelscope/models/Qwen/demo", "api/models/modelscope/Qwen/demo"}, + {"alice", "demo", "files/models/dingo-local/alice/demo", "api/models/dingo-local/alice/demo"}, + {"bob", "demo", "files/models/dingo-local/bob/demo", "api/models/dingo-local/bob/demo"}, + {Company, "demo", "files/models/dingo-local/datacanvas/demo", "api/models/dingo-local/datacanvas/demo"}, + } { + k := RepoKey{Namespace: tc.namespace, RepoType: "models", Repo: tc.repo} + if k.FilesRoot(root) != filepath.Join(root, filepath.FromSlash(tc.files)) || k.APIRoot(root) != filepath.Join(root, filepath.FromSlash(tc.api)) { + t.Fatalf("wrong layout for %+v: %s %s", k, k.FilesRoot(root), k.APIRoot(root)) + } + } +} + +func TestLegacyAliasCannotOverlapPersonalRepository(t *testing.T) { + root := t.TempDir() + if err := Register(root, Hosted(RepoKey{Namespace: "alice", RepoType: "models", Repo: "demo"})); err != nil { + t.Fatal(err) + } + if err := Register(root, Hosted(RepoKey{Namespace: Local, RepoType: "models", Repo: "alice/demo"})); err == nil { + t.Fatal("legacy alias shares physical storage with another namespace") + } +} diff --git a/pkg/repository/register_cost_test.go b/pkg/repository/register_cost_test.go new file mode 100644 index 0000000..04795df --- /dev/null +++ b/pkg/repository/register_cost_test.go @@ -0,0 +1,90 @@ +package repository + +import ( + "strconv" + "testing" +) + +// buildNamespacedRegistry creates n namespaces with perNamespace hosted +// repositories each, so scoping can be observed against namespace count. +func buildNamespacedRegistry(t *testing.T, root string, namespaces, perNamespace int) { + t.Helper() + for n := 0; n < namespaces; n++ { + ns := "ns-" + strconv.Itoa(n) + for i := 0; i < perNamespace; i++ { + k := RepoKey{Namespace: ns, RepoType: "models", Repo: "team/repo-" + strconv.Itoa(i)} + if err := Register(root, Hosted(k)); err != nil { + t.Fatal(err) + } + } + } +} + +// The scoped conflict view must not grow with repositories in unrelated +// namespaces, while the whole-registry view does. +func TestScopedConflictSetDoesNotGrowWithUnrelatedNamespaces(t *testing.T) { + candidate := Hosted(RepoKey{Namespace: "target", RepoType: "models", Repo: "team/new"}) + + measure := func(namespaces, perNamespace int) (scopedCount, wholeCount int) { + root := t.TempDir() + buildNamespacedRegistry(t, root, namespaces, perNamespace) + scoped, err := listConflicts(root, candidate) + if err != nil { + t.Fatal(err) + } + whole, err := List(root) + if err != nil { + t.Fatal(err) + } + return len(scoped), len(whole) + } + + scopedSmall, wholeSmall := measure(2, 10) + scopedLarge, wholeLarge := measure(40, 10) + t.Logf("2 namespaces : scoped=%d whole=%d", scopedSmall, wholeSmall) + t.Logf("40 namespaces: scoped=%d whole=%d", scopedLarge, wholeLarge) + + if wholeLarge <= wholeSmall { + t.Fatalf("whole registry did not grow: %d -> %d", wholeSmall, wholeLarge) + } + if scopedLarge != scopedSmall { + t.Fatalf("scoped conflict set grew with unrelated namespaces: %d -> %d", scopedSmall, scopedLarge) + } +} + +// End-to-end: registering a batch against an existing registry must not rescan +// the whole registry per call. Work may grow quadratically within the candidate's +// own namespace, but must never include the unrelated namespaces. +func TestRegisterBatchDoesNotRescanWholeRegistry(t *testing.T) { + const seedNamespaces, seedPerNamespace = 20, 10 + const targetPerNamespace = 5 + + work := func(batch int) int { + root := t.TempDir() + buildNamespacedRegistry(t, root, seedNamespaces, seedPerNamespace) + total := 0 + for i := 0; i < batch; i++ { + candidate := Hosted(RepoKey{Namespace: "target", RepoType: "models", Repo: "team/new-" + strconv.Itoa(i)}) + scoped, err := listConflicts(root, candidate) + if err != nil { + t.Fatal(err) + } + total += len(scoped) + if err := Register(root, candidate); err != nil { + t.Fatal(err) + } + } + return total + } + + small := work(targetPerNamespace) + large := work(targetPerNamespace * 3) + // The seeded registry alone holds 200 repositories. If any step rescanned it, + // these totals would be at least that large. + const seededRepositories = seedNamespaces * seedPerNamespace + t.Logf("conflict entries examined: batch %d -> %d, batch %d -> %d (seeded registry holds %d)", + targetPerNamespace, small, targetPerNamespace*3, large, seededRepositories) + if small >= seededRepositories || large >= seededRepositories { + t.Fatalf("batch work included the whole seeded registry: %d/%d, seeded=%d", small, large, seededRepositories) + } +} \ No newline at end of file diff --git a/pkg/repository/register_fastpath_test.go b/pkg/repository/register_fastpath_test.go new file mode 100644 index 0000000..975860a --- /dev/null +++ b/pkg/repository/register_fastpath_test.go @@ -0,0 +1,34 @@ +package repository + +import ( + "os" + "path/filepath" + "testing" +) + +func TestRegisteredRepositoryDoesNotEnumerateUnrelatedMarkers(t *testing.T) { + root := t.TempDir() + d := Hosted(RepoKey{"alice", "models", "existing"}) + if err := Register(root, d); err != nil { + t.Fatal(err) + } + bad := filepath.Join(root, "api", "models", Local, "alice", "unrelated", Marker) + if err := os.MkdirAll(filepath.Dir(bad), 0755); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(bad, []byte("invalid"), 0644); err != nil { + t.Fatal(err) + } + if err := Register(root, d); err != nil { + t.Fatalf("idempotent registration scanned unrelated repository: %v", err) + } + if err := Register(root, Hosted(RepoKey{"alice", "models", "new"})); err == nil { + t.Fatal("new registration bypassed conflict discovery") + } + if err := os.WriteFile(filepath.Join(d.APIRoot(root), Marker), []byte("invalid"), 0644); err != nil { + t.Fatal(err) + } + if err := Register(root, d); err == nil { + t.Fatal("accepted damaged own descriptor") + } +} diff --git a/pkg/repository/repository.go b/pkg/repository/repository.go new file mode 100644 index 0000000..73abec4 --- /dev/null +++ b/pkg/repository/repository.go @@ -0,0 +1,625 @@ +// Package repository owns platform repository identity, layout and registration. +package repository + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "github.com/gofrs/flock" + "io/fs" + "net/url" + "os" + "path/filepath" + "strings" + "sync" + "time" + "unicode" + "unicode/utf8" +) + +const Marker = "repository.json" +const HuggingFace = "huggingface" +const ModelScope = "modelscope" +const Local = "dingo-local" +const Company = "datacanvas" + +type RepoKey struct { + Namespace string `json:"namespace"` + RepoType string `json:"repoType"` + Repo string `json:"repo"` +} + +func (k RepoKey) ID() string { return k.Namespace + "/" + k.Repo } +func (k RepoKey) Validate() error { + if k.RepoType != "models" && k.RepoType != "datasets" && k.RepoType != "spaces" { + return fmt.Errorf("invalid repoType") + } + if err := Segment(k.Namespace); err != nil { + return fmt.Errorf("namespace: %w", err) + } + if err := Relative(k.Repo); err != nil { + return fmt.Errorf("repo: %w", err) + } + return nil +} + +// ParseID is only for values already known to be full platform repository IDs. +func ParseID(repoType, id string) (RepoKey, error) { + p := strings.SplitN(id, "/", 2) + if len(p) != 2 { + return RepoKey{}, fmt.Errorf("repository ID must include namespace") + } + k := RepoKey{p[0], repoType, p[1]} + return k, k.Validate() +} + +func Segment(s string) error { + if s == "" || s == "." || s == ".." || !utf8.ValidString(s) || len(s) > 255 || strings.ContainsAny(s, `/\:*?"<>|`) || strings.TrimSpace(s) != s || strings.HasSuffix(s, ".") { + return fmt.Errorf("unsafe path segment %q", s) + } + for _, r := range s { + if unicode.IsControl(r) { + return fmt.Errorf("control character in path") + } + } + base := strings.ToLower(strings.SplitN(s, ".", 2)[0]) + if base == "con" || base == "prn" || base == "aux" || base == "nul" || (len(base) == 4 && (strings.HasPrefix(base, "com") || strings.HasPrefix(base, "lpt")) && base[3] >= '1' && base[3] <= '9') { + return fmt.Errorf("reserved platform path name") + } + return nil +} +func Relative(s string) error { + if len(s) > 1024 { + return fmt.Errorf("relative path exceeds 1024 bytes") + } + for _, p := range strings.Split(s, "/") { + if err := Segment(p); err != nil { + return err + } + } + return nil +} +func (k RepoKey) FilesRoot(root string) string { + if k.Namespace == ModelScope { + return filepath.Join(root, "modelscope", k.RepoType, filepath.FromSlash(k.Repo)) + } + return filepath.Join(root, "files", k.RepoType, k.storagePath()) +} +func (k RepoKey) APIRoot(root string) string { + if k.Namespace == ModelScope { + return filepath.Join(root, "api", k.RepoType, ModelScope, filepath.FromSlash(k.Repo)) + } + return filepath.Join(root, "api", k.RepoType, k.storagePath()) +} + +// Provider identity is virtual: HF keeps its original upstream directory and +// hosted namespaces live exclusively below the local upload root. +func (k RepoKey) storagePath() string { + if k.Namespace == HuggingFace { + return filepath.FromSlash(k.Repo) + } + // Existing pre-namespace uploads remain addressable until explicitly migrated. + if k.Namespace == Local { + return filepath.Join(Local, filepath.FromSlash(k.Repo)) + } + return filepath.Join(Local, k.Namespace, filepath.FromSlash(k.Repo)) +} +func (k RepoKey) Blob(root, sha string) string { return filepath.Join(k.FilesRoot(root), "blobs", sha) } +func (k RepoKey) Resolve(root, commit, path string) string { + return filepath.Join(k.FilesRoot(root), "resolve", commit, filepath.FromSlash(path)) +} +func EscapeURLPath(value string) string { + parts := strings.Split(value, "/") + for i := range parts { + parts[i] = url.PathEscape(parts[i]) + } + return strings.Join(parts, "/") +} +func (k RepoKey) PathsInfo(root, commit, path string) string { + return filepath.Join(k.APIRoot(root), "paths-info", commit, filepath.FromSlash(path)) +} +func (k RepoKey) Revision(root, revision string) string { + return filepath.Join(k.APIRoot(root), "revision", revision) +} + +type Descriptor struct { + RepoKey + Version int `json:"version"` + Source string `json:"source"` + Persistent bool `json:"persistent"` + Provider string `json:"provider,omitempty"` + UpstreamRepo string `json:"upstreamRepo,omitempty"` + Format string `json:"format"` +} + +func Hosted(k RepoKey) Descriptor { + return Descriptor{RepoKey: k, Version: 1, Source: "hosted", Persistent: true, Format: "dingcache"} +} +func Remote(k RepoKey) Descriptor { + return Descriptor{RepoKey: k, Version: 1, Source: "remote", Provider: k.Namespace, UpstreamRepo: k.Repo, Format: "dingcache"} +} +func (d Descriptor) Validate() error { + if err := d.RepoKey.Validate(); err != nil { + return err + } + if d.Version != 1 { + return fmt.Errorf("unsupported repository descriptor version") + } + switch d.Source { + case "hosted": + if !d.Persistent || d.Provider != "" || d.UpstreamRepo != "" || d.Format != "dingcache" || strings.EqualFold(d.Namespace, HuggingFace) || strings.EqualFold(d.Namespace, ModelScope) { + return fmt.Errorf("invalid hosted source policy") + } + case "remote": + if d.Persistent || (d.Provider != HuggingFace && d.Provider != ModelScope) || d.Namespace != d.Provider || d.UpstreamRepo != d.Repo { + return fmt.Errorf("invalid remote source policy") + } + if d.Format != "dingcache" { + return fmt.Errorf("invalid remote format") + } + default: + return fmt.Errorf("unknown repository source") + } + return nil +} +func Read(root string, k RepoKey) (Descriptor, error) { + if err := k.Validate(); err != nil { + return Descriptor{}, err + } + if err := SafePath(root, filepath.Join(k.APIRoot(root), Marker)); err != nil { + return Descriptor{}, err + } + b, err := os.ReadFile(filepath.Join(k.APIRoot(root), Marker)) + if err != nil { + return Descriptor{}, err + } + var d Descriptor + if err = json.Unmarshal(b, &d); err != nil { + return d, err + } + if d.RepoKey != k { + return d, fmt.Errorf("repository marker identity mismatch") + } + return d, d.Validate() +} + +func List(root string) ([]Descriptor, error) { + return listIn(filepath.Join(root, "api"), root) +} + +// ListHosted enumerates only hosted repositories, which live exclusively below +// api//dingo-local. Remote upstream caches share the same api/ tree and +// outnumber uploads by orders of magnitude, so scanning all of api/ to then drop +// every non-hosted descriptor costs a full-tree walk per call on cheap triggers +// such as the inventory ticker. Walking the hosted roots instead keeps the scan +// proportional to the uploads we actually report. +// +// This is an enumeration-scope change only: a descriptor the caller could not +// reach before is unreachable now, because hosted keys always resolve below the +// roots walked here. +func ListHosted(root string) ([]Descriptor, error) { + result := []Descriptor{} + for _, repoType := range []string{"models", "datasets", "spaces"} { + found, err := listIn(filepath.Join(root, "api", repoType, Local), root) + if err != nil { + return result, err + } + result = append(result, found...) + } + return result, nil +} + +func isDataSubtree(name string) bool { + switch name { + case "paths-info", "revision", "recycle", "blobs", "resolve": + return true + } + return false +} + +// isRepositoryDataSubtree reports whether a directory walk must stop descending +// here because the directory is platform-owned repository data rather than a +// candidate repository root. +// +// A data subtree is identified by the marker above it, never by its own name: +// repository names are user input, so "team/resolve/model" and even +// "paths-info/repo" are legitimate repositories whose own paths contain a +// data-subtree word. Inside a recognised repository, though, these names are +// reserved platform layout, so nothing below one can be a repository root. +func isRepositoryDataSubtree(scanRoot, p, name string) bool { + if p == scanRoot { + return false + } + if !isDataSubtree(name) { + return false + } + // Either this directory itself is data (a marker above it), or it is nested + // inside a data subtree that the walk already declined to enter. + return hasRepositoryAncestor(scanRoot, p) || hasDataSubtreeAncestor(scanRoot, p) +} + +// hasDataSubtreeAncestor reports whether any strict ancestor of p below the scan +// root is a data subtree. +func hasDataSubtreeAncestor(scanRoot, p string) bool { + scanRoot = filepath.Clean(scanRoot) + for dir := filepath.Dir(filepath.Clean(p)); dir != scanRoot; dir = filepath.Dir(dir) { + if len(dir) <= len(scanRoot) { + return false + } + if isDataSubtree(filepath.Base(dir)) { + return true + } + } + return false +} + +// hasRepositoryAncestor reports whether any strict ancestor of p below the scan +// root carries a repository marker. Walking a handful of parents through Lstat +// is far cheaper than descending a data subtree that can hold one directory per +// stored file. +func hasRepositoryAncestor(scanRoot, p string) bool { + scanRoot = filepath.Clean(scanRoot) + for dir := filepath.Dir(p); ; dir = filepath.Dir(dir) { + if len(dir) < len(scanRoot) || (dir != scanRoot && !strings.HasPrefix(dir, scanRoot+string(filepath.Separator))) { + return false + } + if _, err := os.Lstat(filepath.Join(dir, Marker)); err == nil { + return true + } + if dir == scanRoot { + return false + } + } +} + +func listIn(scanRoot, root string) ([]Descriptor, error) { + result := []Descriptor{} + err := filepath.WalkDir(scanRoot, func(p string, e fs.DirEntry, err error) error { + if err != nil { + return err + } + if !e.IsDir() { + return nil + } + // Data subtrees under a repository are not repositories themselves, and + // descending them turns discovery into a walk of every stored file. + if isRepositoryDataSubtree(scanRoot, p, e.Name()) { + return filepath.SkipDir + } + marker := filepath.Join(p, Marker) + info, statErr := os.Lstat(marker) + if errors.Is(statErr, os.ErrNotExist) { + return nil + } + if statErr != nil { + return statErr + } + if info.IsDir() { + return nil + } + b, err := os.ReadFile(marker) + if errors.Is(err, os.ErrNotExist) { + return nil + } + if err != nil { + return err + } + if err = SafePath(root, marker); err != nil { + return err + } + var d Descriptor + if err = json.Unmarshal(b, &d); err != nil { + return err + } + if err = d.Validate(); err != nil { + return err + } + if filepath.Clean(filepath.Join(d.APIRoot(root), Marker)) != filepath.Clean(marker) { + return fmt.Errorf("repository marker path mismatch: %s", marker) + } + result = append(result, d) + return filepath.SkipDir + }) + if errors.Is(err, os.ErrNotExist) { + return result, nil + } + return result, err +} + +var registrationMu sync.Mutex + +// listConflicts returns exactly the registered repositories that the conflict +// checks in Register can act on, without scanning the whole registry. +// +// Every check in registrationConflicts compares against the candidate's namespace +// or one of its own path prefixes, so a repository can only conflict when its +// storage path can overlap the candidate's. Hosted namespaces occupy disjoint +// directories below dingo-local, so repositories under other namespaces can never +// overlap - with one exception: a bare local-namespace repository is stored at +// dingo-local/, so its path can be a prefix of any other namespace's +// directory and must be considered for every hosted candidate. +func listConflicts(root string, d Descriptor) ([]Descriptor, error) { + result := []Descriptor{} + seen := map[RepoKey]struct{}{} + scan := func(scanRoot string) error { + found, err := listIn(scanRoot, root) + if err != nil { + return err + } + for _, old := range found { + if _, ok := seen[old.RepoKey]; ok { + continue + } + seen[old.RepoKey] = struct{}{} + result = append(result, old) + } + return nil + } + for _, repoType := range []string{"models", "datasets", "spaces"} { + // Namespace case conflicts cross repository types, so the candidate's + // namespace directory is visited under every type as well. + for _, scanRoot := range conflictScanRoots(root, repoType, d.Namespace) { + if err := scan(scanRoot); err != nil { + return result, err + } + } + if d.Namespace != Local && d.Namespace != HuggingFace && d.Namespace != ModelScope { + bare, err := listBareLocalRepositories(root, repoType, d.Namespace) + if err != nil { + return result, err + } + for _, old := range bare { + if _, ok := seen[old.RepoKey]; ok { + continue + } + seen[old.RepoKey] = struct{}{} + result = append(result, old) + } + } + } + return result, nil +} + +// conflictScanRoots lists the directories that can hold a repository whose stored +// path can overlap a candidate of the given repository type and namespace. The +// paths mirror APIRoot/storagePath; a path that does not exist yields nothing. +func conflictScanRoots(root, repoType, namespace string) []string { + switch namespace { + case ModelScope: + // modelscope/: disjoint from everything else. + return []string{filepath.Join(root, "api", repoType, ModelScope)} + case HuggingFace: + // directly below the type directory: disjoint from namespaces. + return []string{filepath.Join(root, "api", repoType)} + case Local: + // dingo-local/: only the shared upload root can hold overlaps. + return []string{filepath.Join(root, "api", repoType, Local)} + default: + // dingo-local//. Siblings in the same namespace can + // overlap the candidate. A bare local repository is stored at + // dingo-local/, so it can also be a prefix of dingo-local/; + // only namespaces without a marker of their own are those bare repositories, + // which is why the second root is read shallowly rather than walked. + return []string{filepath.Join(root, "api", repoType, Local, namespace)} + } +} + +// listBareLocalRepositories returns repositories stored directly below the local +// upload root, i.e. those using the local namespace. They are the only hosted +// repositories outside the candidate's namespace whose path can overlap it. +func listBareLocalRepositories(root, repoType, namespace string) ([]Descriptor, error) { + result := []Descriptor{} + dir := filepath.Join(root, "api", repoType, Local) + entries, err := os.ReadDir(dir) + if err != nil { + if errors.Is(err, os.ErrNotExist) { + return result, nil + } + return result, err + } + for _, entry := range entries { + if !entry.IsDir() { + continue + } + // Case-folded namespace aliases are physically separate on Linux but + // still forbidden by the registry contract. Inspect that matching tree + // only, not every unrelated namespace. + if entry.Name() != namespace && strings.EqualFold(entry.Name(), namespace) { + aliases, err := listIn(filepath.Join(dir, entry.Name()), root) + if err != nil { + return nil, err + } + result = append(result, aliases...) + } + marker := filepath.Join(dir, entry.Name(), Marker) + b, err := os.ReadFile(marker) + if err != nil { + continue // a namespace directory, or a bare tree without a marker + } + var d Descriptor + if err = json.Unmarshal(b, &d); err != nil || d.Validate() != nil { + continue + } + if d.Namespace == Local { + result = append(result, d) + } + } + return result, nil +} + +// registrationConflicts applies the same conflict rules Register enforces, given +// a candidate and the repositories visible to it. It is shared with the +// whole-registry reference used in tests so both paths cannot drift. +func registrationConflicts(root string, d Descriptor, all []Descriptor) error { + for _, old := range all { + if old.RepoKey != d.RepoKey { + for _, pair := range [][2]string{{old.FilesRoot(root), d.FilesRoot(root)}, {old.APIRoot(root), d.APIRoot(root)}} { + a, b := strings.ToLower(filepath.Clean(pair[0])), strings.ToLower(filepath.Clean(pair[1])) + if a == b || strings.HasPrefix(a, b+string(filepath.Separator)) || strings.HasPrefix(b, a+string(filepath.Separator)) { + return fmt.Errorf("repository physical path conflict") + } + } + } + if strings.EqualFold(old.Namespace, d.Namespace) && old.Namespace != d.Namespace { + return fmt.Errorf("namespace case conflict") + } + if old.RepoType != d.RepoType || old.Namespace != d.Namespace { + continue + } + oldParts, newParts := strings.Split(old.Repo, "/"), strings.Split(d.Repo, "/") + for i := 0; i < len(oldParts) && i < len(newParts); i++ { + if !strings.EqualFold(oldParts[i], newParts[i]) { + break + } + if oldParts[i] != newParts[i] { + return fmt.Errorf("repository parent directory case conflict") + } + } + a, b := strings.ToLower(old.Repo), strings.ToLower(d.Repo) + if a == b { + if old.RepoKey == d.RepoKey && old == d { + continue + } + return fmt.Errorf("repository identity or source conflict") + } + if strings.HasPrefix(a, b+"/") || strings.HasPrefix(b, a+"/") { + return fmt.Errorf("repository ancestor/descendant conflict") + } + } + return nil +} + +// Register serializes case and ancestor checks with marker creation. The lock +// file also rejects concurrent registration by another process sharing root. +func Register(root string, d Descriptor) error { + if err := d.Validate(); err != nil { + return err + } + if err := SafePath(root, filepath.Join(d.APIRoot(root), Marker)); err != nil { + return err + } + if err := SafePath(root, filepath.Join(d.FilesRoot(root), ".check")); err != nil { + return err + } + registrationMu.Lock() + defer registrationMu.Unlock() + if err := os.MkdirAll(filepath.Join(root, "api"), 0755); err != nil { + return err + } + lockPath := filepath.Join(root, "api", ".repository-registration.lock") + fileLock := flock.New(lockPath) + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + locked, err := fileLock.TryLockContext(ctx, 10*time.Millisecond) + if err != nil { + return err + } + if !locked { + return fmt.Errorf("repository registration busy") + } + defer fileLock.Unlock() + // An existing descriptor is immutable through Register. Validate it directly + // under both registration locks instead of enumerating unrelated repositories. + if old, readErr := Read(root, d.RepoKey); readErr == nil { + if old == d { + return nil + } + return fmt.Errorf("repository identity or source conflict") + } else if !errors.Is(readErr, os.ErrNotExist) { + return readErr + } + // Conflict detection only ever compares against repositories that share this + // repository type and namespace, plus this descriptor's own ancestors. Those + // are reachable directly from the storage root, so registration cost stays + // independent of how many unrelated repositories a node accumulates. + all, err := listConflicts(root, d) + if err != nil { + return err + } + for _, old := range all { + if old.RepoKey == d.RepoKey && old == d { + return nil // already registered; nothing to do + } + } + if err = registrationConflicts(root, d, all); err != nil { + return err + } + dest := filepath.Join(d.APIRoot(root), Marker) + if err = os.MkdirAll(filepath.Dir(dest), 0755); err != nil { + return err + } + b, err := json.MarshalIndent(d, "", " ") + if err != nil { + return err + } + f, err := os.CreateTemp(filepath.Dir(dest), ".repository-*.tmp") + if err != nil { + return err + } + defer os.Remove(f.Name()) + if _, err = f.Write(b); err != nil { + f.Close() + return err + } + if err = f.Sync(); err != nil { + f.Close() + return err + } + if err = f.Close(); err != nil { + return err + } + return os.Rename(f.Name(), dest) +} + +// Locate finds the descriptor for a file by registered roots, never by guessing +// how many repository path segments precede blobs/resolve/provider-cache. +func Locate(root, file string) (Descriptor, string, error) { + all, err := List(root) + if err != nil { + return Descriptor{}, "", err + } + return LocateRegistered(root, file, all) +} + +// LocateRegistered reuses one validated registry snapshot for batch traversal. +func LocateRegistered(root, file string, all []Descriptor) (Descriptor, string, error) { + for _, d := range all { + rel, err := filepath.Rel(d.FilesRoot(root), file) + if err == nil && rel != "." && rel != ".." && !strings.HasPrefix(rel, ".."+string(filepath.Separator)) && !filepath.IsAbs(rel) { + return d, filepath.ToSlash(rel), nil + } + } + return Descriptor{}, "", os.ErrNotExist +} + +func SafePath(root, target string) error { + rootAbs, err := filepath.Abs(root) + if err != nil { + return err + } + targetAbs, err := filepath.Abs(target) + if err != nil { + return err + } + rel, err := filepath.Rel(rootAbs, targetAbs) + if err != nil { + return err + } + if rel == ".." || strings.HasPrefix(rel, ".."+string(filepath.Separator)) || filepath.IsAbs(rel) { + return fmt.Errorf("path escapes repository root") + } + for current := targetAbs; current != rootAbs; current = filepath.Dir(current) { + info, err := os.Lstat(current) + if os.IsNotExist(err) { + continue + } + if err != nil { + return err + } + if info.Mode()&os.ModeSymlink != 0 { + return fmt.Errorf("symlink in repository path: %s", current) + } + } + return nil +} diff --git a/pkg/repository/repository_test.go b/pkg/repository/repository_test.go new file mode 100644 index 0000000..ad73f13 --- /dev/null +++ b/pkg/repository/repository_test.go @@ -0,0 +1,489 @@ +package repository + +import ( + "encoding/json" + "os" + "path/filepath" + "sync" + "testing" +) + +func TestMultiLevelIdentityAndNamespaceIsolation(t *testing.T) { + root := t.TempDir() + keys := []RepoKey{ + {Namespace: HuggingFace, RepoType: "models", Repo: "Qwen/Qwen3"}, + {Namespace: ModelScope, RepoType: "models", Repo: "Qwen/Qwen3"}, + {Namespace: "alice", RepoType: "models", Repo: "team/model-a"}, + {Namespace: "bob", RepoType: "models", Repo: "team/model-a"}, + } + paths := map[string]bool{} + for _, k := range keys { + parsed, err := ParseID(k.RepoType, k.ID()) + if err != nil || parsed != k { + t.Fatalf("lost identity %v %v", parsed, err) + } + blob := k.Blob(root, "same-content") + if paths[blob] { + t.Fatal("namespaces share a blob path") + } + paths[blob] = true + d := Hosted(k) + if k.Namespace == HuggingFace || k.Namespace == ModelScope { + d = Remote(k) + } + if err = Register(root, d); err != nil { + t.Fatal(err) + } + got, err := Read(root, k) + if err != nil || got != d { + t.Fatalf("descriptor roundtrip: %v %v", got, err) + } + located, rel, err := Locate(root, k.Resolve(root, "commit", "weights/a.bin")) + if err != nil || located != d || rel != "resolve/commit/weights/a.bin" { + t.Fatalf("wrong root: %+v %s %v", located, rel, err) + } + } + all, err := List(root) + if err != nil || len(all) != 4 { + t.Fatalf("enumeration: %d %v", len(all), err) + } +} +func TestInvalidRepositoryLocators(t *testing.T) { + invalid := []RepoKey{ + {Namespace: "", RepoType: "models", Repo: "repo"}, + {Namespace: "alice/bob", RepoType: "models", Repo: "repo"}, + {Namespace: "alice", RepoType: "unknown", Repo: "repo"}, + } + for _, repo := range []string{"", "/repo", "repo/", "a//b", "a/../b", "./a", `a`, "a/nul.txt", "a/name.", "a/\x00bad"} { + invalid = append(invalid, RepoKey{Namespace: "alice", RepoType: "models", Repo: repo}) + } + for _, k := range invalid { + if k.Validate() == nil { + t.Fatalf("accepted %+v", k) + } + } + for _, name := range []string{"team/resolve/model", "tree/revision/model", "Qwen/Qwen3"} { + if err := (RepoKey{Namespace: "alice", RepoType: "models", Repo: name}).Validate(); err != nil { + t.Fatalf("HTTP operation name incorrectly reserved: %s %v", name, err) + } + } +} +func TestRegistrationRejectsCaseAndAncestorConflicts(t *testing.T) { + for _, repo := range []string{"team", "team/model-a/child", "TEAM/model-a"} { + root := t.TempDir() + first := Hosted(RepoKey{Namespace: "alice", RepoType: "models", Repo: "team/model-a"}) + if err := Register(root, first); err != nil { + t.Fatal(err) + } + next := Hosted(RepoKey{Namespace: "alice", RepoType: "models", Repo: repo}) + if err := Register(root, next); err == nil { + t.Fatalf("accepted conflict %s", repo) + } + } + root := t.TempDir() + if err := Register(root, Hosted(RepoKey{Namespace: "Alice", RepoType: "models", Repo: "a"})); err != nil { + t.Fatal(err) + } + if err := Register(root, Hosted(RepoKey{Namespace: "alice", RepoType: "datasets", Repo: "b"})); err == nil { + t.Fatal("namespace case collision crosses repo types") + } + // Segment-aware comparison must allow sibling names. + if err := Register(root, Hosted(RepoKey{Namespace: "Alice", RepoType: "models", Repo: "a-b"})); err != nil { + t.Fatal(err) + } +} +func TestConcurrentIdempotentRegistrationAndConflicts(t *testing.T) { + root := t.TempDir() + d := Hosted(RepoKey{Namespace: "alice", RepoType: "models", Repo: "team/model-a"}) + var wg sync.WaitGroup + errs := make(chan error, 12) + for i := 0; i < 12; i++ { + wg.Add(1) + go func() { defer wg.Done(); errs <- Register(root, d) }() + } + wg.Wait() + close(errs) + for err := range errs { + if err != nil { + t.Fatal(err) + } + } + all, err := List(root) + if err != nil || len(all) != 1 { + t.Fatalf("concurrent registration duplicated: %v %v", all, err) + } + root = t.TempDir() + results := make(chan error, 2) + for _, repo := range []string{"team", "team/model-a"} { + wg.Add(1) + go func(name string) { + defer wg.Done() + results <- Register(root, Hosted(RepoKey{Namespace: "alice", RepoType: "models", Repo: name})) + }(repo) + } + wg.Wait() + close(results) + success := 0 + for err := range results { + if err == nil { + success++ + } + } + if success != 1 { + t.Fatalf("conflicting concurrent registrations succeeded %d times", success) + } +} +func TestInvalidDescriptorDoesNotBecomeRemoteOrDeletable(t *testing.T) { + root := t.TempDir() + k := RepoKey{Namespace: "alice", RepoType: "models", Repo: "model"} + d := Hosted(k) + d.Persistent = false + if Register(root, d) == nil { + t.Fatal("accepted nonpersistent hosted data") + } + d = Remote(k) + if Register(root, d) == nil { + t.Fatal("accepted user namespace as remote") + } + d = Hosted(RepoKey{Namespace: HuggingFace, RepoType: "models", Repo: "Qwen/Qwen3"}) + if Register(root, d) == nil { + t.Fatal("accepted hosted upload into reserved remote namespace") + } + if err := Register(root, Hosted(k)); err != nil { + t.Fatal(err) + } + marker := filepath.Join(k.APIRoot(root), Marker) + if err := os.WriteFile(marker, []byte("{broken"), 0644); err != nil { + t.Fatal(err) + } + if _, err := Read(root, k); err == nil { + t.Fatal("read ignored bad descriptor") + } + if _, err := List(root); err == nil { + t.Fatal("scan ignored bad descriptor") + } +} +func TestMarkerIdentityMustMatchItsLocation(t *testing.T) { + root := t.TempDir() + k := RepoKey{Namespace: "alice", RepoType: "models", Repo: "model"} + if err := Register(root, Hosted(k)); err != nil { + t.Fatal(err) + } + wrong := Hosted(RepoKey{Namespace: "bob", RepoType: "models", Repo: "model"}) + b, _ := json.Marshal(wrong) + if err := os.WriteFile(filepath.Join(k.APIRoot(root), Marker), b, 0644); err != nil { + t.Fatal(err) + } + if _, err := Read(root, k); err == nil { + t.Fatal("accepted marker belonging to another namespace") + } + if _, err := List(root); err == nil { + t.Fatal("scan accepted wrong marker location") + } +} +func TestSafePathAndSymlinkRegistration(t *testing.T) { + root := t.TempDir() + if err := SafePath(root, filepath.Join(root, "..", "escape")); err == nil { + t.Fatal("allowed traversal") + } + if err := SafePath(root, filepath.Join(root, "new", "repo")); err != nil { + t.Fatal(err) + } + outside := t.TempDir() + parent := filepath.Join(root, "api", "models", "dingo-local") + if err := os.MkdirAll(parent, 0755); err != nil { + t.Fatal(err) + } + if err := os.Symlink(outside, filepath.Join(parent, "alice")); err != nil { + t.Skipf("symlink creation unavailable: %v", err) + } + if err := Register(root, Hosted(RepoKey{Namespace: "alice", RepoType: "models", Repo: "model"})); err == nil { + t.Fatal("registered through an escaping directory symlink") + } + entries, err := os.ReadDir(outside) + if err != nil || len(entries) != 0 { + t.Fatal("registration wrote outside storage root") + } +} +func TestNestedMetadataMarkerNameDoesNotCreateAnotherRepository(t *testing.T) { + root := t.TempDir() + k := RepoKey{Namespace: "alice", RepoType: "models", Repo: "team/resolve/model"} + if err := Register(root, Hosted(k)); err != nil { + t.Fatal(err) + } + nested := filepath.Join(k.APIRoot(root), "revision", "commit", Marker) + if err := os.MkdirAll(filepath.Dir(nested), 0755); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(nested, []byte("ordinary nested metadata"), 0644); err != nil { + t.Fatal(err) + } + all, err := List(root) + if err != nil || len(all) != 1 { + t.Fatalf("scanned repository contents as roots: %v %v", all, err) + } +} + +func TestMarkerWordIsAllowedAsRepositoryOrNamespaceDirectory(t *testing.T) { + root := t.TempDir() + keys := []RepoKey{ + {Namespace: "alice", RepoType: "models", Repo: "team/repository.json/model"}, + {Namespace: "alice", RepoType: "models", Repo: "repository.json"}, + {Namespace: "repository.json", RepoType: "models", Repo: "model"}, + } + for _, k := range keys { + if err := Register(root, Hosted(k)); err != nil { + t.Fatalf("register %s: %v", k.ID(), err) + } + if _, err := Read(root, k); err != nil { + t.Fatalf("read %s: %v", k.ID(), err) + } + } + all, err := List(root) + if err != nil || len(all) != len(keys) { + t.Fatalf("marker-named directories broke discovery: %v %v", all, err) + } +} + +// ListHosted must return exactly the hosted subset List returns, so callers that +// only consume hosted descriptors can stop paying for a whole api/ walk. +func TestListHostedMatchesHostedSubsetOfList(t *testing.T) { + root := t.TempDir() + hosted := []RepoKey{ + {Namespace: "alice", RepoType: "models", Repo: "team/model-a"}, + {Namespace: "alice", RepoType: "datasets", Repo: "corpus"}, + {Namespace: "bob", RepoType: "spaces", Repo: "demo"}, + } + for _, k := range hosted { + if err := Register(root, Hosted(k)); err != nil { + t.Fatalf("register hosted %s: %v", k.ID(), err) + } + } + for _, k := range []RepoKey{ + {Namespace: HuggingFace, RepoType: "models", Repo: "Qwen/Qwen3-32B"}, + {Namespace: ModelScope, RepoType: "models", Repo: "Qwen/Qwen3"}, + } { + if err := Register(root, Remote(k)); err != nil { + t.Fatalf("register remote %s: %v", k.ID(), err) + } + } + + all, err := List(root) + if err != nil { + t.Fatal(err) + } + want := map[string]Descriptor{} + for _, d := range all { + if d.Source == "hosted" { + want[d.ID()] = d + } + } + if len(want) != len(hosted) { + t.Fatalf("fixture mismatch: List returned %d hosted of %d total", len(want), len(all)) + } + + got, err := ListHosted(root) + if err != nil { + t.Fatal(err) + } + if len(got) != len(want) { + t.Fatalf("ListHosted returned %d descriptors, want %d: %+v", len(got), len(want), got) + } + for _, d := range got { + expected, ok := want[d.ID()] + if !ok { + t.Fatalf("ListHosted returned non-hosted or unknown descriptor: %+v", d) + } + if d != expected { + t.Fatalf("descriptor mismatch for %s: %+v want %+v", d.ID(), d, expected) + } + } +} + +// Remote cache trees dominate api/ on a heavily used node. A hosted-scoped scan +// must not descend into them at all. +func TestListHostedDoesNotDescendIntoRemoteNamespace(t *testing.T) { + root := t.TempDir() + hostedKey := RepoKey{Namespace: "alice", RepoType: "models", Repo: "team/model-a"} + if err := Register(root, Hosted(hostedKey)); err != nil { + t.Fatal(err) + } + remote := RepoKey{Namespace: HuggingFace, RepoType: "models", Repo: "Qwen/Qwen3-32B"} + if err := Register(root, Remote(remote)); err != nil { + t.Fatal(err) + } + // A marker planted deep inside the remote cache must never be enumerated by + // the hosted scan, even though it is a structurally valid repository root. + decoy := filepath.Join(remote.APIRoot(root), "paths-info", "commit", "nested", Marker) + if err := os.MkdirAll(filepath.Dir(decoy), 0o755); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(decoy, []byte("{}"), 0o644); err != nil { + t.Fatal(err) + } + + got, err := ListHosted(root) + if err != nil { + t.Fatal(err) + } + if len(got) != 1 || got[0].ID() != hostedKey.ID() { + t.Fatalf("hosted scan leaked remote cache entries: %+v", got) + } +} + +// Data subtrees are not repositories. Descending them is what turns discovery +// into a walk of every stored file. +func TestListStopsAtRepositoryDataSubtrees(t *testing.T) { + root := t.TempDir() + k := RepoKey{Namespace: "alice", RepoType: "models", Repo: "team/model-a"} + if err := Register(root, Hosted(k)); err != nil { + t.Fatal(err) + } + for _, name := range []string{"paths-info", "revision", "recycle"} { + decoy := filepath.Join(k.APIRoot(root), name, "commit", "nested", Marker) + if err := os.MkdirAll(filepath.Dir(decoy), 0o755); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(decoy, []byte("{}"), 0o644); err != nil { + t.Fatal(err) + } + } + for _, list := range []struct { + name string + fn func(string) ([]Descriptor, error) + }{{"List", List}, {"ListHosted", ListHosted}} { + all, err := list.fn(root) + if err != nil { + t.Fatalf("%s: %v", list.name, err) + } + if len(all) != 1 || all[0].ID() != k.ID() { + t.Fatalf("%s descended into repository data subtrees: %+v", list.name, all) + } + } +} + +// Empty or absent hosted roots are normal, not an error. +func TestListHostedToleratesMissingRoots(t *testing.T) { + root := t.TempDir() + got, err := ListHosted(root) + if err != nil { + t.Fatalf("missing api/ must not fail: %v", err) + } + if len(got) != 0 { + t.Fatalf("expected no repositories, got %+v", got) + } + k := RepoKey{Namespace: "alice", RepoType: "models", Repo: "only-models"} + if err := Register(root, Hosted(k)); err != nil { + t.Fatal(err) + } + got, err = ListHosted(root) + if err != nil { + t.Fatalf("absent datasets/spaces roots must not fail: %v", err) + } + if len(got) != 1 || got[0].ID() != k.ID() { + t.Fatalf("unexpected hosted enumeration: %+v", got) + } +} + +// Orphan residue is a deleted repository whose marker is gone but whose +// paths-info tree was left behind. It contributes nothing to enumeration. +// +// It also cannot be pruned by directory name alone: a repository may legitimately +// be named "paths-info/repo", and that tree is byte-for-byte indistinguishable +// from orphan residue under a markerless parent. What matters here is that the +// residue is not enumerated as a repository and does not fail the scan. +func TestListDoesNotEnumerateOrphanResidue(t *testing.T) { + root := t.TempDir() + keep := RepoKey{Namespace: "alice", RepoType: "models", Repo: "team/keeper"} + if err := Register(root, Hosted(keep)); err != nil { + t.Fatal(err) + } + // A removed repository: api/models/alice/ghost/ with paths-info but no marker. + orphanRoot := filepath.Join(root, "api", "models", "alice", "ghost") + deep := filepath.Join(orphanRoot, "paths-info", "commit", "weights", "file.bin") + if err := os.MkdirAll(deep, 0o755); err != nil { + t.Fatal(err) + } + // A marker planted at the bottom of the orphan tree is not a valid + // descriptor, so enumeration must either skip it or report the malformed + // marker; it must never materialise a bogus repository. A valid marker there + // would describe a repository whose path does not match, which is rejected + // separately. Here the residue is simply never enumerated as a repository. + for _, list := range []struct { + name string + fn func(string) ([]Descriptor, error) + }{{"List", List}, {"ListHosted", ListHosted}} { + all, _ := list.fn(root) + for _, d := range all { + if d.ID() == "alice/ghost" || d.Repo == "ghost" { + t.Fatalf("%s enumerated orphan residue as a repository: %+v", list.name, d) + } + } + } +} + +// Even with no marker anywhere, a bare data subtree must not yield descriptors. +func TestListDoesNotDescendBareDataSubtree(t *testing.T) { + root := t.TempDir() + orphan := filepath.Join(root, "api", "models", "alice", "ghost") + deep := filepath.Join(orphan, "paths-info", "commit", "weights", "file.bin") + if err := os.MkdirAll(deep, 0o755); err != nil { + t.Fatal(err) + } + all, err := List(root) + if err != nil { + t.Fatal(err) + } + if len(all) != 0 { + t.Fatalf("bare data subtree produced descriptors: %+v", all) + } +} + +// A repository whose own name contains a data-subtree word stays discoverable. +// This is why pruning can never be decided by directory name alone. +func TestListKeepsRepositoryNamedLikeDataSubtree(t *testing.T) { + root := t.TempDir() + for _, repo := range []string{"team/resolve/model", "team/revision/model", "team/recycle/model", "paths-info/repo"} { + k := RepoKey{Namespace: "alice", RepoType: "models", Repo: repo} + if err := Register(root, Hosted(k)); err != nil { + t.Fatalf("register %s: %v", repo, err) + } + } + all, err := List(root) + if err != nil { + t.Fatal(err) + } + if len(all) != 4 { + t.Fatalf("data-subtree-worded repositories were pruned: %+v", all) + } + hosted, err := ListHosted(root) + if err != nil { + t.Fatal(err) + } + if len(hosted) != 4 { + t.Fatalf("ListHosted pruned data-subtree-worded repositories: %+v", hosted) + } +} + +func TestCommonDirectoryCasingCannotAliasAcrossRepositoryRoots(t *testing.T) { + root := t.TempDir() + base := RepoKey{Namespace: "alice", RepoType: "models", Repo: "team/model-a"} + if err := Register(root, Hosted(base)); err != nil { + t.Fatal(err) + } + for _, k := range []RepoKey{ + {Namespace: "ALICE", RepoType: "models", Repo: "team/model-b"}, + {Namespace: "alice", RepoType: "models", Repo: "Team/model-b"}, + } { + if err := Register(root, Hosted(k)); err == nil { + t.Fatalf("accepted casing alias in common directory: %s", k.ID()) + } + } + if err := Register(root, Hosted(RepoKey{Namespace: "alice", RepoType: "models", Repo: "team/model-b"})); err != nil { + t.Fatalf("ordinary sibling rejected: %v", err) + } + all, err := List(root) + if err != nil || len(all) != 2 { + t.Fatalf("directory casing damaged registry: %+v %v", all, err) + } +} diff --git a/pkg/repository/scan_cost_test.go b/pkg/repository/scan_cost_test.go new file mode 100644 index 0000000..e6788de --- /dev/null +++ b/pkg/repository/scan_cost_test.go @@ -0,0 +1,153 @@ +package repository + +import ( + "io/fs" + "os" + "path/filepath" + "strconv" + "testing" +) + +// countVisits measures how many directory entries are opened by a walk, which is +// what actually costs FUSE round trips on a network filesystem. +func countVisits(t *testing.T, scanRoot string) (dirs int, files int) { + t.Helper() + err := filepath.WalkDir(scanRoot, func(p string, e fs.DirEntry, err error) error { + if err != nil { + return nil + } + if e.IsDir() { + dirs++ + } else { + files++ + } + return nil + }) + if err != nil && !os.IsNotExist(err) { + t.Fatal(err) + } + return dirs, files +} + +// buildHeavyNode models a production node: a few hosted repositories plus a large +// remote cache whose repositories each carry a deep per-file paths-info tree. +func buildHeavyNode(t *testing.T, hostedCount, remoteCount, filesPerRepo int) string { + t.Helper() + root := t.TempDir() + for i := 0; i < hostedCount; i++ { + k := RepoKey{Namespace: "datacanvas", RepoType: "models", Repo: "team/hosted-" + strconv.Itoa(i)} + if err := Register(root, Hosted(k)); err != nil { + t.Fatal(err) + } + for j := 0; j < filesPerRepo; j++ { + writeFixture(t, filepath.Join(k.PathsInfo(root, "commit1", "dir/file-"+strconv.Itoa(j)), "paths-info_post.json")) + } + } + for i := 0; i < remoteCount; i++ { + k := RepoKey{Namespace: HuggingFace, RepoType: "models", Repo: "Qwen/cached-" + strconv.Itoa(i)} + if err := Register(root, Remote(k)); err != nil { + t.Fatal(err) + } + for j := 0; j < filesPerRepo; j++ { + writeFixture(t, filepath.Join(k.PathsInfo(root, "commit1", "dir/file-"+strconv.Itoa(j)), "paths-info_post.json")) + } + } + return root +} + +func writeFixture(t *testing.T, path string) { + t.Helper() + if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(path, []byte("{}"), 0o644); err != nil { + t.Fatal(err) + } +} + +// The hosted-scoped scan must return the same descriptors as filtering List, +// while touching dramatically fewer directories. +func TestListHostedIsEquivalentButFarCheaper(t *testing.T) { + const hosted = 3 + const remote = 20 + const filesPerRepo = 25 + root := buildHeavyNode(t, hosted, remote, filesPerRepo) + + // Ground truth: the pre-change behaviour, filtered to hosted. + all, err := List(root) + if err != nil { + t.Fatal(err) + } + want := map[string]Descriptor{} + for _, d := range all { + if d.Source == "hosted" { + want[d.ID()] = d + } + } + if len(want) != hosted { + t.Fatalf("fixture built %d hosted repositories, List saw %d", hosted, len(want)) + } + + got, err := ListHosted(root) + if err != nil { + t.Fatal(err) + } + if len(got) != len(want) { + t.Fatalf("ListHosted returned %d, want %d", len(got), len(want)) + } + for _, d := range got { + if want[d.ID()] != d { + t.Fatalf("descriptor mismatch for %s: %+v", d.ID(), d) + } + } + + // Cost comparison: full api/ walk versus hosted-only walk. + fullDirs, fullFiles := countVisits(t, filepath.Join(root, "api")) + hostedDirs, hostedFiles := countVisits(t, filepath.Join(root, "api", "models", Local)) + + t.Logf("full api/ walk: %d dirs, %d files", fullDirs, fullFiles) + t.Logf("hosted root walk: %d dirs, %d files", hostedDirs, hostedFiles) + + if hostedDirs >= fullDirs { + t.Fatalf("hosted scan visited %d dirs, not fewer than full scan %d", hostedDirs, fullDirs) + } + if hostedFiles >= fullFiles { + t.Fatalf("hosted scan visited %d files, not fewer than full scan %d", hostedFiles, fullFiles) + } +} + +// The whole point of the change: a node that caches heavily must not pay for the +// cache when reporting uploads. Remote repositories must contribute nothing. +func TestListHostedCostIsIndependentOfRemoteCacheSize(t *testing.T) { + small := buildHeavyNode(t, 2, 2, 10) + large := buildHeavyNode(t, 2, 60, 10) + + for _, c := range []struct{ name, root string }{{"small", small}, {"large", large}} { + got, err := ListHosted(c.root) + if err != nil { + t.Fatalf("%s: %v", c.name, err) + } + if len(got) != 2 { + t.Fatalf("%s: hosted enumeration returned %d, want 2", c.name, len(got)) + } + for _, d := range got { + if d.Source != "hosted" { + t.Fatalf("%s: non-hosted descriptor leaked: %+v", c.name, d) + } + } + } + + _, smallFiles := countVisits(t, filepath.Join(small, "api", "models", Local)) + _, largeFiles := countVisits(t, filepath.Join(large, "api", "models", Local)) + if smallFiles != largeFiles { + t.Fatalf("hosted scan cost changed with remote cache size: %d vs %d", smallFiles, largeFiles) + } + + // Meanwhile the full walk grows with the cache, which is the old behaviour. + _, smallFull := countVisits(t, filepath.Join(small, "api")) + _, largeFull := countVisits(t, filepath.Join(large, "api")) + if largeFull <= smallFull { + t.Fatalf("full walk did not grow with cache size: %d vs %d", smallFull, largeFull) + } + t.Logf("full walk files: small=%d large=%d; hosted scan files: %d", smallFull, largeFull, smallFiles) +} \ No newline at end of file diff --git a/pkg/repository/storage.go b/pkg/repository/storage.go new file mode 100644 index 0000000..0d3f602 --- /dev/null +++ b/pkg/repository/storage.go @@ -0,0 +1,47 @@ +package repository + +import ( + "fmt" + "strings" +) + +// FromStorage decodes unchanged remote identities and explicitly encoded hosted ones. +func FromStorage(repoType, org, repo string) (RepoKey, error) { + k := RepoKey{Namespace: "huggingface", RepoType: repoType, Repo: repo} + if strings.HasPrefix(org, "dingo-local/") { + k.Namespace = strings.TrimPrefix(org, "dingo-local/") + } else if strings.HasPrefix(org, "modelscope/") { + k.Namespace = "modelscope" + k.Repo = strings.TrimPrefix(org, "modelscope/") + "/" + repo + } else if org != "" { + k.Repo = org + "/" + repo + } + return k, k.Validate() +} + +// Storage leaves HF owner/repo exactly as old SQL and old peers expect. +func (k RepoKey) Storage() (org, repo string, err error) { + if err = k.Validate(); err != nil { + return + } + repo = k.Repo + switch k.Namespace { + case "huggingface": + parts := strings.SplitN(repo, "/", 2) + if len(parts) == 2 { + org, repo = parts[0], parts[1] + } + case "modelscope": + parts := strings.SplitN(repo, "/", 2) + if len(parts) != 2 { + return "", "", fmt.Errorf("ModelScope needs owner/repo") + } + org, repo = "modelscope/"+parts[0], parts[1] + default: + org = "dingo-local/" + k.Namespace + } + if k.Namespace != "huggingface" && (len(org) > 100 || len(repo) > 100 || len(org)+1+len(repo) > 100) { + err = fmt.Errorf("repository identity exceeds unchanged SQL VARCHAR(100) capacity") + } + return +} diff --git a/pkg/storageprobe/evidence_test.go b/pkg/storageprobe/evidence_test.go new file mode 100644 index 0000000..e4f0671 --- /dev/null +++ b/pkg/storageprobe/evidence_test.go @@ -0,0 +1,25 @@ +package storageprobe + +import ( + "dingospeed/pkg/dependency" + "testing" + "time" +) + +func TestImpactUsesBlockedEvidenceWithoutStartingIO(t *testing.T) { + p := &Probe{} + p.status[1] = status{inFlight: true, timedOut: true} + now := time.Now() + p.monitor.Observe(dependency.MetadataWrite, false, now.Add(-3*time.Minute)) + evidence, blocked := p.Evidence(dependency.MetadataWrite, now) + if !blocked || evidence.State != dependency.Unknown { + t.Fatal("stale blockage lost") + } + r := dependency.AssessStorage(dependency.Snapshot{}, evidence, true, blocked, now) + if r.Assessment != dependency.ImpactProbeFailure { + t.Fatalf("lost ongoing blockage: %+v", r) + } + if _, blocked := p.Evidence(dependency.Scheduler, now); blocked { + t.Fatal("scheduler mixed into storage") + } +} diff --git a/pkg/storageprobe/filesystem.go b/pkg/storageprobe/filesystem.go new file mode 100644 index 0000000..ffa5744 --- /dev/null +++ b/pkg/storageprobe/filesystem.go @@ -0,0 +1,98 @@ +package storageprobe + +import ( + "context" + "errors" + "io" + "os" + "path/filepath" +) + +const Directory = ".dingospeed-probe" +const Sentinel = "dingospeed-storage-probe-v1\n" + +var errSentinel = errors.New("invalid probe sentinel or directory") + +// Filesystem performs no I/O during construction. The dedicated directory and +// sentinel must be provisioned beforehand on the intended filesystem. +func Filesystem(repos string) (Operation, Operation) { + dir := filepath.Join(repos, Directory) + read := func(ctx context.Context) error { + if err := ctx.Err(); err != nil { + return err + } + info, err := os.Lstat(dir) + if err != nil { + return err + } + if !info.IsDir() || info.Mode()&os.ModeSymlink != 0 { + return errSentinel + } + path := filepath.Join(dir, "sentinel") + info, err = os.Lstat(path) + if err != nil { + return err + } + if !info.Mode().IsRegular() { + return errSentinel + } + if err := ctx.Err(); err != nil { + return err + } + f, err := os.Open(path) + if err != nil { + return err + } + data, readErr := io.ReadAll(io.LimitReader(f, int64(len(Sentinel)+1))) + closeErr := f.Close() + if err := errors.Join(readErr, closeErr); err != nil { + return err + } + if string(data) != Sentinel { + return errSentinel + } + return ctx.Err() + } + // Only one write worker uses this closure. A failed cleanup blocks new file + // creation on subsequent attempts, bounding leftover files to one per process. + var pending string + write := func(ctx context.Context) (err error) { + if err = read(ctx); err != nil { + return err + } + if err = ctx.Err(); err != nil { + return err + } + if pending != "" { + if err = os.Remove(pending); err != nil && !errors.Is(err, os.ErrNotExist) { + return err + } + pending = "" + } + f, err := os.CreateTemp(dir, "write-*") + if err != nil { + return err + } + pending = f.Name() + defer func() { + closeErr := f.Close() + removeErr := os.Remove(pending) + if removeErr == nil || errors.Is(removeErr, os.ErrNotExist) { + pending = "" + removeErr = nil + } + err = errors.Join(err, closeErr, removeErr) + }() + if err = ctx.Err(); err != nil { + return err + } + if _, err = f.WriteString(Sentinel); err != nil { + return err + } + if err = ctx.Err(); err != nil { + return err + } + return f.Sync() + } + return read, write +} diff --git a/pkg/storageprobe/linux_integration_test.go b/pkg/storageprobe/linux_integration_test.go new file mode 100644 index 0000000..6cfb471 --- /dev/null +++ b/pkg/storageprobe/linux_integration_test.go @@ -0,0 +1,270 @@ +//go:build linux + +package storageprobe + +import ( + "context" + "errors" + "net" + "os" + "path/filepath" + "strconv" + "strings" + "syscall" + "testing" + "time" + + "dingospeed/pkg/dependency" + "github.com/prometheus/client_golang/prometheus" +) + +// Only the isolated runner supplies these paths; normal go test skips real mounts. +func TestLocalLinuxFilesystemFaults(t *testing.T) { + root := os.Getenv("DINGOSPEED_PROBE_TEST_ROOT") + if root == "" { + t.Skip("run test/storage-probe/run-local.sh in an isolated Linux namespace") + } + if os.Geteuid() == 0 { + t.Fatal("permission tests must run as an unprivileged user") + } + control := os.Getenv("DINGOSPEED_PROBE_TEST_CONTROL") + backing := os.Getenv("DINGOSPEED_PROBE_TEST_BACKING") + if control == "" || backing == "" { + t.Fatal("incomplete test environment") + } + setMode := func(mode string) { + t.Helper() + if err := os.WriteFile(control, []byte(mode), 0600); err != nil { + t.Fatal(err) + } + } + read, write := Filesystem(root) + ctx := context.Background() + t.Run("normal_read_write_sync_cleanup", func(t *testing.T) { + setMode("normal") + for _, op := range []Operation{read, write} { + if err := op(ctx); err != nil { + t.Fatal(err) + } + } + entries, err := os.ReadDir(filepath.Join(backing, Directory)) + if err != nil || len(entries) != 1 { + t.Fatalf("cleanup: %v %v", entries, err) + } + }) + t.Run("readable_but_not_writable", func(t *testing.T) { + dir := filepath.Join(backing, Directory) + if err := os.Chmod(dir, 0500); err != nil { + t.Fatal(err) + } + defer os.Chmod(dir, 0700) + if err := read(ctx); err != nil { + t.Fatalf("read affected by write-only restriction: %v", err) + } + err := write(ctx) + if !errors.Is(err, syscall.EACCES) { + t.Fatalf("expected actual EACCES, got %v", err) + } + p := &Probe{} + p.observe(1, err) + if p.snapshot(1).errors[0] != 1 || p.monitor.Snapshot(dependency.MetadataRead, time.Now()).Unresolved { + t.Fatal("permission classification/scope") + } + t.Logf("uid=%d write errno=%v; read passed", os.Geteuid(), err) + }) + t.Run("read_permission_denied", func(t *testing.T) { + sentinel := filepath.Join(backing, Directory, "sentinel") + if err := os.Chmod(sentinel, 0000); err != nil { + t.Fatal(err) + } + defer os.Chmod(sentinel, 0600) + if err := read(ctx); !errors.Is(err, syscall.EACCES) { + t.Fatalf("expected actual EACCES: %v", err) + } + }) + for _, mode := range []string{"eio-read", "eio-write", "eio-sync"} { + t.Run(mode, func(t *testing.T) { + setMode(mode) + defer setMode("normal") + op := write + if mode == "eio-read" { + op = read + } + err := op(ctx) + if !errors.Is(err, syscall.EIO) { + t.Fatalf("expected actual EIO: %v", err) + } + p := &Probe{} + p.observe(0, err) + if p.snapshot(0).errors[2] != 1 { + t.Fatal("EIO not classified") + } + t.Logf("kernel returned %v", err) + }) + } + t.Run("failed_cleanup_does_not_accumulate_files", func(t *testing.T) { + setMode("eio-unlink") + defer setMode("normal") + for i := 0; i < 5; i++ { + if err := write(ctx); !errors.Is(err, syscall.EIO) { + t.Fatalf("cleanup error lost: %v", err) + } + entries, err := os.ReadDir(filepath.Join(backing, Directory)) + if err != nil || len(entries) != 2 { + t.Fatalf("temporary files accumulated: %v %v", entries, err) + } + } + setMode("normal") + if err := write(ctx); err != nil { + t.Fatalf("retry cleanup failed: %v", err) + } + entries, _ := os.ReadDir(filepath.Join(backing, Directory)) + if len(entries) != 1 { + t.Fatal("recovery left temporary files") + } + }) + t.Run("blocked_write_real_syscall_is_bounded", func(t *testing.T) { + setMode("block-write") + defer setMode("normal") + p, err := New(read, write, 10*time.Second, 3*time.Second) + if err != nil { + t.Fatal(err) + } + p.Start(ctx) + defer p.Stop() + registry := prometheus.NewPedanticRegistry() + registry.MustRegister(Collector(p)) + deadline := time.Now().Add(6 * time.Second) + for s := p.snapshot(1); !s.inFlight || !s.timedOut; s = p.snapshot(1) { + if time.Now().After(deadline) { + t.Fatal("default probe timeout was not observed") + } + time.Sleep(10 * time.Millisecond) + } + before := p.snapshot(1) + // Observe several would-be retry intervals while the kernel call remains pending. + for i := 0; i < 32; i++ { + families, err := registry.Gather() + if err != nil || len(families) != 4 { + t.Fatalf("metrics blocked: %v", err) + } + entries, err := os.ReadDir(filepath.Join(backing, Directory)) + if err != nil || len(entries) != 2 { + t.Fatalf("blocked write tasks accumulated temporary files: %v %v", entries, err) + } + time.Sleep(time.Second) + } + after := p.snapshot(1) + if after != before || after.errors[6] != 1 { + t.Fatalf("blocked work retried: before=%+v after=%+v", before, after) + } + if s := p.monitor.Snapshot(dependency.MetadataRead, time.Now()); s.State != dependency.ObservedHealthy { + t.Fatalf("write blockage affected reads: %+v", s) + } + start := time.Now() + p.Stop() + stopDuration := time.Since(start) + if stopDuration > time.Second { + t.Fatal("shutdown blocked on filesystem") + } + failed := p.monitor.Snapshot(dependency.MetadataWrite, time.Now()) + setMode("normal") + await(t, func() bool { return !p.snapshot(1).inFlight }) + if got := p.monitor.Snapshot(dependency.MetadataWrite, time.Now()); got != failed { + t.Fatalf("late success altered evidence: %+v -> %+v", failed, got) + } + t.Logf("default 3s timeout, 10s interval: blocked for >35s, one timeout and one temporary file; read and metrics stayed available; stop=%s", stopDuration) + }) + t.Run("real_time_positive_recovery", func(t *testing.T) { + setMode("eio-write") + p, err := New(read, write, 5*time.Second, 3*time.Second) + if err != nil { + t.Fatal(err) + } + p.Start(ctx) + defer p.Stop() + defer setMode("normal") + deadline := time.Now().Add(16 * time.Second) + for p.monitor.Snapshot(dependency.MetadataWrite, time.Now()).State != dependency.Degraded { + if time.Now().After(deadline) { + t.Fatal("ongoing EIO never degraded") + } + time.Sleep(50 * time.Millisecond) + } + setMode("normal") + deadline = time.Now().Add(17 * time.Second) + for p.monitor.Snapshot(dependency.MetadataWrite, time.Now()).Unresolved { + if time.Now().After(deadline) { + t.Fatal("positive evidence never recovered") + } + time.Sleep(50 * time.Millisecond) + } + if s := p.monitor.Snapshot(dependency.MetadataWrite, time.Now()); s.State != dependency.ObservedHealthy { + t.Fatalf("unexpected recovery %+v", s) + } + }) + // Last: a separate root controller kills only this disposable FUSE daemon. + t.Run("host_remount_healthy_but_old_namespace_ENOTCONN", func(t *testing.T) { + request := os.Getenv("DINGOSPEED_PROBE_TEST_DISCONNECT") + if request == "" { + t.Fatal("missing disconnect controller") + } + if err := os.WriteFile(request, []byte("disconnect"), 0600); err != nil { + t.Fatal(err) + } + await(t, func() bool { + data, err := os.ReadFile(request + ".healthy") + return err == nil && string(data) == Sentinel + }) + deadline := time.Now().Add(3 * time.Second) + var err error + for time.Now().Before(deadline) { + err = read(ctx) + if errors.Is(err, syscall.ENOTCONN) { + break + } + time.Sleep(10 * time.Millisecond) + } + if !errors.Is(err, syscall.ENOTCONN) { + t.Fatalf("expected real ENOTCONN after daemon exit: %v", err) + } + p := &Probe{} + p.observe(0, err) + if p.snapshot(0).errors[1] != 1 { + t.Fatal("disconnected mount not classified") + } + if writeErr := write(ctx); !errors.Is(writeErr, syscall.ENOTCONN) { + t.Fatalf("old namespace write unexpectedly usable: %v", writeErr) + } + t.Logf("parent namespace remount read succeeded as uid 65534; old namespace returned %v", err) + }) +} + +func TestLocalNetworkIsolation(t *testing.T) { + if os.Getenv("DINGOSPEED_PROBE_TEST_ROOT") == "" { + t.Skip("local runner only") + } + data, err := os.ReadFile("/proc/net/route") + if err != nil { + t.Fatal(err) + } + fields := strings.Fields(string(data)) + if len(fields) != 0 && (len(fields) != 11 || fields[0] != "Iface") { + t.Fatalf("unexpected IPv4 route table: %s", data) + } + entries, err := net.Interfaces() + if err != nil { + t.Fatal(err) + } + for _, entry := range entries { + if entry.Name != "lo" || entry.Flags&net.FlagUp != 0 { + t.Fatalf("external or active interface present: %s", entry.Name) + } + } + if fd, err := syscall.Socket(syscall.AF_INET, syscall.SOCK_DGRAM, 0); err != nil { + t.Fatal(err) + } else { + syscall.Close(fd) + } + t.Log("no external interfaces or IPv4 routes; unprivileged test uid=" + strconv.Itoa(os.Geteuid())) +} diff --git a/pkg/storageprobe/metrics.go b/pkg/storageprobe/metrics.go new file mode 100644 index 0000000..aefe469 --- /dev/null +++ b/pkg/storageprobe/metrics.go @@ -0,0 +1,47 @@ +package storageprobe + +import ( + "dingospeed/pkg/dependency" + "github.com/prometheus/client_golang/prometheus" + "time" +) + +type collector struct { + p *Probe + state, unresolved, blocked, failures *prometheus.Desc +} + +func Collector(p *Probe) prometheus.Collector { + labels := []string{"operation"} + return &collector{p, + prometheus.NewDesc("dingospeed_storage_probe_state", "Dedicated probe evidence; same state values as dependency_state; not business readiness.", labels, nil), + prometheus.NewDesc("dingospeed_storage_probe_unresolved", "Probe failure lacking positive recovery evidence.", labels, nil), + prometheus.NewDesc("dingospeed_storage_probe_blocked", "Timed-out probe operation still in flight; no replacement is launched.", labels, nil), + prometheus.NewDesc("dingospeed_storage_probe_errors_total", "Failed probe attempts, counted once, with bounded classification.", []string{"operation", "kind"}, nil), + } +} +func (c *collector) Describe(ch chan<- *prometheus.Desc) { + ch <- c.state + ch <- c.unresolved + ch <- c.blocked + ch <- c.failures +} +func (c *collector) Collect(ch chan<- prometheus.Metric) { + for id, op := range []string{"read", "write"} { + s := c.p.monitor.Snapshot(dependency.ID(id), time.Now()) + status := c.p.snapshot(id) + u, b := 0., 0. + if s.Unresolved { + u = 1 + } + if status.inFlight && status.timedOut { + b = 1 + } + ch <- prometheus.MustNewConstMetric(c.state, prometheus.GaugeValue, float64(s.State), op) + ch <- prometheus.MustNewConstMetric(c.unresolved, prometheus.GaugeValue, u, op) + ch <- prometheus.MustNewConstMetric(c.blocked, prometheus.GaugeValue, b, op) + for i, kind := range kinds { + ch <- prometheus.MustNewConstMetric(c.failures, prometheus.CounterValue, float64(status.errors[i]), op, kind) + } + } +} diff --git a/pkg/storageprobe/probe.go b/pkg/storageprobe/probe.go new file mode 100644 index 0000000..f49e005 --- /dev/null +++ b/pkg/storageprobe/probe.go @@ -0,0 +1,165 @@ +// Package storageprobe gathers bounded, independent storage evidence. +// Probe results never change request handling or clear business observations. +package storageprobe + +import ( + "context" + "errors" + "os" + "sync" + "time" + + "dingospeed/pkg/dependency" + "dingospeed/pkg/util" +) + +type Operation func(context.Context) error + +var kinds = [...]string{string(util.FileAccessPermissionDenied), string(util.FileAccessMountDisconnected), string(util.FileAccessIOFailure), string(util.FileAccessUnavailable), "not_found", "invalid_sentinel", "timeout"} + +type status struct { + inFlight, timedOut bool + errors [len(kinds)]uint64 +} + +type Probe struct { + monitor dependency.Monitor + mu sync.Mutex + status [2]status + once sync.Once + workers sync.WaitGroup // Controllers only; blocked filesystem calls are not awaited. + cancel context.CancelFunc + stopped bool + interval, timeout time.Duration + ops [2]Operation +} + +func New(read, write Operation, interval, timeout time.Duration) (*Probe, error) { + if read == nil || write == nil || timeout <= 0 || interval <= timeout { + return nil, errors.New("probe requires two operations and interval > timeout > 0") + } + return &Probe{interval: interval, timeout: timeout, ops: [2]Operation{read, write}}, nil +} + +// Start is one-shot: restarting the controllers cannot replace a stuck worker. +func (p *Probe) Start(ctx context.Context) { + p.once.Do(func() { + ctx, cancel := context.WithCancel(ctx) + p.mu.Lock() + if p.stopped { + p.mu.Unlock() + cancel() + return + } + p.cancel = cancel + p.workers.Add(2) + p.mu.Unlock() + for i := range p.ops { + go p.run(ctx, i) + } + }) +} + +// Stop stops scheduling immediately; it cannot interrupt an OS filesystem call. +// An already blocked call retains its slot until it returns or the process exits. +func (p *Probe) Stop() { + p.mu.Lock() + p.stopped = true + if p.cancel != nil { + p.cancel() + } + p.mu.Unlock() + p.workers.Wait() +} + +func (p *Probe) run(ctx context.Context, id int) { + defer p.workers.Done() + for ctx.Err() == nil { + opCtx, cancel := context.WithTimeout(ctx, p.timeout) + result := make(chan error, 1) + p.mu.Lock() + p.status[id].inFlight = true + p.status[id].timedOut = false + p.mu.Unlock() + go func() { + err := p.ops[id](opCtx) + p.mu.Lock() + p.status[id].inFlight = false + p.mu.Unlock() + result <- err + }() + select { + case err := <-result: + if ctx.Err() == nil { + if opCtx.Err() != nil { + err = opCtx.Err() + } + p.observe(id, err) + } + case <-opCtx.Done(): + if ctx.Err() != nil { + cancel() + return + } + p.observe(id, context.DeadlineExceeded) + // Never launch a replacement while the timed-out operation is alive. + select { + case <-result: // A late result cannot count as recovery evidence. + case <-ctx.Done(): + cancel() + return + } + } + cancel() + timer := time.NewTimer(p.interval) + select { + case <-ctx.Done(): + timer.Stop() + return + case <-timer.C: + } + } +} + +func (p *Probe) observe(id int, err error) { + if err != nil { + kind := string(util.FileAccessUnavailable) + switch { + case errors.Is(err, context.DeadlineExceeded): + kind = "timeout" + case errors.Is(err, errSentinel): + kind = "invalid_sentinel" + case errors.Is(err, os.ErrNotExist): + kind = "not_found" + default: + if classified, ok := util.ClassifyFileAccessError("", err); ok { + kind = string(classified.Kind) + } + } + p.mu.Lock() + p.status[id].timedOut = kind == "timeout" + for i, k := range kinds { + if k == kind { + p.status[id].errors[i]++ + break + } + } + p.mu.Unlock() + } + p.monitor.Observe(dependency.ID(id), err == nil, time.Now()) +} + +func (p *Probe) snapshot(id int) status { + p.mu.Lock() + defer p.mu.Unlock() + return p.status[id] +} + +// Evidence reads existing probe observations; it never starts filesystem I/O. +func (p *Probe) Evidence(id dependency.ID, now time.Time) (dependency.Snapshot, bool) { + if id != dependency.MetadataRead && id != dependency.MetadataWrite { + return dependency.Snapshot{}, false + } + s := p.snapshot(int(id)) + return p.monitor.Snapshot(id, now), s.inFlight && s.timedOut +} diff --git a/pkg/storageprobe/probe_test.go b/pkg/storageprobe/probe_test.go new file mode 100644 index 0000000..dbd13da --- /dev/null +++ b/pkg/storageprobe/probe_test.go @@ -0,0 +1,184 @@ +package storageprobe + +import ( + "context" + "errors" + "os" + "path/filepath" + "runtime" + "sync" + "sync/atomic" + "syscall" + "testing" + "time" + + "dingospeed/pkg/dependency" + "github.com/prometheus/client_golang/prometheus" +) + +func await(t *testing.T, condition func() bool) { + t.Helper() + deadline := time.Now().Add(3 * time.Second) + for !condition() { + if time.Now().After(deadline) { + t.Fatal("condition timed out") + } + time.Sleep(time.Millisecond) + } +} + +func TestBlockedCallHasOneSlotAndStopDoesNotWaitForIO(t *testing.T) { + var reads, writes atomic.Int32 + release := make(chan struct{}) + var releaseOnce sync.Once + unblock := func() { releaseOnce.Do(func() { close(release) }) } + defer unblock() + p, err := New(func(context.Context) error { reads.Add(1); <-release; return nil }, func(context.Context) error { writes.Add(1); return nil }, 30*time.Millisecond, 10*time.Millisecond) + if err != nil { + t.Fatal(err) + } + p.Start(context.Background()) + defer p.Stop() + registry := prometheus.NewPedanticRegistry() + registry.MustRegister(Collector(p)) + await(t, func() bool { return writes.Load() >= 4 }) + for i := 0; i < 20; i++ { + if _, err := registry.Gather(); err != nil { + t.Fatal(err) + } + } + s := p.snapshot(0) + if reads.Load() != 1 || !s.inFlight || !s.timedOut || s.errors[len(kinds)-1] != 1 { + t.Fatalf("blocked operation replaced or counted repeatedly: reads=%d status=%+v", reads.Load(), s) + } + if got := p.monitor.Snapshot(dependency.MetadataRead, time.Now()); !got.Unresolved { + t.Fatal("timeout unobserved") + } + if got := p.monitor.Snapshot(dependency.MetadataWrite, time.Now()); got.Unresolved || got.State != dependency.ObservedHealthy { + t.Fatal("read blockage affected write evidence") + } + stopped := make(chan struct{}) + go func() { p.Stop(); close(stopped) }() + select { + case <-stopped: + case <-time.After(time.Second): + t.Fatal("Stop waited for blocked I/O") + } + unblock() + await(t, func() bool { return !p.snapshot(0).inFlight }) + if !p.monitor.Snapshot(dependency.MetadataRead, time.Now()).Unresolved { + t.Fatal("late result cleared timeout") + } + p.Start(context.Background()) + if reads.Load() != 1 { + t.Fatal("restarted stopped controller") + } +} + +func TestStopBeforeStart(t *testing.T) { + op := func(context.Context) error { t.Error("operation ran after stop"); return nil } + p, _ := New(op, op, time.Second, time.Millisecond) + p.Stop() + p.Start(context.Background()) + p.Stop() +} + +func TestProbeEvidenceCannotClearBusinessFailure(t *testing.T) { + now := time.Now() + dependency.Default.Observe(dependency.MetadataRead, false, now) + before := dependency.Default.Snapshot(dependency.MetadataRead, now) + p := &Probe{} + for i := 0; i < 3; i++ { + p.monitor.Observe(dependency.MetadataRead, true, now.Add(time.Duration(i)*5*time.Second)) + } + if got := dependency.Default.Snapshot(dependency.MetadataRead, now); got != before { + t.Fatal("probe mutated business evidence") + } +} + +func TestFaultClassification(t *testing.T) { + cases := []struct { + name string + err error + index int + }{ + {"permission", &os.PathError{Op: "open", Path: "private", Err: os.ErrPermission}, 0}, + {"missing", &os.PathError{Op: "open", Path: "sentinel", Err: os.ErrNotExist}, 4}, + {"invalid", errSentinel, 5}, + {"deadline", context.DeadlineExceeded, 6}, + } + if runtime.GOOS == "linux" { + cases = append(cases, struct { + name string + err error + index int + }{"io", &os.PathError{Op: "read", Err: syscall.Errno(5)}, 2}, struct { + name string + err error + index int + }{"disconnected", &os.PathError{Op: "read", Err: syscall.Errno(107)}, 1}) + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + p := &Probe{} + p.observe(0, tc.err) + if p.snapshot(0).errors[tc.index] != 1 || !p.monitor.Snapshot(dependency.MetadataRead, time.Now()).Unresolved { + t.Fatal("fault not classified") + } + }) + } +} + +func provision(t *testing.T) string { + t.Helper() + root := t.TempDir() + dir := filepath.Join(root, Directory) + if err := os.Mkdir(dir, 0700); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(dir, "sentinel"), []byte(Sentinel), 0600); err != nil { + t.Fatal(err) + } + return root +} + +func TestFilesystemChecksAndCleanup(t *testing.T) { + root := provision(t) + read, write := Filesystem(root) + if err := read(context.Background()); err != nil { + t.Fatal(err) + } + if err := write(context.Background()); err != nil { + t.Fatal(err) + } + entries, err := os.ReadDir(filepath.Join(root, Directory)) + if err != nil || len(entries) != 1 || entries[0].Name() != "sentinel" { + t.Fatalf("probe file leaked: %v %v", entries, err) + } + if err := os.WriteFile(filepath.Join(root, Directory, "sentinel"), []byte("wrong"), 0600); err != nil { + t.Fatal(err) + } + if err := write(context.Background()); !errors.Is(err, errSentinel) { + t.Fatalf("invalid sentinel accepted: %v", err) + } + empty := t.TempDir() + _, write = Filesystem(empty) + if err := write(context.Background()); !errors.Is(err, os.ErrNotExist) { + t.Fatalf("missing probe root accepted: %v", err) + } + entries, err = os.ReadDir(empty) + if err != nil || len(entries) != 0 { + t.Fatal("probe created absent root") + } +} + +func TestCancelledProbeDoesNoIO(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + read, write := Filesystem(filepath.Join(t.TempDir(), "absent")) + for _, op := range []Operation{read, write} { + if err := op(ctx); !errors.Is(err, context.Canceled) { + t.Fatalf("ignored cancellation: %v", err) + } + } +} diff --git a/pkg/storageprobe/promtool_test.go b/pkg/storageprobe/promtool_test.go new file mode 100644 index 0000000..065e36e --- /dev/null +++ b/pkg/storageprobe/promtool_test.go @@ -0,0 +1,21 @@ +package storageprobe + +import ( + "os" + "os/exec" + "testing" +) + +func TestProbeAlertRules(t *testing.T) { + tool := os.Getenv("PROMTOOL") + if tool == "" { + t.Skip("set PROMTOOL to validate alert timing with Prometheus") + } + for _, args := range [][]string{{"check", "rules", "../../config/prometheus/storage_probe_alerts.yml"}, {"test", "rules", "../../config/prometheus/tests/storage_probe.test.yml"}} { + out, err := exec.Command(tool, args...).CombinedOutput() + if err != nil { + t.Fatalf("promtool: %v\n%s", err, out) + } + t.Log(string(out)) + } +} diff --git a/pkg/transferlimit/limit.go b/pkg/transferlimit/limit.go new file mode 100644 index 0000000..e828ec9 --- /dev/null +++ b/pkg/transferlimit/limit.go @@ -0,0 +1,79 @@ +// Package transferlimit bounds active HTTP transfers until their bodies close. +package transferlimit + +import ( + "context" + "io" + "net/http" + "sync" +) + +type Gate struct { + mu sync.Mutex + active int + changed chan struct{} +} + +func (g *Gate) Acquire(ctx context.Context, limit func() int) (func(), error) { + for { + if err := ctx.Err(); err != nil { + return nil, err + } + g.mu.Lock() + if g.changed == nil { + g.changed = make(chan struct{}) + } + if g.active < max(1, limit()) { + g.active++ + g.mu.Unlock() + var once sync.Once + return func() { + once.Do(func() { g.mu.Lock(); g.active--; close(g.changed); g.changed = make(chan struct{}); g.mu.Unlock() }) + }, nil + } + changed := g.changed + g.mu.Unlock() + select { + case <-ctx.Done(): + return nil, ctx.Err() + case <-changed: + } + } +} + +func (g *Gate) Wake() { + g.mu.Lock() + if g.changed != nil { + close(g.changed) + } + g.changed = make(chan struct{}) + g.mu.Unlock() +} + +type body struct { + io.ReadCloser + release func() +} + +func (b *body) Close() error { defer b.release(); return b.ReadCloser.Close() } +func (b *body) Read(p []byte) (int, error) { + n, err := b.ReadCloser.Read(p) + if err == io.EOF { + b.release() + } + return n, err +} + +func (g *Gate) Do(client *http.Client, req *http.Request, limit func() int) (*http.Response, error) { + release, err := g.Acquire(req.Context(), limit) + if err != nil { + return nil, err + } + resp, err := client.Do(req) + if err != nil { + release() + return resp, err + } + resp.Body = &body{resp.Body, release} + return resp, nil +} diff --git a/pkg/transfersettings/settings.go b/pkg/transfersettings/settings.go new file mode 100644 index 0000000..f4719ad --- /dev/null +++ b/pkg/transfersettings/settings.go @@ -0,0 +1,122 @@ +package transfersettings + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "os" + "path/filepath" + "sync" + "sync/atomic" + + "dingospeed/pkg/config" + "dingospeed/pkg/transferlimit" +) + +type Settings struct { + Upload int `json:"upload"` + HuggingFace int `json:"huggingface"` + ModelScope int `json:"modelscope"` + Peer int `json:"peer"` + Download int `json:"download"` +} + +var initial sync.Once +var value atomic.Pointer[Settings] +var writeMu sync.Mutex +var loadErr error +var gates [5]transferlimit.Gate + +func AcquireDownload(ctx context.Context) (func(), error) { + return gates[4].Acquire(ctx, func() int { return Current().Download }) +} + +func AcquireUpload(ctx context.Context) (func(), error) { + return gates[0].Acquire(ctx, func() int { return Current().Upload }) +} + +func settingsPath() string { return filepath.Join(config.SysConfig.Repos(), ".transfer-settings.json") } +func Current() Settings { + initial.Do(func() { + s := Settings{Upload: max(1, min(config.SysConfig.Upload.ConcurrentLimit, config.SysConfig.GetUploadChunkConcurrentLimit())), HuggingFace: 8, ModelScope: 8, Peer: 8, Download: 8} + b, err := os.ReadFile(settingsPath()) + if err == nil { + err = json.Unmarshal(b, &s) + if err == nil { + err = Validate(s) + } + } + if err != nil && !os.IsNotExist(err) { + loadErr = err + s = Settings{1, 1, 1, 1, 1} + } + value.Store(&s) + }) + return *value.Load() +} +func Validate(s Settings) error { + for _, n := range []int{s.Upload, s.HuggingFace, s.ModelScope, s.Peer, s.Download} { + if n < 1 || n > 64 { + return fmt.Errorf("HTTP concurrency must be between 1 and 64") + } + } + if s.Upload > 8 { + return fmt.Errorf("upload concurrency must be between 1 and 8 (chunk memory budget)") + } + return nil +} +func LoadError() error { Current(); return loadErr } +func Save(s Settings) error { + if err := Validate(s); err != nil { + return err + } + Current() + writeMu.Lock() + defer writeMu.Unlock() + b, err := json.MarshalIndent(s, "", " ") + if err != nil { + return err + } + f, err := os.CreateTemp(filepath.Dir(settingsPath()), ".transfer-settings-*") + if err != nil { + return err + } + name := f.Name() + defer os.Remove(name) + if _, err = f.Write(b); err == nil { + err = f.Sync() + } + closeErr := f.Close() + if err == nil { + err = closeErr + } + if err != nil { + return err + } + if err = os.Rename(name, settingsPath()); err != nil { + return err + } + value.Store(&s) + for i := range gates { + gates[i].Wake() + } + return nil +} +func Do(kind string, client *http.Client, req *http.Request) (*http.Response, error) { + index := 0 + switch kind { + case "huggingface": + index = 1 + case "modelscope": + index = 2 + case "peer": + index = 3 + case "download": + index = 4 + } + return gates[index].Do(client, req, func() int { + s := Current() + return []int{s.Upload, s.HuggingFace, s.ModelScope, s.Peer, s.Download}[index] + }) +} diff --git a/pkg/util/download_retry.go b/pkg/util/download_retry.go new file mode 100644 index 0000000..752252d --- /dev/null +++ b/pkg/util/download_retry.go @@ -0,0 +1,95 @@ +package util + +import ( + "context" + "errors" + "io" + "net" + "net/http" + "os" + "syscall" + "time" + + "dingospeed/pkg/config" + myerr "dingospeed/pkg/error" + "github.com/avast/retry-go" +) + +// RetryDownloadContext applies only to idempotent ranged reads. The caller must +// advance its Range after partial delivery before the next attempt. +func RetryDownloadContext(ctx context.Context, read func() error) error { + attempts := max(config.SysConfig.Retry.Attempts, 1) + return retry.Do(func() error { + if err := ctx.Err(); err != nil { + return retry.Unrecoverable(err) + } + err := read() + if ctx.Err() != nil { + return retry.Unrecoverable(ctx.Err()) + } + if err != nil && !retryableDownloadError(err) { + return retry.Unrecoverable(err) + } + return err + }, retry.Attempts(attempts), + retry.Delay(time.Duration(config.SysConfig.Retry.Delay)*time.Second), + retry.DelayType(retry.BackOffDelay), retry.MaxDelay(30*time.Second), + retry.Context(ctx), retry.LastErrorOnly(true)) +} + +func retryableDownloadError(err error) bool { + // A child HTTP client's timeout is retryable while the task context lives. + // Explicit cancellation and local filesystem failures are never retried. + if errors.Is(err, context.Canceled) { + return false + } + var pathErr *os.PathError + if errors.As(err, &pathErr) { + return false + } + var status myerr.Error + if errors.As(err, &status) { + switch status.StatusCode() { + case http.StatusRequestTimeout, http.StatusTooManyRequests, http.StatusInternalServerError, + http.StatusBadGateway, http.StatusServiceUnavailable, http.StatusGatewayTimeout: + return true + default: + return false + } + } + if errors.Is(err, context.DeadlineExceeded) || errors.Is(err, io.ErrUnexpectedEOF) || + errors.Is(err, syscall.ECONNRESET) || errors.Is(err, syscall.ECONNREFUSED) || errors.Is(err, syscall.EPIPE) { + return true + } + var network net.Error + return errors.As(err, &network) && (network.Timeout() || network.Temporary()) +} + +// CacheFailureReason is safe for an HTTP trailer and for display in a task. +// Raw errors can contain provider URLs, credentials or local storage paths. +func CacheFailureReason(err error) string { + switch { + case errors.Is(err, context.DeadlineExceeded): + return "download timed out (context deadline exceeded)" + case errors.Is(err, context.Canceled): + return "download canceled" + case errors.Is(err, io.ErrUnexpectedEOF): + return "upstream response ended before the file was complete" + case errors.Is(err, syscall.ENOSPC): + return "insufficient cache disk space" + case errors.Is(err, os.ErrPermission): + return "cache storage permission denied" + } + var status myerr.Error + if errors.As(err, &status) && status.StatusCode() >= 400 { + return "upstream returned HTTP " + Itoa(status.StatusCode()) + } + var network net.Error + if errors.As(err, &network) { + if network.Timeout() { + return "download timed out" + } + return "upstream network request failed" + } + return "cache write or transfer failed; see node logs for details" +} diff --git a/pkg/util/download_retry_test.go b/pkg/util/download_retry_test.go new file mode 100644 index 0000000..cf08e30 --- /dev/null +++ b/pkg/util/download_retry_test.go @@ -0,0 +1,128 @@ +package util + +import ( + "context" + "errors" + "fmt" + "io" + "net" + "net/http" + "net/http/httptest" + "os" + "strings" + "syscall" + "testing" + "time" + + "dingospeed/pkg/config" + myerr "dingospeed/pkg/error" +) + +func TestDownloadRetryPolicy(t *testing.T) { + old := config.SysConfig + t.Cleanup(func() { config.SysConfig = old }) + config.SysConfig = &config.Config{Retry: config.Retry{Attempts: 3}} + for _, tc := range []struct { + name string + err error + calls int + }{ + {"http-timeout", fmt.Errorf("body: %w", context.DeadlineExceeded), 3}, + {"short-body", io.ErrUnexpectedEOF, 3}, + {"connection-reset", &net.OpError{Op: "read", Net: "tcp", Err: syscall.ECONNRESET}, 3}, + {"temporary-http", myerr.NewAppendCode(503, "unavailable"), 3}, + {"rate-limit", myerr.NewAppendCode(429, "busy"), 3}, + {"unauthorized", myerr.NewAppendCode(401, "unauthorized"), 1}, + {"forbidden", myerr.NewAppendCode(403, "forbidden"), 1}, + {"missing-file", myerr.NewAppendCode(404, "not found"), 1}, + {"disk-full", &os.PathError{Op: "write", Path: "cache", Err: syscall.ENOSPC}, 1}, + {"permission", os.ErrPermission, 1}, + {"canceled", context.Canceled, 1}, + {"invalid-range", errors.New("source ignored Range request"), 1}, + } { + t.Run(tc.name, func(t *testing.T) { + calls := 0 + err := RetryDownloadContext(context.Background(), func() error { calls++; return tc.err }) + if !errors.Is(err, tc.err) || calls != tc.calls { + t.Fatalf("calls=%d want=%d error=%v", calls, tc.calls, err) + } + }) + } + config.SysConfig.Retry.Attempts = 0 + calls := 0 + if err := RetryDownloadContext(context.Background(), func() error { calls++; return io.ErrUnexpectedEOF }); err == nil || calls != 1 { + t.Fatalf("zero configuration: calls=%d error=%v", calls, err) + } +} + +func TestDownloadRetryHTTPClientTimeoutWithLiveParent(t *testing.T) { + old := config.SysConfig + t.Cleanup(func() { config.SysConfig = old }) + config.SysConfig = &config.Config{Retry: config.Retry{Attempts: 2}} + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Length", "2") + w.WriteHeader(200) + w.(http.Flusher).Flush() + <-r.Context().Done() + })) + defer server.Close() + client := &http.Client{Timeout: 100 * time.Millisecond} + calls := 0 + err := RetryDownloadContext(context.Background(), func() error { + calls++ + resp, err := client.Get(server.URL) + if err != nil { + return err + } + defer resp.Body.Close() + _, err = io.Copy(io.Discard, resp.Body) + return err + }) + if !errors.Is(err, context.DeadlineExceeded) || calls != 2 { + t.Fatalf("client timeout did not retry: calls=%d error=%v", calls, err) + } +} + +func TestDownloadRetryStopsDuringBackoff(t *testing.T) { + old := config.SysConfig + t.Cleanup(func() { config.SysConfig = old }) + config.SysConfig = &config.Config{Retry: config.Retry{Attempts: 3, Delay: 30}} + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + first := make(chan struct{}) + done := make(chan error, 1) + calls := 0 + go func() { + done <- RetryDownloadContext(ctx, func() error { + calls++ + if calls == 1 { + close(first) + } + return context.DeadlineExceeded + }) + }() + <-first + cancel() + select { + case err := <-done: + if !errors.Is(err, context.Canceled) || calls != 1 { + t.Fatalf("calls=%d error=%v", calls, err) + } + case <-time.After(time.Second): + t.Fatal("cancellation waited for backoff") + } + ctx, cancel = context.WithDeadline(context.Background(), time.Now().Add(-time.Second)) + defer cancel() + err := RetryDownloadContext(ctx, func() error { t.Fatal("expired parent made a request"); return nil }) + if !errors.Is(err, context.DeadlineExceeded) { + t.Fatal(err) + } +} + +func TestCacheFailureReasonDoesNotExposeProviderCredentials(t *testing.T) { + err := fmt.Errorf("https://private.invalid/file?token=secret: %w", context.DeadlineExceeded) + reason := CacheFailureReason(err) + if !strings.Contains(reason, "timed out") || strings.Contains(reason, "secret") || strings.Contains(reason, "private.invalid") { + t.Fatalf("unsafe or unhelpful reason %q", reason) + } +} diff --git a/pkg/util/file_access_metrics.go b/pkg/util/file_access_metrics.go new file mode 100644 index 0000000..87c287d --- /dev/null +++ b/pkg/util/file_access_metrics.go @@ -0,0 +1,33 @@ +package util + +import "github.com/prometheus/client_golang/prometheus" + +// Like the existing application metrics, register once at package initialization. +// The existing server.metrics flag controls exposure of /metrics. Router creation +// never registers this collector, so multiple routers cannot register it twice. +func init() { + prometheus.MustRegister(newFileAccessCollector()) +} + +type fileAccessCollector struct { + desc *prometheus.Desc +} + +func newFileAccessCollector() *fileAccessCollector { + return &fileAccessCollector{desc: prometheus.NewDesc( + "dingospeed_storage_access_errors_total", + "Observed filesystem access errors in instrumented metadata operations since process start; not a storage health check.", + []string{"operation", "kind"}, nil, + )} +} + +func (c *fileAccessCollector) Describe(ch chan<- *prometheus.Desc) { ch <- c.desc } + +func (c *fileAccessCollector) Collect(ch chan<- prometheus.Metric) { + // Snapshot existing atomics only: scraping performs no filesystem access, + // logging, retries, or changes to request handling. + for _, observation := range FileAccessObservations() { + ch <- prometheus.MustNewConstMetric(c.desc, prometheus.CounterValue, + float64(observation.Count), observation.Operation, string(observation.Kind)) + } +} diff --git a/pkg/util/file_access_metrics_test.go b/pkg/util/file_access_metrics_test.go new file mode 100644 index 0000000..7cc7854 --- /dev/null +++ b/pkg/util/file_access_metrics_test.go @@ -0,0 +1,93 @@ +package util + +import ( + "errors" + "os" + "reflect" + "sync" + "testing" + + "github.com/prometheus/client_golang/prometheus" + dto "github.com/prometheus/client_model/go" +) + +func TestFileAccessMetricsSnapshot(t *testing.T) { + r := prometheus.NewPedanticRegistry() + collector := newFileAccessCollector() + r.MustRegister(collector) + var duplicate prometheus.AlreadyRegisteredError + if err := r.Register(newFileAccessCollector()); !errors.As(err, &duplicate) { + t.Fatalf("expected duplicate descriptor detection: %v", err) + } + before := FileAccessObservations() + ObserveFileAccessFailure("read", &os.PathError{Op: "read", Path: "/private/repo/file", Err: os.ErrPermission}) + ObserveFileAccessFailure("read", &os.PathError{Op: "read", Path: "missing", Err: os.ErrNotExist}) + for scrape := 0; scrape < 2; scrape++ { + families, err := r.Gather() + if err != nil { + t.Fatal(err) + } + if len(families) != 1 { + t.Fatalf("families=%d", len(families)) + } + f := families[0] + if f.GetName() != "dingospeed_storage_access_errors_total" || f.GetType() != dto.MetricType_COUNTER || len(f.Metric) != 16 { + t.Fatalf("unexpected family: %v", f) + } + seen := map[string]bool{} + for _, m := range f.Metric { + labels := map[string]string{} + for _, l := range m.Label { + labels[l.GetName()] = l.GetValue() + } + if len(labels) != 2 { + t.Fatalf("unexpected labels: %v", labels) + } + key := labels["operation"] + "/" + labels["kind"] + if seen[key] { + t.Fatal("duplicate series") + } + seen[key] = true + found := false + for _, b := range before { + if b.Operation == labels["operation"] && string(b.Kind) == labels["kind"] { + want := b.Count + if b.Operation == "read" && b.Kind == FileAccessPermissionDenied { + want++ + } + if m.Counter.GetValue() != float64(want) { + t.Fatalf("%s=%v want %d", key, m.Counter.GetValue(), want) + } + found = true + } + } + if !found { + t.Fatalf("unbounded label: %v", labels) + } + } + } + first := FileAccessObservations() + _, _ = r.Gather() + if !reflect.DeepEqual(first, FileAccessObservations()) { + t.Fatal("scraping changes observation counts") + } +} + +func TestFileAccessMetricsConcurrentScrapes(t *testing.T) { + r := prometheus.NewPedanticRegistry() + r.MustRegister(newFileAccessCollector()) + var wg sync.WaitGroup + for i := 0; i < 8; i++ { + wg.Add(1) + go func() { + defer wg.Done() + for j := 0; j < 20; j++ { + ObserveFileAccessFailure("write", &os.PathError{Op: "write", Path: "p", Err: os.ErrPermission}) + if _, err := r.Gather(); err != nil { + t.Error(err) + } + } + }() + } + wg.Wait() +} diff --git a/pkg/util/file_access_observation.go b/pkg/util/file_access_observation.go new file mode 100644 index 0000000..3a79571 --- /dev/null +++ b/pkg/util/file_access_observation.go @@ -0,0 +1,59 @@ +package util + +import ( + "dingospeed/pkg/dependency" + "sync/atomic" + "time" +) + +var fileAccessOperations = [...]string{"stat", "read", "mkdir", "write"} +var fileAccessKinds = [...]FileAccessErrorKind{ + FileAccessPermissionDenied, FileAccessMountDisconnected, FileAccessIOFailure, FileAccessUnavailable, +} +var fileAccessCounts [4][4]atomic.Uint64 + +// ObserveFileAccessFailure records only filesystem failures already returned by +// an operation. It performs no I/O, logging, retries or background work, and +// retains no paths or error objects. Counts describe failed operations, not +// failed mounts or nodes. Unknown operations and content errors are ignored. +func ObserveFileAccessFailure(operation string, err error) { + accessErr, ok := ClassifyFileAccessError("", err) + if !ok { + return + } + for i, op := range fileAccessOperations { + if op != operation { + continue + } + for j, kind := range fileAccessKinds { + if kind == accessErr.Kind { + fileAccessCounts[i][j].Add(1) + id := dependency.MetadataRead + if operation == "mkdir" || operation == "write" { + id = dependency.MetadataWrite + } + dependency.Default.Observe(id, false, time.Now()) + return + } + } + } +} + +type FileAccessObservation struct { + Operation string + Kind FileAccessErrorKind + Count uint64 +} + +// FileAccessObservations returns process-lifetime counters for diagnostics. +// Each count is atomic; the collection is not a transactional snapshot. +// The metrics collector exports these counts through the existing /metrics route. +func FileAccessObservations() []FileAccessObservation { + result := make([]FileAccessObservation, 0, 16) + for i, op := range fileAccessOperations { + for j, kind := range fileAccessKinds { + result = append(result, FileAccessObservation{op, kind, fileAccessCounts[i][j].Load()}) + } + } + return result +} diff --git a/pkg/util/file_create_test.go b/pkg/util/file_create_test.go new file mode 100644 index 0000000..c84911b --- /dev/null +++ b/pkg/util/file_create_test.go @@ -0,0 +1,32 @@ +package util + +import ( + "bytes" + "os" + "path/filepath" + "sync" + "testing" +) + +func TestCreateFileIfNotExistPreservesConcurrentBlob(t *testing.T) { + p := filepath.Join(t.TempDir(), "blob") + content := bytes.Repeat([]byte("cache block"), 512) + if err := os.WriteFile(p, content, 0600); err != nil { + t.Fatal(err) + } + var wg sync.WaitGroup + for i := 0; i < 32; i++ { + wg.Add(1) + go func() { + defer wg.Done() + if err := CreateFileIfNotExist(p); err != nil { + t.Error(err) + } + }() + } + wg.Wait() + b, err := os.ReadFile(p) + if err != nil || !bytes.Equal(b, content) { + t.Fatal("existing blob was truncated", err) + } +} diff --git a/pkg/util/file_range.go b/pkg/util/file_range.go new file mode 100644 index 0000000..8955c75 --- /dev/null +++ b/pkg/util/file_range.go @@ -0,0 +1,53 @@ +package util + +import ( + "fmt" + "strconv" + "strings" +) + +// FileRange converts a single HTTP byte range to the downloader's [start,end). +func FileRange(value string, size int64) (int64, int64, error) { + if size < 0 { + return 0, 0, fmt.Errorf("unknown file size") + } + if value == "" { + return 0, size, nil + } + if size == 0 || !strings.HasPrefix(value, "bytes=") || strings.Contains(value, ",") { + return 0, 0, fmt.Errorf("unsupported byte range") + } + parts := strings.Split(strings.TrimPrefix(value, "bytes="), "-") + if len(parts) != 2 { + return 0, 0, fmt.Errorf("invalid byte range") + } + number := func(s string) (int64, error) { + n, e := strconv.ParseInt(s, 10, 64) + if e != nil || n < 0 { + return 0, fmt.Errorf("invalid byte offset") + } + return n, nil + } + if parts[0] == "" { + n, e := number(parts[1]) + if e != nil || n == 0 { + return 0, 0, fmt.Errorf("invalid suffix range") + } + return max(0, size-n), size, nil + } + start, e := number(parts[0]) + if e != nil || start >= size { + return 0, 0, fmt.Errorf("range out of bounds") + } + end := size + if parts[1] != "" { + last, e := number(parts[1]) + if e != nil || last < start { + return 0, 0, fmt.Errorf("invalid range end") + } + if last < size-1 { + end = last + 1 + } + } + return start, end, nil +} diff --git a/pkg/util/http_cancel_test.go b/pkg/util/http_cancel_test.go new file mode 100644 index 0000000..3cddb90 --- /dev/null +++ b/pkg/util/http_cancel_test.go @@ -0,0 +1,72 @@ +package util + +import ( + "context" + "dingospeed/pkg/common" + "dingospeed/pkg/config" + "errors" + "io" + "net/http" + "net/http/httptest" + "testing" + "time" +) + +func TestStreamCancellationClosesBlockedRequest(t *testing.T) { + old := config.SysConfig + config.SysConfig = &config.Config{} + defer func() { config.SysConfig = old }() + for _, headers := range []bool{false, true} { + t.Run(map[bool]string{false: "headers", true: "body"}[headers], func(t *testing.T) { + entered, canceled := make(chan struct{}), make(chan struct{}) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if headers { + w.Header().Set("Content-Length", "100") + w.WriteHeader(200) + w.(http.Flusher).Flush() + } + close(entered) + <-r.Context().Done() + close(canceled) + })) + defer server.Close() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + done := make(chan error, 1) + go func() { + _, err := doGetStreamContext(ctx, server.Client(), server.URL, nil, func(resp *http.Response) error { _, err := io.Copy(io.Discard, resp.Body); return err }) + done <- err + }() + <-entered + cancel() + select { + case err := <-done: + if !errors.Is(err, context.Canceled) { + t.Fatalf("error=%v", err) + } + case <-time.After(time.Second): + t.Fatal("HTTP request did not stop") + } + select { + case <-canceled: + case <-time.After(time.Second): + t.Fatal("source did not observe cancellation") + } + }) + } +} + +func TestRetryCancellationInterruptsBackoff(t *testing.T) { + old := config.SysConfig + config.SysConfig = &config.Config{} + defer func() { config.SysConfig = old }() + config.SysConfig.Retry.Attempts = 5 + config.SysConfig.Retry.Delay = 30 + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + calls := 0 + _, err := RetryRequestContext(ctx, func() (*common.Response, error) { calls++; cancel(); return nil, errors.New("retryable failure") }) + if !errors.Is(err, context.Canceled) || calls != 1 { + t.Fatalf("calls=%d error=%v", calls, err) + } +} diff --git a/pkg/util/http_util.go b/pkg/util/http_util.go index c6477f6..8b9c435 100644 --- a/pkg/util/http_util.go +++ b/pkg/util/http_util.go @@ -17,6 +17,7 @@ package util import ( "bytes" "context" + "dingospeed/pkg/transfersettings" "errors" "fmt" "io" @@ -47,16 +48,27 @@ var ( ) func RetryRequest(f func() (*common.Response, error)) (*common.Response, error) { + return RetryRequestContext(context.Background(), f) +} + +func RetryRequestContext(ctx context.Context, f func() (*common.Response, error)) (*common.Response, error) { var resp *common.Response err := retry.Do( func() error { var err error + if err := ctx.Err(); err != nil { + return retry.Unrecoverable(err) + } resp, err = f() + if isClientCanceled(err) { + return retry.Unrecoverable(err) + } return err }, retry.Delay(time.Duration(config.SysConfig.Retry.Delay)*time.Second), retry.Attempts(config.SysConfig.Retry.Attempts), retry.DelayType(retry.FixedDelay), + retry.Context(ctx), ) return resp, err } @@ -181,25 +193,36 @@ func constructRoute(method string) (route, error) { } func Head(requestUri string, headers map[string]string) (*common.Response, error) { + return HeadContext(context.Background(), requestUri, headers) +} +func HeadContext(ctx context.Context, requestUri string, headers map[string]string) (*common.Response, error) { r, err := constructRoute(http.MethodHead) if err != nil { return nil, fmt.Errorf("construct http client err: %v", err) } requestURL := fmt.Sprintf("%s%s", r.domain, requestUri) - resp, err := doHead(r.client, requestURL, headers) - r.reportResp(resp, err) + resp, err := doHeadContext(ctx, r.client, requestURL, headers) + if !isClientCanceled(err) { + r.reportResp(resp, err) + } return resp, err } func doHead(client *http.Client, targetURL string, headers map[string]string) (*common.Response, error) { - req, err := http.NewRequest("HEAD", targetURL, nil) + return doHeadContext(context.Background(), client, targetURL, headers) +} +func doHeadContext(ctx context.Context, client *http.Client, targetURL string, headers map[string]string) (*common.Response, error) { + req, err := http.NewRequestWithContext(ctx, "HEAD", targetURL, nil) if err != nil { return nil, fmt.Errorf("创建HEAD请求失败: %v", err) } for key, value := range headers { + if strings.EqualFold(key, "X-Dingo-Service-Token") { + continue + } req.Header.Set(key, value) } - resp, err := client.Do(req) + resp, err := transfersettings.Do("huggingface", client, req) if err != nil { zap.S().Warnf("URL请求失败: %s, 错误: %v", targetURL, err) return nil, fmt.Errorf("执行HEAD请求失败: %v", err) @@ -221,25 +244,36 @@ func doHead(client *http.Client, targetURL string, headers map[string]string) (* } func Get(requestUri string, headers map[string]string) (*common.Response, error) { + return GetContext(context.Background(), requestUri, headers) +} +func GetContext(ctx context.Context, requestUri string, headers map[string]string) (*common.Response, error) { r, err := constructRoute(http.MethodGet) if err != nil { return nil, fmt.Errorf("construct http client err: %v", err) } requestURL := fmt.Sprintf("%s%s", r.domain, requestUri) - resp, err := doGet(r.client, requestURL, headers) - r.reportResp(resp, err) + resp, err := doGetContext(ctx, r.client, requestURL, headers) + if !isClientCanceled(err) { + r.reportResp(resp, err) + } return resp, err } func doGet(client *http.Client, targetURL string, headers map[string]string) (*common.Response, error) { - req, err := http.NewRequest("GET", targetURL, nil) + return doGetContext(context.Background(), client, targetURL, headers) +} +func doGetContext(ctx context.Context, client *http.Client, targetURL string, headers map[string]string) (*common.Response, error) { + req, err := http.NewRequestWithContext(ctx, "GET", targetURL, nil) if err != nil { return nil, fmt.Errorf("创建GET请求失败: %v", err) } for key, value := range headers { + if strings.EqualFold(key, "X-Dingo-Service-Token") { + continue + } req.Header.Set(key, value) } - resp, err := client.Do(req) + resp, err := transfersettings.Do("huggingface", client, req) if err != nil { zap.S().Warnf("URL请求失败: %s, 错误: %v", targetURL, err) return nil, fmt.Errorf("执行GET请求失败: %v", err) @@ -270,6 +304,10 @@ func doGet(client *http.Client, targetURL string, headers map[string]string) (*c } func GetStream(domain, uri string, headers map[string]string, f func(r *http.Response) error) error { + return GetStreamContext(context.Background(), domain, uri, headers, f) +} + +func GetStreamContext(ctx context.Context, domain, uri string, headers map[string]string, f func(r *http.Response) error) error { // 内网目标（兄弟节点互拉、回环上传口）必须旁路代理： // 这些地址走公网出口必然失败，且会把好出口误判成坏出口。 if IsInnerDomain(domain) || ProxyPool().ShouldBypass(domain) { @@ -278,7 +316,7 @@ func GetStream(domain, uri string, headers map[string]string, f func(r *http.Res return fmt.Errorf("construct http client err: %v", err) } headers[consts.RequestSourceInner] = Itoa(1) - _, err = doGetStream(client, fmt.Sprintf("%s%s", domain, uri), headers, f) + _, err = doGetStreamContext(ctx, client, fmt.Sprintf("%s%s", domain, uri), headers, f) return err } r, err := constructRoute(http.MethodGet) @@ -286,7 +324,7 @@ func GetStream(domain, uri string, headers map[string]string, f func(r *http.Res return fmt.Errorf("construct http client err: %v", err) } requestURL := fmt.Sprintf("%s%s", r.domain, uri) - code, err := doGetStream(r.client, requestURL, headers, f) + code, err := doGetStreamContext(ctx, r.client, requestURL, headers, f) // 这里上报的 err 包含了流式传输中途的失败（f 回调里读 body 断掉）， // 而这正是 gost 的 TCP 层健康检查看不见的那一类故障。 // 但客户端主动取消不能算在出口头上，否则用户中断下载会误伤健康出口。 @@ -304,15 +342,22 @@ func isClientCanceled(err error) bool { // doGetStream 返回响应状态码与错误，状态码供代理池计分使用；未拿到响应时返回 0。 func doGetStream(client *http.Client, targetURL string, headers map[string]string, f func(r *http.Response) error) (int, error) { + return doGetStreamContext(context.Background(), client, targetURL, headers, f) +} + +func doGetStreamContext(ctx context.Context, client *http.Client, targetURL string, headers map[string]string, f func(r *http.Response) error) (int, error) { escapedURL := strings.ReplaceAll(targetURL, "#", "%23") - req, err := http.NewRequest("GET", escapedURL, nil) + req, err := http.NewRequestWithContext(ctx, "GET", escapedURL, nil) if err != nil { return 0, fmt.Errorf("创建GET请求失败: %v", err) } for key, value := range headers { + if strings.EqualFold(key, "X-Dingo-Service-Token") { + continue + } req.Header.Set(key, value) } - resp, err := client.Do(req) + resp, err := transfersettings.Do("huggingface", client, req) if err != nil { return 0, err } @@ -325,28 +370,39 @@ func doGetStream(client *http.Client, targetURL string, headers map[string]strin } func Post(requestUri string, contentType string, data []byte, headers map[string]string) (*common.Response, error) { + return PostContext(context.Background(), requestUri, contentType, data, headers) +} +func PostContext(ctx context.Context, requestUri string, contentType string, data []byte, headers map[string]string) (*common.Response, error) { r, err := constructRoute(http.MethodPost) if err != nil { return nil, fmt.Errorf("construct http client err: %v", err) } requestURL := fmt.Sprintf("%s%s", r.domain, requestUri) - resp, err := doPost(r.client, requestURL, contentType, data, headers) - r.reportResp(resp, err) + resp, err := doPostContext(ctx, r.client, requestURL, contentType, data, headers) + if !isClientCanceled(err) { + r.reportResp(resp, err) + } return resp, err } func doPost(client *http.Client, targetURL string, contentType string, data []byte, headers map[string]string) (*common.Response, error) { - req, err := http.NewRequest("POST", targetURL, bytes.NewBuffer(data)) + return doPostContext(context.Background(), client, targetURL, contentType, data, headers) +} +func doPostContext(ctx context.Context, client *http.Client, targetURL string, contentType string, data []byte, headers map[string]string) (*common.Response, error) { + req, err := http.NewRequestWithContext(ctx, "POST", targetURL, bytes.NewBuffer(data)) if err != nil { return nil, fmt.Errorf("创建POST请求失败: %v", err) } req.Header.Set("Content-Type", contentType) for key, value := range headers { + if strings.EqualFold(key, "X-Dingo-Service-Token") { + continue + } req.Header.Set(key, value) } - resp, err := client.Do(req) + resp, err := transfersettings.Do("huggingface", client, req) if err != nil { zap.S().Warnf("URL请求失败: %s, 错误: %v", targetURL, err) return nil, fmt.Errorf("执行POST请求失败: %v", err) @@ -397,6 +453,8 @@ func ResponseStream(c echo.Context, fileName string, headers map[string]string, } for { select { + case <-c.Request().Context().Done(): + return c.Request().Context().Err() case b, ok := <-content: if !ok { zap.S().Infof("ResponseStream complete, %s", fileName) @@ -437,16 +495,19 @@ func ForwardRequest(originalReq echo.Context) (*http.Response, error) { Path: forwardPath, RawQuery: originalReq.Request().URL.RawQuery, } - proxyReq, err := http.NewRequest(originalReq.Request().Method, forwardURL.String(), originalReq.Request().Body) + proxyReq, err := http.NewRequestWithContext(originalReq.Request().Context(), originalReq.Request().Method, forwardURL.String(), originalReq.Request().Body) if err != nil { return nil, fmt.Errorf("创建转发请求失败: %v", err) } for key, values := range originalReq.Request().Header { + if strings.EqualFold(key, "X-Dingo-Service-Token") { + continue + } for _, value := range values { proxyReq.Header.Add(key, value) } } - resp, err := client.Do(proxyReq) + resp, err := transfersettings.Do("huggingface", client, proxyReq) if err != nil { r.report(0, err) zap.S().Warnf("转发请求失败: %s, 错误: %v", targetURL, err) diff --git a/pkg/util/modelscope_retry_test.go b/pkg/util/modelscope_retry_test.go new file mode 100644 index 0000000..f494ccd --- /dev/null +++ b/pkg/util/modelscope_retry_test.go @@ -0,0 +1,55 @@ +package util + +import ( + "context" + "errors" + "io" + "net/http" + "net/http/httptest" + "testing" + "time" + + "dingospeed/pkg/config" +) + +func TestModelScopeZeroRetryStillRequests(t *testing.T) { + old := config.SysConfig + t.Cleanup(func() { config.SysConfig = old }) + config.SysConfig = &config.Config{} + calls := 0 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { calls++; io.WriteString(w, "ok") })) + defer server.Close() + req, _ := http.NewRequest("GET", server.URL, nil) + resp, err := DoRequestWithRetry(req) + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + if calls != 1 || resp.StatusCode != 200 { + t.Fatalf("initial request suppressed: calls=%d", calls) + } +} + +type msTimeoutTransport struct{ calls int } + +func (t *msTimeoutTransport) RoundTrip(*http.Request) (*http.Response, error) { + t.calls++ + return nil, context.DeadlineExceeded +} + +func TestModelScopeRetryDelayUsesSecondsAndCancels(t *testing.T) { + old := config.SysConfig + client := CreateHTTPClient() + oldTransport := client.Transport + transport := &msTimeoutTransport{} + client.Transport = transport + t.Cleanup(func() { config.SysConfig = old; client.Transport = oldTransport }) + config.SysConfig = &config.Config{Modelscope: config.Modelscope{MaxRetry: 3, RetryDelay: 1}} + ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond) + defer cancel() + req, _ := http.NewRequestWithContext(ctx, "GET", "http://modelscope.test", nil) + _, err := DoRequestWithRetry(req) + if !errors.Is(err, context.DeadlineExceeded) || transport.calls != 1 { + t.Fatalf("retry did not wait for seconds/cancellation: calls=%d error=%v", transport.calls, err) + } +} diff --git a/pkg/util/modelscope_util.go b/pkg/util/modelscope_util.go index 7c96fbd..0b98c59 100644 --- a/pkg/util/modelscope_util.go +++ b/pkg/util/modelscope_util.go @@ -2,6 +2,7 @@ package util import ( "crypto/tls" + "dingospeed/pkg/transfersettings" "fmt" "net/http" "os" @@ -77,35 +78,6 @@ func EnsureDir(path string) error { return nil } -func GetCachePath(repoType, repoId, revision, filePath string) (string, bool) { - parts := strings.Split(repoId, "/") - if len(parts) != 2 { - zap.S().Errorf("无效的repoId格式: %s，需为 org/repo 格式", repoId) - return "", false - } - - var cacheRoot string - switch repoType { - case "datasets": - cacheRoot = config.SysConfig.GetDatasetCacheRoot() - case "models", "": - cacheRoot = config.SysConfig.GetModelCacheRoot() - default: - zap.S().Warnf("未知的repoType: %s，默认使用models缓存目录", repoType) - cacheRoot = config.SysConfig.GetModelCacheRoot() - } - - targetCachePath := filepath.Join(cacheRoot, parts[0], parts[1], revision, filepath.Clean(filePath)) - fileInfo, err := os.Stat(targetCachePath) - if err == nil { - zap.S().Debugf("缓存文件存在: %s, 大小: %d字节", targetCachePath, fileInfo.Size()) - return targetCachePath, true - } - - _ = EnsureDir(targetCachePath) - return targetCachePath, false -} - var ( httpClientOnce sync.Once globalHTTPClient *http.Client @@ -147,23 +119,36 @@ func DoRequestWithRetry(req *http.Request) (*http.Response, error) { client := CreateHTTPClient() var resp *http.Response var err error + attempts := config.SysConfig.Modelscope.MaxRetry + if attempts < 1 { + attempts = 1 // A retry setting must never suppress the initial request. + } - for i := 0; i < config.SysConfig.Modelscope.MaxRetry; i++ { - resp, err = client.Do(req) + for i := 0; i < attempts; i++ { + resp, err = transfersettings.Do("modelscope", client, req) if err == nil { return resp, nil } if strings.Contains(err.Error(), "timeout") || strings.Contains(err.Error(), "deadline exceeded") { - zap.S().Warnf("⚠️ Retry %d/%d: request timeout - %v", i+1, config.SysConfig.Modelscope.MaxRetry, err) - time.Sleep(time.Duration(config.SysConfig.Modelscope.RetryDelay) * time.Duration(i+1)) + zap.S().Warnf("⚠️ Retry %d/%d: request timeout - %v", i+1, attempts, err) + if i+1 == attempts { + break + } + timer := time.NewTimer(time.Duration(config.SysConfig.Modelscope.RetryDelay) * time.Second * time.Duration(i+1)) + select { + case <-req.Context().Done(): + timer.Stop() + return nil, req.Context().Err() + case <-timer.C: + } continue } return nil, err } - return nil, fmt.Errorf("failed after %d retries: %v", config.SysConfig.Modelscope.MaxRetry, err) + return nil, fmt.Errorf("failed after %d attempts: %w", attempts, err) } // ParseRangeHeader 解析Range请求头，返回起始字节和结束字节（-1表示到末尾） diff --git a/pkg/util/peer_http.go b/pkg/util/peer_http.go new file mode 100644 index 0000000..e406346 --- /dev/null +++ b/pkg/util/peer_http.go @@ -0,0 +1,32 @@ +package util + +import ( + "context" + "dingospeed/pkg/config" + "dingospeed/pkg/consts" + "dingospeed/pkg/transfersettings" + "net/http" +) + +// GetPeerStream reads from the selected peer without following redirects. +func GetPeerStream(domain, uri string, headers map[string]string, consume func(*http.Response) error) error { + return GetPeerStreamContext(context.Background(), domain, uri, headers, consume) +} + +func GetPeerStreamContext(ctx context.Context, domain, uri string, headers map[string]string, consume func(*http.Response) error) error { + req, err := http.NewRequestWithContext(ctx, http.MethodGet, domain+uri, nil) + if err != nil { + return err + } + for key, value := range headers { + req.Header.Set(key, value) + } + req.Header.Set(consts.RequestSourceInner, "1") + client := &http.Client{Timeout: config.SysConfig.GetReqTimeOut(), CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }} + resp, err := transfersettings.Do("peer", client, req) + if err != nil { + return err + } + defer resp.Body.Close() + return consume(resp) +} diff --git a/pkg/util/repo_util.go b/pkg/util/repo_util.go index 65f554a..8baaabe 100644 --- a/pkg/util/repo_util.go +++ b/pkg/util/repo_util.go @@ -72,19 +72,19 @@ func PathExists(filePath string) (bool, error) { if os.IsNotExist(err) { return false, nil } - accessErr, _ := ClassifyFileAccessError(filePath, err) - return false, accessErr + return false, &FileAccessError{Kind: classifyFileAccessError(err), Path: filePath, Err: err} } // ClassifyFileAccessError recognizes filesystem access failures even when // another layer has wrapped the original *os.PathError. Content and parsing // errors are deliberately left unclassified. func ClassifyFileAccessError(filePath string, err error) (*FileAccessError, bool) { - if err == nil || os.IsNotExist(err) { + if err == nil || errors.Is(err, os.ErrNotExist) { return nil, false } var pathErr *os.PathError - if !errors.As(err, &pathErr) { + var linkErr *os.LinkError + if !errors.As(err, &pathErr) && !errors.As(err, &linkErr) { return nil, false } return &FileAccessError{Kind: classifyFileAccessError(err), Path: filePath, Err: err}, true @@ -95,7 +95,9 @@ func classifyFileAccessError(err error) FileAccessErrorKind { } func classifyFileAccessErrorForOS(err error, goos string) FileAccessErrorKind { - if goos != "linux" && os.IsPermission(err) { + var errno syscall.Errno + hasErrno := errors.As(err, &errno) + if !hasErrno && errors.Is(err, os.ErrPermission) { return FileAccessPermissionDenied } @@ -106,16 +108,17 @@ func classifyFileAccessErrorForOS(err error, goos string) FileAccessErrorKind { linuxENOTCONN syscall.Errno = 107 linuxESTALE syscall.Errno = 116 ) - var errno syscall.Errno - if errors.As(err, &errno) { + if goos == "linux" && hasErrno { switch errno { + case 1, 13: // EPERM, EACCES; use Linux semantics even in cross-platform tests. + return FileAccessPermissionDenied case linuxENOTCONN, linuxESTALE: return FileAccessMountDisconnected case linuxEIO: return FileAccessIOFailure } } - if os.IsPermission(err) { + if (goos == "windows" && hasErrno && errno == 5) || (goos == runtime.GOOS && errors.Is(err, os.ErrPermission)) { return FileAccessPermissionDenied } return FileAccessUnavailable @@ -130,7 +133,7 @@ func GetOrgRepo(org, repo string) string { } func SplitOrgRepo(orgRepo string) (string, string) { - splits := strings.Split(orgRepo, "/") + splits := strings.SplitN(orgRepo, "/", 2) if len(splits) == 0 { return "", "" } else if len(splits) == 1 { @@ -186,10 +189,15 @@ func CreateFile(filePath string) error { } func CreateFileIfNotExist(filePath string) error { - if exists := FileExists(filePath); !exists { - return CreateFile(filePath) + // Exclusive creation never truncates a blob another request already opened. + f, err := os.OpenFile(filePath, os.O_CREATE|os.O_EXCL|os.O_WRONLY, 0644) + if os.IsExist(err) { + return nil } - return nil + if err != nil { + return err + } + return f.Close() } func DeleteFile(filePath string) error { diff --git a/pkg/util/repo_util_test.go b/pkg/util/repo_util_test.go index c4a7fc7..7127aa5 100644 --- a/pkg/util/repo_util_test.go +++ b/pkg/util/repo_util_test.go @@ -5,6 +5,8 @@ import ( "fmt" "os" "path/filepath" + "runtime" + "sync" "syscall" "testing" ) @@ -27,17 +29,77 @@ func TestPathExistsResults(t *testing.T) { } func TestClassifyWrappedFileAccessError(t *testing.T) { - original := &os.PathError{Op: "read", Path: "/repos/meta.json", Err: syscall.Errno(107)} + original := &os.PathError{Op: "read", Path: "/repos/meta.json", Err: os.ErrPermission} accessErr, ok := ClassifyFileAccessError("/repos/meta.json", fmt.Errorf("cache read: %w", original)) if !ok { t.Fatal("expected wrapped filesystem error to be classified") } - if accessErr.Kind != FileAccessMountDisconnected { - t.Fatalf("got kind %q, want %q", accessErr.Kind, FileAccessMountDisconnected) + if accessErr.Kind != FileAccessPermissionDenied || !errors.Is(accessErr, original) { + t.Fatalf("classification or original cause lost: %v", accessErr) } if _, ok = ClassifyFileAccessError("/repos/meta.json", errors.New("invalid metadata")); ok { t.Fatal("content error must not be classified as a storage access failure") } + link := &os.LinkError{Op: "rename", Old: "tmp", New: "meta.json", Err: os.ErrPermission} + classified, ok := ClassifyFileAccessError("meta.json", fmt.Errorf("atomic write: %w", link)) + if !ok || classified.Kind != FileAccessPermissionDenied || !errors.Is(classified, link) { + t.Fatal("rename failure or original error lost") + } +} + +func TestFileAccessPlatformSemantics(t *testing.T) { + for _, tc := range []struct { + platform string + code syscall.Errno + want FileAccessErrorKind + }{ + {"linux", 5, FileAccessIOFailure}, {"linux", 13, FileAccessPermissionDenied}, + {"linux", 1, FileAccessPermissionDenied}, {"linux", 107, FileAccessMountDisconnected}, + {"linux", 116, FileAccessMountDisconnected}, {"linux", 999, FileAccessUnavailable}, + {"windows", 5, FileAccessPermissionDenied}, {"windows", 107, FileAccessUnavailable}, + {"windows", 116, FileAccessUnavailable}, {"darwin", 107, FileAccessUnavailable}, + } { + t.Run(fmt.Sprintf("%s/%d", tc.platform, tc.code), func(t *testing.T) { + err := fmt.Errorf("wrapped: %w", &os.PathError{Op: "read", Path: "p", Err: tc.code}) + if got := classifyFileAccessErrorForOS(err, tc.platform); got != tc.want { + t.Fatalf("got %s want %s", got, tc.want) + } + if tc.platform == runtime.GOOS && classifyFileAccessError(err) != tc.want { + t.Fatal("native classification differs") + } + }) + } +} + +func TestFileAccessObservationIsBoundedAndIgnoresNonAccessErrors(t *testing.T) { + before := FileAccessObservations() + for _, err := range []error{nil, errors.New("invalid JSON"), fmt.Errorf("wrapped: %w", &os.PathError{Op: "read", Path: "p", Err: os.ErrNotExist})} { + if _, ok := ClassifyFileAccessError("p", err); ok { + t.Fatalf("unexpected classification: %v", err) + } + ObserveFileAccessFailure("read", err) + } + err := fmt.Errorf("wrapped: %w", &os.PathError{Op: "read", Path: "p", Err: os.ErrPermission}) + ObserveFileAccessFailure("unbounded-operation", err) + var wg sync.WaitGroup + for i := 0; i < 32; i++ { + wg.Add(1) + go func() { defer wg.Done(); ObserveFileAccessFailure("read", err) }() + } + wg.Wait() + after := FileAccessObservations() + if len(after) != 16 { + t.Fatalf("unbounded observations: %d", len(after)) + } + for i, observation := range after { + want := before[i].Count + if observation.Operation == "read" && observation.Kind == FileAccessPermissionDenied { + want += 32 + } + if observation.Count != want { + t.Fatalf("%+v want count %d", observation, want) + } + } } func TestClassifyFileAccessError(t *testing.T) { diff --git a/test/cachemount/accept_tree.py b/test/cachemount/accept_tree.py new file mode 100644 index 0000000..99357d1 --- /dev/null +++ b/test/cachemount/accept_tree.py @@ -0,0 +1,221 @@ +#!/usr/bin/env python3 +"""User-operated acceptance, never run automatically by the build. + +demo --work NEW_DIRECTORY creates isolated HF/hosted/partial examples. +check --report MOUNT_REPORT --output NEW_HTML [--full] inspects a live mount. +The report is real filesystem evidence, not pre-rendered expected success. +""" +import argparse +import errno +import hashlib +import html +import json +import os +from pathlib import Path +import struct + + +def wrapper(path, body): + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps({"status_code": 200, "content": json.dumps(body).encode().hex()})) + + +def demo(work): + work.mkdir(parents=True, exist_ok=False) + root = work / "repos" + selections = [] + for namespace, repo, revision, broken in [("huggingface", "demo/HFModel", "main", False), + ("datacanvas", "team/UploadModel", "v1", False), + ("datacanvas", "team/Incomplete", "v1", True)]: + storage = repo if namespace == "huggingface" else "dingo-local/" + namespace + "/" + repo + api = root / "api/models" / storage + blobs = root / "files/models" / storage / "blobs" + blobs.mkdir(parents=True) + data = {"config.json": json.dumps({"model_type": "acceptance-demo", "namespace": namespace}, indent=2).encode(), + "tokenizer/tokenizer.json": b'{"demo": "nested path retained"}', + "model.safetensors": struct.pack("

{esc(title)}

{esc(v['error'])}

{esc(result)}

") + continue + expected = {f["path"]: f for f in v["files"]} + actual = {} + errors = [] + try: + independent = original_manifest(v) + assert set(independent) == set(expected), "挂载清单与原仓库文件清单不一致" + for name, (size, identity) in independent.items(): + if size >= 0: + assert expected[name]["size"] == size, name + " 清单大小不一致" + if identity: + assert expected[name]["content_id"] == identity, name + " 内容身份不一致" + except (OSError, AssertionError, ValueError, KeyError) as e: + errors.append("独立读取原仓库清单失败：" + str(e)) + try: + for base, _, files in os.walk(root, onerror=lambda e: errors.append(str(e))): + for name in files: + p = Path(base)/name + actual[p.relative_to(root).as_posix()] = p.stat().st_size + except OSError as e: + errors.append(str(e)) + missing = sorted(set(expected)-set(actual)) + extra = sorted(set(actual)-set(expected)) + rows = [] + for name, f in expected.items(): + status = "未读取" + try: + p = root/name + size = p.stat().st_size + assert size == f["size"] or f["size"] == -1, "大小不同" + mode_full = args.full or size <= 16*1024*1024 + with open(p, "rb") as view: + if mode_full: + digest = hashlib.sha256() + # Independently compare against backing payload, not the + # mount's own read routine or its claimed checksum. + at = 0 + while b := view.read(4*1024*1024): + assert b == read_payload(f["backing"], at, len(b)), "正文不同" + digest.update(b) + at += len(b) + assert at == size, "正文读取长度不足" + content_id = f["content_id"] + if len(content_id) == 64: + assert digest.hexdigest() == content_id, "与仓库 SHA256 不同" + status = "全量正文一致" + else: + for at in [0, max(0, size//2-2048), max(0, size-4096)]: + view.seek(at) + assert view.read(4096) == read_payload(f["backing"], at, 4096), "抽样不同" + status = "头/中/尾抽样一致；未做全量校验" + if args.originals: + original = Path(args.originals)/title/name + assert original.read_bytes() == p.read_bytes(), "与独立原始样本不同" + status += "；与原始样本一致" + except (OSError, AssertionError) as e: + status = "失败：" + str(e) + errors.append(name + ": " + str(e)) + rows.append(f"{esc(name)}{esc(f['size'])}{esc(actual.get(name,'缺失'))}{esc(status)}") + passed = not (missing or extra or errors) + all_pass &= passed + preview = "" + if "config.json" in expected and "config.json" in actual: + f = expected["config.json"] + with open(f["backing"], "rb") as stream: + head = stream.read(36) + with open(root/"config.json", "rb") as stream: + body = stream.read(2048) + preview = f"

直接看差别

内部缓存开头：

{esc(head.hex(' '))}

上层读取 config.json：

{esc(body.decode('utf-8',errors='replace'))}
" + sections.append(f"

{esc(title)} — {'通过' if passed else '失败'}

source_file_path：{esc(root)}

" + f"

仓库清单（挂载时固定）

{esc(chr(10).join(sorted(expected)))}
" + f"

实际文件夹（现场读取）

{esc(chr(10).join(sorted(actual)))}
" + f"

已独立读取原仓库元数据核对映射。缺失：{esc(missing)}；额外：{esc(extra)}；错误：{esc(errors)}

" + f"{''.join(rows)}
文件路径预期大小实际大小读取结果
{preview}
") + page = "模型目录现场验收" + page += f"

模型目录现场验收：{'通过' if all_pass else '有失败项'}

报告由当前真实挂载读取生成。预期树来自挂载时固定的版本清单；可与仓库页面核对。合成样本不证明真实训练通过。大文件未指定 --full 时仅抽样，不能视为全量验收。

" + ''.join(sections) + with open(args.output, "x", encoding="utf-8") as f: + f.write(page) + print(args.output) + if not all_pass: + raise SystemExit(1) + + +def main(): + ap = argparse.ArgumentParser() + sub = ap.add_subparsers(dest="command", required=True) + d = sub.add_parser("demo") + d.add_argument("--work", required=True) + c = sub.add_parser("check") + c.add_argument("--report", required=True) + c.add_argument("--output", required=True) + c.add_argument("--originals") + c.add_argument("--full", action="store_true") + args = ap.parse_args() + if args.command == "demo": + demo(Path(args.work).resolve()) + else: + check(args) + + +if __name__ == "__main__": + main() diff --git a/test/cachemount/validate.py b/test/cachemount/validate.py new file mode 100644 index 0000000..ae53873 --- /dev/null +++ b/test/cachemount/validate.py @@ -0,0 +1,221 @@ +#!/usr/bin/env python3 +"""Isolated, real-FUSE validation. No production cache or global cache eviction. + +python3 validate.py --binary /path/cachemount --work /tmp/new-validation-dir +Optional torch installation enables actual mmap load/train/checkpoint/resume. +""" +import argparse +import concurrent.futures +import hashlib +import gc +import json +import mmap +import os +from pathlib import Path +import statistics +import struct +import subprocess +import time + + +def digest(path): + h = hashlib.sha256() + with open(path, "rb") as f: + while b := f.read(4 * 1024 * 1024): + h.update(b) + return h.hexdigest() + + +def wrap(source, target): + size = source.stat().st_size + block = 1024 * 1024 + count = (size + block - 1) // block + bits = max(8, ((count + 7) // 8) * 8) + with open(source, "rb") as src, open(target, "wb") as dst: + dst.write(b"OLAH" + struct.pack(" deadline: + raise RuntimeError("mount failed; inspect mount.log") + time.sleep(0.01) + report["mount_ready_seconds"] = time.perf_counter() - started + for name in files: + assert digest(raw / name) == digest(mount / name), name + assert (raw / name).stat().st_size == (mount / name).stat().st_size + with open(mount / "weights.bin", "rb") as f, mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ) as mm: + for off in [0, 1, 4095, 1024 * 1024 - 3, len(mm) - 9]: + assert mm[off:off+9] == bytes((i % 256 for i in range(off, off+9))) + with concurrent.futures.ProcessPoolExecutor(max_workers=4) as pool: + hashes = list(pool.map(digest, [mount / "weights.bin"] * 8)) + assert len(set(hashes)) == 1 + for operation in [lambda: (mount / "new").write_bytes(b"x"), lambda: (mount / "config.json").unlink()]: + try: + operation() + except OSError: + pass + else: + raise AssertionError("write accepted") + timings = {"plain": [], "fuse": []} + for _ in range(5): + timings["plain"].append(scan(raw / "weights.bin")) + timings["fuse"].append(scan(mount / "weights.bin")) + report["warm_scan_seconds"] = timings + report["warm_scan_median_seconds"] = {k: statistics.median(v) for k, v in timings.items()} + start = time.perf_counter() + with open(cache / "weights.bin", "rb") as src, open(work / "export.bin", "wb") as dst: + src.seek(offsets["weights.bin"]) + while b := src.read(4 * 1024 * 1024): + dst.write(b) + dst.flush() + os.fsync(dst.fileno()) + report["export_fsync_seconds"] = time.perf_counter() - start + report["export_extra_bytes"] = (work / "export.bin").stat().st_size + if torch: + state = torch.load(mount / "model.pt", mmap=True, weights_only=True) + model.load_state_dict(state) + reference = torch.nn.Sequential(torch.nn.Linear(32, 64), torch.nn.ReLU(), torch.nn.Linear(64, 8)) + reference.load_state_dict(torch.load(raw / "model.pt", mmap=True, weights_only=True)) + x, y = torch.randn(4, 32), torch.randn(4, 8) + assert torch.equal(model(x), reference(x)) + opt = torch.optim.SGD(model.parameters(), lr=0.01) + losses = [] + for _ in range(5): + opt.zero_grad() + loss = ((model(x)-y)**2).mean() + loss.backward() + opt.step() + losses.append(loss.item()) + torch.save({"model": model.state_dict(), "optimizer": opt.state_dict()}, work / "checkpoint.pt") + saved = torch.load(work / "checkpoint.pt", weights_only=True) + reference.load_state_dict(saved["model"]) + resumed_opt = torch.optim.SGD(reference.parameters(), lr=0.01) + resumed_opt.load_state_dict(saved["optimizer"]) + assert torch.equal(model(x), reference(x)) + resumed_opt.zero_grad() + ((reference(x)-y)**2).mean().backward() + resumed_opt.step() + report["training"] = {"steps": 5, "losses": losses, "checkpoint_resume": "passed"} + # torch.load(mmap=True) keeps its source mapping alive with state. + del state + gc.collect() + if torch and AutoModelForCausalLM: + hf_model = AutoModelForCausalLM.from_pretrained(mount, local_files_only=True) + hf_reference = AutoModelForCausalLM.from_pretrained(raw, local_files_only=True) + hf_model.eval() + hf_reference.eval() + ids = torch.arange(16).reshape(2, 8) + with torch.no_grad(): + assert torch.equal(hf_model(ids).logits, hf_reference(ids).logits) + hf_model.train() + hf_opt = torch.optim.SGD(hf_model.parameters(), lr=0.001) + hf_losses = [] + for _ in range(3): + hf_opt.zero_grad() + loss = hf_model(input_ids=ids, labels=ids).loss + loss.backward() + hf_opt.step() + hf_losses.append(loss.item()) + checkpoint = work / "hf-checkpoint" + hf_model.save_pretrained(checkpoint) + hf_resumed = AutoModelForCausalLM.from_pretrained(checkpoint, local_files_only=True) + hf_resumed.eval() + hf_model.eval() + with torch.no_grad(): + assert torch.equal(hf_model(ids).logits, hf_resumed(ids).logits) + hf_resumed.train() + resume_opt = torch.optim.SGD(hf_resumed.parameters(), lr=0.001) + resume_opt.load_state_dict(hf_opt.state_dict()) + hf_resumed(input_ids=ids, labels=ids).loss.backward() + resume_opt.step() + report["transformers"] = {"offline_from_pretrained": "passed", "format": "safetensors", "steps": 3, + "losses": hf_losses, "checkpoint_resume": "passed", "model": "random tiny GPT2"} + del hf_model, hf_reference, hf_resumed, hf_opt, resume_opt, loss + gc.collect() + after = {p.name: (digest(p), p.stat().st_size, p.stat().st_mtime_ns) for p in cache.iterdir()} + assert before == after, "backing cache changed" + report["source_unchanged"] = True + report["validation"] = "passed" + finally: + if os.path.ismount(mount): + subprocess.run(["fusermount3", "-u", str(mount)], check=True) + proc.wait(timeout=10) + log.close() + (work / "report.json").write_text(json.dumps(report, indent=2)) + print(json.dumps(report, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/test/cachemount/validate_existing.py b/test/cachemount/validate_existing.py new file mode 100644 index 0000000..9bc480b --- /dev/null +++ b/test/cachemount/validate_existing.py @@ -0,0 +1,134 @@ +#!/usr/bin/env python3 +"""Read-only validation of an existing, complete OLAH cache. No payload export. + +Full direct-payload and FUSE SHA256 passes plus bounded random/mmap/concurrent +reads. Source is opened O_RDONLY only; all output goes to a new work directory. +No global cache eviction. Timings include the host's existing cache state. +""" +import argparse +import concurrent.futures +import hashlib +import json +import mmap +import os +from pathlib import Path +import random +import struct +import subprocess +import time + + +def stamp(path): + s = os.stat(path) + return {"size": s.st_size, "mtime_ns": s.st_mtime_ns, "inode": s.st_ino, + "allocated_bytes": s.st_blocks * 512} + + +def full_pass(path, offset, size, label): + h = hashlib.sha256() + read_seconds = 0 + count = 0 + next_log = 4 * 1024**3 + start = time.perf_counter() + with open(path, "rb", buffering=0) as f: + f.seek(offset) + while count < size: + tick = time.perf_counter() + b = f.read(min(4 * 1024**2, size-count)) + read_seconds += time.perf_counter()-tick + if not b: + raise RuntimeError("unexpected EOF") + h.update(b) + count += len(b) + if count >= next_log: + print(f"{label}: {count/1024**3:.0f}/{size/1024**3:.0f} GiB, elapsed {time.perf_counter()-start:.1f}s", flush=True) + next_log += 4 * 1024**3 + elapsed = time.perf_counter()-start + result = {"bytes": count, "sha256": h.hexdigest(), "elapsed_seconds": elapsed, + "read_call_seconds": read_seconds, "read_call_mib_per_second": size/1024**2/read_seconds, + "hash_and_read_mib_per_second": size/1024**2/elapsed} + print(label + " complete: " + json.dumps(result), flush=True) + return result + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--source", required=True) + ap.add_argument("--binary", required=True) + ap.add_argument("--work", required=True) + ap.add_argument("--expected-sha256", required=True) + args = ap.parse_args() + source = Path(args.source).resolve(strict=True) + before = stamp(source) + with open(source, "rb") as f: + header = f.read(36) + version, block, size, bits = struct.unpack(" 0 and 0 < bits <= 8*1024*1024 + mask = f.read((bits+7)//8) + offset = 36+len(mask) + count = (size+block-1)//block + assert len(mask) == (bits+7)//8 and count <= bits and before["size"] >= offset+size + assert all(mask[i//8] & (1 << (i%8)) for i in range(count)), "incomplete cache" + work = Path(args.work) + work.mkdir(parents=True, exist_ok=False) + mount = work/"mount" + mount.mkdir() + manifest = work/"manifest.json" + manifest.write_text(json.dumps({"payload.bin": str(source)})) + report = {"source": str(source), "source_before": before, "payload_bytes": size, + "header_bytes": offset, "block_bytes": block, "complete_blocks": count, + "environment": "WSL Ubuntu; backing file on Windows D: via /mnt/d (DrvFS/9p)", + "cache_policy": "no cache flush; full passes direct then FUSE, not controlled cold-disk tests", + "expected_sha256": args.expected_sha256} + log = open(work/"mount.log", "w") + start = time.perf_counter() + proc = subprocess.Popen([args.binary, "-manifest", str(manifest), str(mount)], stdout=log, stderr=log) + try: + while not os.path.ismount(mount): + if proc.poll() is not None or time.perf_counter()-start > 15: + raise RuntimeError("mount failed") + time.sleep(0.001) + report["mount_ready_seconds"] = time.perf_counter()-start + mounted = mount/"payload.bin" + assert mounted.stat().st_size == size + print(f"Mounted {size/1024**3:.0f} GiB in {report['mount_ready_seconds']:.6f}s", flush=True) + report["direct_full"] = full_pass(source, offset, size, "direct") + (work/"progress.json").write_text(json.dumps(report, indent=2)) + report["fuse_full"] = full_pass(mounted, 0, size, "fuse") + assert report["direct_full"]["sha256"] == report["fuse_full"]["sha256"] == args.expected_sha256 + rng = random.Random(20260918) + ranges = [(0, 8192), (size-8192, 8192), (block-4096, 8192), (size, 1)] + ranges += [(rng.randrange(size-1048576), rng.choice([4096, 65536, 1048576])) for _ in range(128)] + with open(source, "rb", buffering=0) as raw, open(mounted, "rb", buffering=0) as view: + for off, n in ranges: + assert os.pread(raw.fileno(), n, offset+off) == os.pread(view.fileno(), n, off) + with mmap.mmap(view.fileno(), 0, access=mmap.ACCESS_READ) as mm: + for off, n in ranges: + assert mm[off:off+n] == os.pread(raw.fileno(), n, offset+off) + report["random_and_mmap_ranges"] = len(ranges) + def check(part): + with open(source, "rb", buffering=0) as a, open(mounted, "rb", buffering=0) as b: + for off, n in part: + assert os.pread(a.fileno(), n, offset+off) == os.pread(b.fileno(), n, off) + return len(part) + with concurrent.futures.ThreadPoolExecutor(max_workers=4) as pool: + report["concurrent_ranges"] = sum(pool.map(check, [ranges[i::4] for i in range(4)])) + report["source_after"] = stamp(source) + assert before == report["source_after"], "source metadata changed during validation" + with open(source, "rb") as f: + assert f.read(offset) == header+mask, "header changed" + report["mount_process_memory"] = [line.strip() for line in Path(f"/proc/{proc.pid}/status").read_text().splitlines() if line.startswith(("VmRSS:", "VmHWM:"))] + report["output_bytes_excluding_mount"] = sum(p.stat().st_size for p in work.iterdir() if p.is_file()) + report["source_unchanged"] = True + report["validation"] = "passed" + finally: + if os.path.ismount(mount): + subprocess.run(["fusermount3", "-u", str(mount)], check=True) + proc.wait(timeout=15) + log.close() + (work/"report.json").write_text(json.dumps(report, indent=2)) + print(json.dumps(report, indent=2), flush=True) + + +if __name__ == "__main__": + main() diff --git a/test/cachemount/verify_tree_live.py b/test/cachemount/verify_tree_live.py new file mode 100644 index 0000000..21ad10c --- /dev/null +++ b/test/cachemount/verify_tree_live.py @@ -0,0 +1,103 @@ +#!/usr/bin/env python3 +"""Mount an existing repository tree, verify it, and load real cached models. +Outputs and small reference models go to a new isolated work directory. +""" +import argparse +import gc +import json +import os +from pathlib import Path +import subprocess +import time + +from accept_tree import read_payload + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--binary", required=True) + ap.add_argument("--root", required=True) + ap.add_argument("--work", required=True) + ap.add_argument("--mount", default="/mnt/dingo-models") + ap.add_argument("--load-models", action="store_true") + args = ap.parse_args() + work = Path(args.work) + work.mkdir(parents=True, exist_ok=False) + log = open(work/"mount.log", "w") + report = work/"mapping.json" + summary = {"root": args.root, "mount": args.mount, "models": []} + proc = subprocess.Popen([args.binary, "-tree-root", args.root, "-report", str(report), args.mount], stdout=log, stderr=log) + try: + for _ in range(600): + if report.exists() and report.stat().st_size: + break + if proc.poll() is not None: + raise RuntimeError((work/"mount.log").read_text()) + time.sleep(.05) + versions = json.loads(report.read_text()) + summary["versions"] = len(versions) + summary["ready"] = sum(v["status"] == "ready" for v in versions) + summary["unavailable"] = [{"path": v["source_file_path"], "reason": v.get("error")} for v in versions if v["status"] != "ready"] + summary["files"] = sum(len(v["files"]) for v in versions if v["status"] == "ready") + subprocess.run(["python3", str(Path(__file__).with_name("accept_tree.py")), "check", "--report", str(report), "--output", str(work/"acceptance.html")], check=True) + summary["tree_and_content_check"] = "passed; >16 MiB files sampled, smaller files fully checked" + if args.load_models: + import torch + from transformers import AutoTokenizer, AutoModelForCausalLM + targets = [("huggingface", "sshleifer/tiny-gpt2"), ("dingo-local", "hf-tiny-gpt2")] + for ns, repo in targets: + v = next(v for v in versions if v["namespace"] == ns and v["repo"] == repo and v["revision"] == "main") + assert v["status"] == "ready" + reference = work/"reference"/ns/repo + reference.mkdir(parents=True) + source_stamps = {} + for f in v["files"]: + source = Path(f["backing"]) + st = source.stat() + source_stamps[str(source)] = (st.st_size, st.st_mtime_ns) + assert 0 <= f["size"] <= 32*1024*1024, "reference copy restricted to small test models" + out = reference/f["path"] + out.parent.mkdir(parents=True, exist_ok=True) + out.write_bytes(read_payload(source, 0, f["size"])) + mounted = Path(v["source_file_path"]) + tokenizer = AutoTokenizer.from_pretrained(mounted, local_files_only=True) + reference_tokenizer = AutoTokenizer.from_pretrained(reference, local_files_only=True) + model = AutoModelForCausalLM.from_pretrained(mounted, local_files_only=True) + original = AutoModelForCausalLM.from_pretrained(reference, local_files_only=True) + x = tokenizer("Hello, this is a local model directory.", return_tensors="pt") + baseline_x = reference_tokenizer("Hello, this is a local model directory.", return_tensors="pt") + assert torch.equal(x["input_ids"], baseline_x["input_ids"]) + model.eval(); original.eval() + with torch.no_grad(): + assert torch.equal(model(**x).logits, original(**x).logits) + generated = tokenizer.decode(model.generate(**x, max_new_tokens=8, do_sample=False, pad_token_id=tokenizer.eos_token_id)[0]) + model.train() + optimizer = torch.optim.SGD(model.parameters(), lr=0.001) + optimizer.zero_grad() + loss = model(**x, labels=x["input_ids"]).loss + loss.backward(); optimizer.step() + checkpoint = work/"checkpoint"/ns/repo + model.save_pretrained(checkpoint) + restored = AutoModelForCausalLM.from_pretrained(checkpoint, local_files_only=True) + model.eval(); restored.eval() + with torch.no_grad(): + assert torch.equal(model(**x).logits, restored(**x).logits) + for source, before in source_stamps.items(): + st = Path(source).stat() + assert before == (st.st_size, st.st_mtime_ns) + summary["models"].append({"source_file_path": str(mounted), "tokenizer": "passed", "logits_identical_to_plain_reference": True, + "generated_text": generated, "training_loss": loss.item(), "training_step": "passed", "checkpoint_reload": "passed", "source_metadata_unchanged": True}) + del model, original, restored, optimizer, loss + gc.collect() + summary["result"] = "passed" + finally: + if os.path.ismount(args.mount): + subprocess.run(["fusermount3", "-u", args.mount], check=True) + proc.wait(timeout=10) + log.close() + (work/"summary.json").write_text(json.dumps(summary, indent=2, ensure_ascii=False)) + print(json.dumps(summary, indent=2, ensure_ascii=False)) + + +if __name__ == "__main__": + main() diff --git a/test/publish-e2e/run-e2e.sh b/test/publish-e2e/run-e2e.sh index dd8be81..84a894b 100644 --- a/test/publish-e2e/run-e2e.sh +++ b/test/publish-e2e/run-e2e.sh @@ -4,8 +4,6 @@ set -uo pipefail DL=http://127.0.0.1:18090 UP=http://127.0.0.1:18091 -TOKEN=debug-upload-token -HDR="X-Dingo-Upload-Token: $TOKEN" SRC=./src OUT=./out FAIL=0 @@ -32,8 +30,8 @@ ITEMS="" for f in $FILES; do sha=$(sha_of "$SRC/$f"); size=$(size_of "$SRC/$f") code=$(curl -s -o ./tmpdir/up.json -w '%{http_code}' -X POST \ - -H "$HDR" --data-binary "@$SRC/$f" \ - "$UP/api/local-upload/models/dingo-local/batch-demo/main/$f?size=$size&sha256=$sha&defer=true") + --data-binary "@$SRC/$f" \ + "$UP/api/uploads/models/dingo-local?repo=batch-demo&revision=main&path=$f&size=$size&sha256=$sha&defer=true") status=$(python -c "import json,sys;print(json.load(open('./tmpdir/up.json')).get('status',''))" 2>/dev/null) commit=$(python -c "import json,sys;print(json.load(open('./tmpdir/up.json')).get('commit',''))" 2>/dev/null) if [ "$code" = "201" ] && [ "$status" = "staged" ] && [ -z "$commit" ]; then @@ -54,10 +52,10 @@ commit_dirs=$(ls ./repos/api/models/dingo-local/batch-demo/revision 2>/dev/null check "no metadata on disk before publish" "$commit_dirs" "0" echo "=== 3. 一条 curl 完成发布 ===" -pub_code=$(curl -s -o ./tmpdir/pub.json -w '%{http_code}' -X POST -H "$HDR" \ +pub_code=$(curl -s -o ./tmpdir/pub.json -w '%{http_code}' -X POST \ -H 'Content-Type: application/json' \ -d "{\"files\":[$ITEMS]}" \ - "$UP/api/local-publish/models/dingo-local/batch-demo/main") + "$UP/api/publish/models/dingo-local?repo=batch-demo&revision=main") check "publish http status" "$pub_code" "201" BATCH_COMMIT=$(python -c "import json;print(json.load(open('./tmpdir/pub.json'))['commit'])") pub_status=$(python -c "import json;print(json.load(open('./tmpdir/pub.json'))['status'])") @@ -72,8 +70,8 @@ echo "=== 4. 顺序路径：同一组文件逐个即时生效上传到另一个 for f in $FILES; do sha=$(sha_of "$SRC/$f"); size=$(size_of "$SRC/$f") code=$(curl -s -o ./tmpdir/up2.json -w '%{http_code}' -X POST \ - -H "$HDR" --data-binary "@$SRC/$f" \ - "$UP/api/local-upload/models/dingo-local/seq-demo/main/$f?size=$size&sha256=$sha") + --data-binary "@$SRC/$f" \ + "$UP/api/uploads/models/dingo-local?repo=seq-demo&revision=main&path=$f&size=$size&sha256=$sha") [ "$code" = "201" ] || fail "sequential upload $f (http $code)" SEQ_COMMIT=$(python -c "import json;print(json.load(open('./tmpdir/up2.json'))['commit'])") done @@ -120,13 +118,13 @@ printf 'extra payload two' > "$SRC/subdir/extra2.bin" ITEMS2="" for f in extra1.bin subdir/extra2.bin; do sha=$(sha_of "$SRC/$f"); size=$(size_of "$SRC/$f") - curl -s -o /dev/null -X POST -H "$HDR" --data-binary "@$SRC/$f" \ - "$UP/api/local-upload/models/dingo-local/batch-demo/main/$f?size=$size&sha256=$sha&defer=true" + curl -s -o /dev/null -X POST --data-binary "@$SRC/$f" \ + "$UP/api/uploads/models/dingo-local?repo=batch-demo&revision=main&path=$f&size=$size&sha256=$sha&defer=true" [ -n "$ITEMS2" ] && ITEMS2="$ITEMS2," ITEMS2="$ITEMS2{\"path\":\"$f\",\"sha256\":\"$sha\",\"size\":$size}" done -curl -s -o ./tmpdir/pub2.json -X POST -H "$HDR" -d "{\"files\":[$ITEMS2]}" \ - "$UP/api/local-publish/models/dingo-local/batch-demo/main" +curl -s -o ./tmpdir/pub2.json -X POST -d "{\"files\":[$ITEMS2]}" \ + "$UP/api/publish/models/dingo-local?repo=batch-demo&revision=main" APPEND_COMMIT=$(python -c "import json;print(json.load(open('./tmpdir/pub2.json'))['commit'])") append_count=$(python -c "import json;print(json.load(open('./tmpdir/pub2.json'))['fileCount'])") check "追加后清单文件总数" "$append_count" "7" @@ -140,11 +138,11 @@ if [ "$code" = "200" ] && cmp -s "$SRC/config.json" ./tmpdir/old.bin; then pass echo "=== 8. 覆盖：未声明覆盖必须整次拒绝，声明后客户端能拿到新内容 ===" printf '{"model_type":"demo","hidden":1024}' > "$SRC/config.json" sha=$(sha_of "$SRC/config.json"); size=$(size_of "$SRC/config.json") -curl -s -o /dev/null -X POST -H "$HDR" --data-binary "@$SRC/config.json" \ - "$UP/api/local-upload/models/dingo-local/batch-demo/main/config.json?size=$size&sha256=$sha&defer=true" -code=$(curl -s -o ./tmpdir/conflict.json -w '%{http_code}' -X POST -H "$HDR" \ +curl -s -o /dev/null -X POST --data-binary "@$SRC/config.json" \ + "$UP/api/uploads/models/dingo-local?repo=batch-demo&revision=main&path=config.json&size=$size&sha256=$sha&defer=true" +code=$(curl -s -o ./tmpdir/conflict.json -w '%{http_code}' -X POST \ -d "{\"files\":[{\"path\":\"config.json\",\"sha256\":\"$sha\",\"size\":$size}]}" \ - "$UP/api/local-publish/models/dingo-local/batch-demo/main") + "$UP/api/publish/models/dingo-local?repo=batch-demo&revision=main") check "未声明覆盖时的响应码" "$code" "409" ccode=$(python -c "import json;print(json.load(open('./tmpdir/conflict.json'))['code'])") check "未声明覆盖时的错误码" "$ccode" "PUBLISH_OVERWRITE_REQUIRED" @@ -152,9 +150,9 @@ curl -s -o ./tmpdir/still.json "$DL/api/models/dingo-local/batch-demo/revision/m still=$(python -c "import json;print(json.load(open('./tmpdir/still.json'))['sha'])") check "被拒绝的发布没有改变快照标识" "$still" "$APPEND_COMMIT" -curl -s -o ./tmpdir/pub3.json -X POST -H "$HDR" \ +curl -s -o ./tmpdir/pub3.json -X POST \ -d "{\"files\":[{\"path\":\"config.json\",\"sha256\":\"$sha\",\"size\":$size}]}" \ - "$UP/api/local-publish/models/dingo-local/batch-demo/main?overwrite=true" + "$UP/api/publish/models/dingo-local?repo=batch-demo&revision=main&overwrite=true" OVER_COMMIT=$(python -c "import json;print(json.load(open('./tmpdir/pub3.json'))['commit'])") if [ "$OVER_COMMIT" != "$APPEND_COMMIT" ]; then pass "覆盖产生了新的快照标识"; else fail "覆盖没有改变快照标识"; fi code=$(curl -s -o ./tmpdir/new.bin -w '%{http_code}' "$DL/models/dingo-local/batch-demo/resolve/main/config.json") @@ -162,9 +160,9 @@ if [ "$code" = "200" ] && cmp -s "$SRC/config.json" ./tmpdir/new.bin; then pass echo "=== 9. 发布前置校验：清单里有没传完的文件 ===" ghost_sha=$(printf 'never uploaded' | sha256sum | cut -d' ' -f1) -code=$(curl -s -o ./tmpdir/ghost.json -w '%{http_code}' -X POST -H "$HDR" \ +code=$(curl -s -o ./tmpdir/ghost.json -w '%{http_code}' -X POST \ -d "{\"files\":[{\"path\":\"ghost.bin\",\"sha256\":\"$ghost_sha\",\"size\":14}]}" \ - "$UP/api/local-publish/models/dingo-local/batch-demo/main") + "$UP/api/publish/models/dingo-local?repo=batch-demo&revision=main") check "缺内容时的响应码" "$code" "409" gcode=$(python -c "import json;print(json.load(open('./tmpdir/ghost.json'))['code'])") check "缺内容时的错误码" "$gcode" "PUBLISH_CONTENT_NOT_READY" @@ -173,11 +171,11 @@ case "$gmsg" in *ghost.bin*) pass "错误信息指出了缺失路径";; *) fail echo "=== 10. 数据集类型 ===" sha=$(sha_of "$SRC/README.md"); size=$(size_of "$SRC/README.md") -curl -s -o /dev/null -X POST -H "$HDR" --data-binary "@$SRC/README.md" \ - "$UP/api/local-upload/datasets/dingo-local/ds-demo/main/data/train.md?size=$size&sha256=$sha&defer=true" -curl -s -o ./tmpdir/pubds.json -X POST -H "$HDR" \ +curl -s -o /dev/null -X POST --data-binary "@$SRC/README.md" \ + "$UP/api/uploads/datasets/dingo-local?repo=ds-demo&revision=main&path=data/train.md&size=$size&sha256=$sha&defer=true" +curl -s -o ./tmpdir/pubds.json -X POST \ -d "{\"files\":[{\"path\":\"data/train.md\",\"sha256\":\"$sha\",\"size\":$size}]}" \ - "$UP/api/local-publish/datasets/dingo-local/ds-demo/main" + "$UP/api/publish/datasets/dingo-local?repo=ds-demo&revision=main" DS_COMMIT=$(python -c "import json;print(json.load(open('./tmpdir/pubds.json')).get('commit',''))") if [ -n "$DS_COMMIT" ]; then pass "数据集发布成功"; else fail "数据集发布失败: $(cat ./tmpdir/pubds.json)"; fi code=$(curl -s -o ./tmpdir/ds.bin -w '%{http_code}' "$DL/datasets/dingo-local/ds-demo/resolve/main/data/train.md") diff --git a/test/remote-compat/compare.py b/test/remote-compat/compare.py new file mode 100644 index 0000000..9c29270 --- /dev/null +++ b/test/remote-compat/compare.py @@ -0,0 +1,198 @@ +"""Compare two built DingoSpeed binaries against a deterministic upstream. + +Only writes to --output; servers bind loopback and processes are stopped on exit. +Usage: python compare.py --baseline old.exe --candidate new.exe --output result-dir +""" +import argparse +import hashlib +import http.client +import json +import os +from pathlib import Path +import shutil +import socket +import subprocess +import threading +import time +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from urllib.parse import urlsplit + +PAYLOAD = b"remote fixture bytes\n" * 200 +SHA = hashlib.sha256(PAYLOAD).hexdigest() +COMMIT = "a" * 40 + + +class Upstream(BaseHTTPRequestHandler): + def log_message(self, *_): + pass + + def do_HEAD(self): + self.do_GET() + + def do_POST(self): + self.rfile.read(int(self.headers.get("Content-Length", "0"))) + self.do_GET() + + def do_GET(self): + path = urlsplit(self.path).path + code = 200 + headers = {"Content-Type": "application/json"} + if "/missing/" in path: + code, body = 404, b'{"error":"missing repository"}' + elif "/denied/" in path: + code, body = 403, b'{"error":"access denied"}' + elif "/resolve/" in path or path.endswith("/repo"): + body = PAYLOAD + headers.update({"Content-Type": "application/octet-stream", "ETag": '"' + SHA + '"', "X-Repo-Commit": COMMIT, "Accept-Ranges": "bytes"}) + if self.headers.get("Range"): + start, end = self.headers["Range"].removeprefix("bytes=").split("-") + start, end = int(start), min(int(end) if end else len(body)-1, len(body)-1) + headers["Content-Range"] = f"bytes {start}-{end}/{len(body)}" + body, code = body[start:end+1], 206 + elif path.startswith("/api/v1/"): + data = {"Files": [{"Type": "blob", "Path": "weights.bin", "Name": "weights.bin", "Size": len(PAYLOAD), "Sha256": SHA}], "RevisionMap": {"Branches": [{"Revision": "master"}], "Tags": []}} + body = json.dumps({"Code": 200, "Success": True, "Data": data}).encode() + elif "/tree/" in path or "/paths-info/" in path: + body = json.dumps([{"type": "file", "path": "weights.bin", "oid": SHA, "size": len(PAYLOAD), "lfs": {"oid": SHA, "size": len(PAYLOAD)}}]).encode() + elif path.endswith("/refs"): + body = json.dumps({"branches": [{"name": "main", "ref": "refs/heads/main", "targetCommit": COMMIT}], "tags": [], "converts": []}).encode() + else: + body = json.dumps({"id": "team/demo", "sha": COMMIT, "siblings": [{"rfilename": "weights.bin"}], "usedStorage": len(PAYLOAD)}).encode() + headers["Content-Length"] = str(len(body)) + self.send_response(code) + for key, value in headers.items(): + self.send_header(key, value) + self.end_headers() + if self.command != "HEAD": + self.wfile.write(body) + + +def port(): + with socket.socket() as sock: + sock.bind(("127.0.0.1", 0)) + return sock.getsockname()[1] + + +def request(base_port, method, path, headers=None): + conn = http.client.HTTPConnection("127.0.0.1", base_port, timeout=20) + try: + conn.request(method, path, headers=headers or {}) + response = conn.getresponse() + body = response.read() + selected = {k.lower(): v for k, v in response.getheaders() if k.lower() in {"content-type", "content-length", "content-range", "etag", "x-repo-commit", "accept-ranges"}} + return {"status": response.status, "headers": selected, "sha256": hashlib.sha256(body).hexdigest(), "size": len(body)} + finally: + conn.close() + + +def start(binary, directory, upstream_port, repos=None, online=True): + directory.mkdir(parents=True, exist_ok=True) + read_port, upload_port = port(), port() + config = {"server": {"host": "127.0.0.1", "port": read_port, "online": online, "repos": str(repos or directory / "repos"), "hfScheme": "http", "hfNetLoc": f"127.0.0.1:{upstream_port}", "bpHfNetLoc": f"127.0.0.1:{upstream_port}"}, "tokenBucketLimit": {"handlerCapacity": 50}, "scheduler": {"mode": "standalone"}, "upload": {"host": "127.0.0.1", "port": upload_port, "namespace": "dingo-local"}, "download": {"blockSize": 1048576, "respChunkSize": 1024, "remoteFileBufferSize": 8192, "goroutineMaxNumPerFile": 2, "reqTimeout": 10}, "cache": {"defaultExpiration": 30, "cleanupInterval": 60, "readBlock": {"collectTimePeriod": 5, "prefetchMemoryUsedThreshold": 90, "prefetchBlocks": 8, "prefetchBlockTTL": 30}}, "retry": {"attempts": 1, "delay": 1}, "modelscope": {"officialBaseURL": f"http://127.0.0.1:{upstream_port}", "chunkSize": 1024, "maxRetry": 1, "retryDelay": 1}} + path = directory / "config.json" + path.write_text(json.dumps(config), encoding="utf-8") + log = (directory / "server.log").open("w", encoding="utf-8") + proc = subprocess.Popen([str(binary), "-config", str(path)], cwd=directory, stdout=log, stderr=subprocess.STDOUT, creationflags=subprocess.CREATE_NO_WINDOW if os.name == "nt" else 0) + for _ in range(100): + if proc.poll() is not None: + log.close() + raise RuntimeError(f"server exited: {directory / 'server.log'}") + try: + assert request(read_port, "GET", "/info")["status"] == 200 + return proc, log, read_port + except (OSError, http.client.HTTPException): + time.sleep(.1) + proc.terminate() + proc.wait() + log.close() + raise RuntimeError("server startup timed out") + + +def main(): + parser = argparse.ArgumentParser() + for name in ("baseline", "candidate", "output"): + parser.add_argument("--" + name, required=True, type=Path) + parser.add_argument("--spinfield", type=Path, help="Also run console HTTP upload integration and HF SDK downloads") + args = parser.parse_args() + output = args.output.resolve() + output.mkdir(parents=True, exist_ok=False) + upstream = ThreadingHTTPServer(("127.0.0.1", 0), Upstream) + threading.Thread(target=upstream.serve_forever, daemon=True).start() + processes = [] + report = [] + try: + ports = [] + for name, binary in (("baseline", args.baseline), ("candidate", args.candidate)): + running = start(binary.resolve(), output / name, upstream.server_port) + processes.append(running) + ports.append(running[2]) + cases = [(method, path, headers) for method, path, headers in [ + ("GET", "/api/models/team/demo/revision/main", {}), + ("HEAD", "/api/models/team/demo/revision/main", {}), + ("GET", "/api/models/team/demo/tree/main", {}), + ("GET", "/api/models/team/demo/refs", {}), + ("HEAD", "/team/demo/resolve/main/weights.bin", {}), + ("GET", "/team/demo/resolve/main/weights.bin", {}), + ("GET", "/team/demo/resolve/main/weights.bin", {"Range": "bytes=7-39"}), + ("GET", "/api/datasets/team/demo/revision/main", {}), + ("GET", "/datasets/team/demo/resolve/main/weights.bin", {}), + ("GET", "/api/models/missing/demo/revision/main", {}), + ("GET", "/api/models/denied/demo/revision/main", {}), + ("GET", "/api/v1/models/team/demo", {}), + ("GET", "/api/v1/models/team/demo/repo?Revision=master&FilePath=weights.bin", {}), + ("GET", "/api/v1/models/team/demo/repo?Revision=master&FilePath=weights.bin", {"Range": "bytes=7-39"}), + ]] + for method, path, headers in cases: + results = [request(p, method, path, headers) for p in ports] + expected_status = 404 if "/missing/" in path else 403 if "/denied/" in path else 206 if "Range" in headers or path.startswith("/api/v1/") and "/repo?" in path else 200 + if results[0]["status"] != expected_status: + raise AssertionError(f"fixture did not exercise expected baseline behavior: {method} {path}: {results}") + report.append({"method": method, "path": path, "requestHeaders": headers, "baseline": results[0], "candidate": results[1], "equal": results[0] == results[1]}) + if args.spinfield: + config = json.loads((output / "candidate/config.json").read_text()) + env = os.environ.copy() + env["DINGOSPEED_E2E_UPLOAD_BASE"] = f"http://127.0.0.1:{config['upload']['port']}" + env["DINGOSPEED_E2E_DOWNLOAD_BASE"] = f"http://127.0.0.1:{ports[1]}" + result = subprocess.run(["go", "test", "./modelfleet/internal/modelasset/api", "-run", "TestDingoSpeedConsoleUploadNamespaceIsolation", "-count=1", "-v"], cwd=args.spinfield.resolve(), env=env, capture_output=True, text=True, timeout=180) + (output / "console-http-test.log").write_text(result.stdout + result.stderr, encoding="utf-8") + if result.returncode: + raise RuntimeError(result.stdout + result.stderr) + from huggingface_hub import snapshot_download + sdk_results = [] + for ns in ("alice", "bob", "datacanvas"): + repo = next((output / "candidate/repos/files/models/dingo-local" / ns).iterdir()).name + dest = output / "sdk-download" / ns + snapshot_download(repo_id=ns+"/"+repo, revision="main", endpoint=env["DINGOSPEED_E2E_DOWNLOAD_BASE"]+"/dingo-local", local_dir=dest, token=False) + data = (dest / "weights.bin").read_bytes() + assert data == (ns+" owns these bytes").encode() + sdk_results.append({"namespace": ns, "repo": repo, "sha256": hashlib.sha256(data).hexdigest()}) + (output / "sdk-results.json").write_text(json.dumps(sdk_results, indent=2), encoding="utf-8") + # Transfer unmodified baseline caches, with no new repository descriptors. + baseline, log, _ = processes.pop(0) + baseline.terminate() + baseline.wait() + log.close() + old_cache = output / "legacy-cache" + shutil.copytree(output / "baseline/repos", old_cache) + running = start(args.candidate.resolve(), output / "legacy-reader", upstream.server_port, old_cache, False) + processes.append(running) + for method, path, headers in (cases[0], cases[4], cases[5], cases[6], cases[8], cases[12]): + value = request(running[2], method, path, headers) + expected = next(row["baseline"] for row in report if row["method"] == method and row["path"] == path and row["requestHeaders"] == headers) + report.append({"legacyCache": True, "method": method, "path": path, "baseline": expected, "candidate": value, "equal": value == expected}) + (output / "report.json").write_text(json.dumps(report, indent=2), encoding="utf-8") + failed = [row for row in report if not row["equal"]] + print(json.dumps({"cases": len(report), "passed": len(report)-len(failed), "report": str(output / "report.json"), "failures": failed}, indent=2)) + if failed: + raise SystemExit(1) + finally: + for proc, log, _ in processes: + proc.terminate() + proc.wait(timeout=10) + log.close() + upstream.shutdown() + upstream.server_close() + + +if __name__ == "__main__": + main() diff --git a/test/uploadbench/README.md b/test/uploadbench/README.md index 454773c..12b9837 100644 --- a/test/uploadbench/README.md +++ b/test/uploadbench/README.md @@ -1,17 +1,21 @@ # uploadbench —— 本地单文件上传接口压测工具 -针对 `POST /api/local-upload/:repoType/:org/:repo/:revision/*` 的压测客户端与实验脚本。 +针对 `POST /api/uploads/:repoType/:namespace?repo=...&revision=...&path=...` 的压测客户端与实验脚本。 结论见 [`docs/upload-load-test-report.md`](../../docs/upload-load-test-report.md)。 ## 构建 ```bash go build -o bin/dingospeed-test.exe ./cmd -go build -o test/uploadbench/uploadbench.exe ./test/uploadbench +cd test/uploadbench +go build -o uploadbench.exe . +cd ../.. ``` `test/uploadbench` 是独立 module，不会被主 module 的 `go build ./...` 编译进去。 +无需新增服务 token 配置。命名空间参数为 `-namespace`，`-repo` 可包含多级目录。 + ## 跑全套实验 ```bash diff --git a/test/uploadbench/config-bench.tmpl.yaml b/test/uploadbench/config-bench.tmpl.yaml index 651db19..b5d6442 100644 --- a/test/uploadbench/config-bench.tmpl.yaml +++ b/test/uploadbench/config-bench.tmpl.yaml @@ -14,7 +14,6 @@ server: upload: host: 127.0.0.1 port: 18091 - token: bench-token namespace: dingo-local concurrentLimit: __LIMIT__ stagingRetentionHours: 168 @@ -71,4 +70,4 @@ diskClean: enabled: false cacheSizeLimit: 41781441855488 cacheCleanStrategy: "LRU" - collectTimePeriod: 1 + collectTimePeriod: 24 diff --git a/test/uploadbench/main.go b/test/uploadbench/main.go index 3fb24bd..c65b2b0 100644 --- a/test/uploadbench/main.go +++ b/test/uploadbench/main.go @@ -1,4 +1,4 @@ -// uploadbench 是针对 dingospeed 本地单文件上传接口 POST /api/local-upload/... 的压测客户端。 +// uploadbench 是针对 dingospeed 本地单文件上传接口 POST /api/uploads/:repoType/:namespace?repo&revision&path 的压测客户端。 // 只依赖标准库，内容按种子确定性生成，不落磁盘，避免客户端磁盘 IO 干扰服务端指标。 package main @@ -12,6 +12,7 @@ import ( "math" "net" "net/http" + "net/url" "os" "sort" "strings" @@ -133,17 +134,16 @@ type sample struct { } type target struct { - base string - token string - repoType string - org string - repo string - revision string + base string + repoType string + namespace string + repo string + revision string } func (t target) upload(cli *http.Client, filePath string, seed uint64, size int64, overwrite bool, body io.Reader) (int, string, string, error) { - u := fmt.Sprintf("%s/api/local-upload/%s/%s/%s/%s/%s?size=%d&sha256=%s", - t.base, t.repoType, t.org, t.repo, t.revision, filePath, size, contentSha(seed, size)) + q := url.Values{"repo": {t.repo}, "revision": {t.revision}, "path": {filePath}, "size": {fmt.Sprint(size)}, "sha256": {contentSha(seed, size)}} + u := strings.TrimRight(t.base, "/") + "/api/uploads/" + url.PathEscape(t.repoType) + "/" + url.PathEscape(t.namespace) + "?" + q.Encode() if overwrite { u += "&overwrite=true" } @@ -152,7 +152,6 @@ func (t target) upload(cli *http.Client, filePath string, seed uint64, size int6 return 0, "", "", err } req.ContentLength = size - req.Header.Set(uploadTokenHeader, t.token) req.Header.Set("Content-Type", "application/octet-stream") resp, err := cli.Do(req) if err != nil { @@ -173,8 +172,6 @@ func (t target) upload(cli *http.Client, filePath string, seed uint64, size int6 return resp.StatusCode, code, upStatus, nil } -const uploadTokenHeader = "X-Dingo-Upload-Token" - // retry429 > 0 时，客户端遇到 429 会退避后重试同一个文件，模拟真实上传工具的行为。 var retry429 time.Duration @@ -300,32 +297,31 @@ func summarize(scenario, label string, conc int, size int64, samples []sample, w func main() { var ( - base = flag.String("base", "http://127.0.0.1:8091", "上传服务地址") - token = flag.String("token", "bench-token", "上传 token") - repoType = flag.String("repoType", "datasets", "models|datasets") - org = flag.String("org", "dingo-local", "命名空间") - repo = flag.String("repo", "bench", "仓库名") - revision = flag.String("revision", "main", "版本") - scenario = flag.String("scenario", "closed", "closed|dataset|slowloris|idem") - label = flag.String("label", "", "结果标签") - conc = flag.Int("c", 4, "并发数") - total = flag.Int("n", 32, "总请求数") - size = flag.Int64("size", 4<<20, "单文件字节数") - prefix = flag.String("prefix", "f", "文件名前缀") - out = flag.String("out", "", "逐请求样本 JSONL 输出路径") - holders = flag.Int("holders", 4, "slowloris 场景占位连接数") - holdRate = flag.Int64("holdRate", 4096, "slowloris 每个占位连接的字节/秒") - holdSize = flag.Int64("holdSize", 8<<20, "slowloris 占位文件字节数") - probeFor = flag.Duration("probeFor", 20*time.Second, "slowloris 探测持续时间") - warmup = flag.Bool("warmup", true, "先发一次请求预热连接与目录") - shards = flag.Int("shardRepos", 1, "dataset 场景把文件散列到 N 个仓库，1 表示全部写同一个仓库") - retry = flag.Duration("retry429", 0, "closed 场景遇到 429 的退避重试间隔，0 表示不重试") + base = flag.String("base", "http://127.0.0.1:8091", "上传服务地址") + repoType = flag.String("repoType", "datasets", "models|datasets") + namespace = flag.String("namespace", "dingo-local", "命名空间") + repo = flag.String("repo", "bench", "仓库名") + revision = flag.String("revision", "main", "版本") + scenario = flag.String("scenario", "closed", "closed|dataset|slowloris|idem") + label = flag.String("label", "", "结果标签") + conc = flag.Int("c", 4, "并发数") + total = flag.Int("n", 32, "总请求数") + size = flag.Int64("size", 4<<20, "单文件字节数") + prefix = flag.String("prefix", "f", "文件名前缀") + out = flag.String("out", "", "逐请求样本 JSONL 输出路径") + holders = flag.Int("holders", 4, "slowloris 场景占位连接数") + holdRate = flag.Int64("holdRate", 4096, "slowloris 每个占位连接的字节/秒") + holdSize = flag.Int64("holdSize", 8<<20, "slowloris 占位文件字节数") + probeFor = flag.Duration("probeFor", 20*time.Second, "slowloris 探测持续时间") + warmup = flag.Bool("warmup", true, "先发一次请求预热连接与目录") + shards = flag.Int("shardRepos", 1, "dataset 场景把文件散列到 N 个仓库，1 表示全部写同一个仓库") + retry = flag.Duration("retry429", 0, "closed 场景遇到 429 的退避重试间隔，0 表示不重试") ) flag.Parse() retry429 = *retry shardRepos = *shards - t := target{base: *base, token: *token, repoType: *repoType, org: *org, repo: *repo, revision: *revision} + t := target{base: *base, repoType: *repoType, namespace: *namespace, repo: *repo, revision: *revision} lbl := *label if lbl == "" { lbl = fmt.Sprintf("%s-c%d-n%d-%dB", *scenario, *conc, *total, *size)