From 25afbe78895189f5c345913e3db0c025027e6a9f Mon Sep 17 00:00:00 2001 From: taoky Date: Thu, 10 Sep 2026 03:23:25 +0800 Subject: [PATCH 1/3] feat: support zero-downtime upgrades across process generations Inherit listeners while older generations drain existing sessions. Coordinate admission and FIFO queues through SQLite, and aggregate connection status, metrics, and log rotation across generations. Add configurable drain deadlines, systemd MAINPID handoff, and package upgrade integration. Cover multi-generation upgrades, failure recovery, and real rsync transfers with process tests. Assisted-by: OpenAI Codex (GPT-6) --- README.md | 28 +- assets/config.example.toml | 14 +- assets/logrotate.conf | 2 +- assets/postinst.sh | 10 +- assets/rsync-proxy.service | 5 +- cmd/cmd.go | 12 +- cmd/daemon_linux.go | 592 +++++++++++++++++++++++++++++++++ cmd/daemon_linux_test.go | 580 ++++++++++++++++++++++++++++++++ cmd/daemon_other.go | 12 + cmd/upgrade.go | 82 +++++ contrib/grafana/dashboard.json | 18 +- go.mod | 15 +- go.sum | 52 ++- pkg/server/admission.go | 157 +++++++++ pkg/server/config.go | 20 +- pkg/server/handoff.go | 95 ++++++ pkg/server/metrics.go | 16 +- pkg/server/server.go | 254 ++++++++++---- pkg/server/snapshot.go | 168 ++++++++++ pkg/server/snapshot_test.go | 59 ++++ pkg/server/utils.go | 28 +- pkg/server/utils_test.go | 23 ++ pkg/state/lock.go | 34 ++ pkg/state/store.go | 351 +++++++++++++++++++ pkg/state/store_test.go | 246 ++++++++++++++ test/systemd/upgrade.py | 157 +++++++++ 26 files changed, 2925 insertions(+), 105 deletions(-) create mode 100644 cmd/daemon_linux.go create mode 100644 cmd/daemon_linux_test.go create mode 100644 cmd/daemon_other.go create mode 100644 cmd/upgrade.go create mode 100644 pkg/server/admission.go create mode 100644 pkg/server/handoff.go create mode 100644 pkg/server/snapshot.go create mode 100644 pkg/server/snapshot_test.go create mode 100644 pkg/state/lock.go create mode 100644 pkg/state/store.go create mode 100644 pkg/state/store_test.go create mode 100644 test/systemd/upgrade.py diff --git a/README.md b/README.md index 8b1726f..ee52a0e 100644 --- a/README.md +++ b/README.md @@ -55,6 +55,30 @@ systemctl enable --now rsync-proxy.service cp logrotate.conf /etc/logrotate.d/ ``` +### 零停机更新（Linux） + +运行中的版本具备升级协议后，可以先原子替换安装路径上的二进制，再执行： + +```shell +sudo install -m 0755 rsync-proxy.new /usr/bin/rsync-proxy.next +sudo mv /usr/bin/rsync-proxy.next /usr/bin/rsync-proxy +sudo /usr/bin/rsync-proxy --config=/etc/rsync-proxy/config.toml upgrade --timeout=30s +``` + +`upgrade` 等待新进程就绪并完成监听交接后返回，不等待旧连接结束。已有明文/TLS 传输、握手和排队连接留在原进程；新连接由新进程接收。可以在旧连接未结束时继续升级，形成多代共存。`reload` 仍然只重读配置；`systemctl restart` 是普通重启，不是零停机更新入口。 + +- `[proxy].state_dir` 默认为 `/run/rsync-proxy`，保存各代共享的 SQLite 运行状态。手动运行时须使用可写目录，各实例须使用独立目录；不要放在 NFS 等网络文件系统上，也不要在进程运行期间删除数据库、WAL 或锁文件。 +- `[proxy].upgrade_drain_timeout` 单位为秒，默认 `0`。非零值从该代停止接收新连接时起计算，到期关闭剩余连接，包括握手和排队。采用本次成功升级的新配置值；后续更新不会改变或重置更早各代的期限。默认不增加期限，原有 idle/max-duration/throughput 规则继续生效。 +- 所有代合计执行上游并发、排队和单 IP 限额，排队按全局 FIFO 晋升。调低限额不会主动断开已有连接；等待占用降到新限额以下再放行。旧连接继续使用已经选定的上游目标。 +- 热更新要求 `listen`、`listen_tls`、`listen_http` 和状态目录保持一致。新程序、配置、证书或协议不兼容，或就绪超过 `--timeout`，会报告失败并保留旧进程。数据库不可用时不绕过配额放行，已建立的传输继续运行。 +- `connections`、`/status` 和 `/metrics` 汇总各代，逐连接数据新增 `generation` 标识；Go/runtime 指标属于当前主进程。跨代快照每秒发布，正常退出前发布最终值，业务累计指标保留已退出代的贡献；异常退出可能损失最后一次未发布的统计。一次完全停止后的启动会重置运行状态。 +- 使用配套 systemd unit，使主 PID 在升级后指向新进程；需要支持 `sd_notify` barrier 的 systemd（246 或更新版本）。普通停止仍受 systemd 的停止超时约束。旧进程仍有连接时会继续占用内存和 FD，可通过退出期限限制。 +- 日志轮转使用 `rsync-proxy --config=... reopen-logs`，让所有存活的代重新打开日志。升级与日志重开命令读取配置中的管理地址，显式 `--host` 可覆盖。 + +**首次从不支持该协议的旧版迁入需要一次普通重启。** deb 升级会对正在运行的服务调用 `upgrade`，失败不会自动强制重启；首次迁入或不兼容升级时，检查错误后手动 `systemctl restart rsync-proxy`，再完成包配置。进程崩溃、主机重启、主动停止和配置的连接超时不属于零停机保证。 + +开发验证：`go test -race ./...` 包含真实子进程的多代交接、SQLite 并发调度及故障测试。真实用户级 systemd 验证可运行 `python3 test/systemd/upgrade.py --binary /absolute/path/to/rsync-proxy`；该脚本只创建临时测试 unit、回环监听和临时目录。 + ### 安装 fail2ban filter ```shell @@ -90,7 +114,7 @@ cp fail2ban/filter.d/* /etc/fail2ban/filter.d/ | `relay_idle_timeout` | int 秒 | 0 | relay 阶段双向无 I/O 多久后关闭连接。语义同 rsyncd `timeout`。 | `600` | | `relay_max_duration` | int 秒 | 0 | relay 阶段总时长上限。超时关闭，rsync 客户端通常会重连续传。 | `14400`（4h） | | `tcp_keepalive` | int 秒 | 0 | 客户端连接和上游连接的 TCP keepalive 周期，0 沿用 OS 默认（通常 ~2h）。 | `120` | -| `per_ip_max_active_connections` | int | 0 | 单 IP 对单上游最多并发 relay 连接数，proxy-wide 默认值。NAT/校园出口 IP 时取值需放宽。 | `4` | +| `per_ip_max_active_connections` | int | 0 | 单 IP 对单上游最多同时持有的连接数（含排队），proxy-wide 默认值。NAT/校园出口 IP 时取值需放宽。 | `4` | | `dial_timeout` | int 秒 | 0 | 拨号上游的超时；0 沿用内核 SYN 重试（~75s）。 | `5`（LAN） | | `min_throughput_bytes` | int64 字节 | 0 | relay 阶段最近 `min_throughput_window` 秒内累计收发须 ≥ 此值，否则视作慢速吸血并关闭。0 关闭整组检查。 | `1048576` | | `min_throughput_window` | int 秒 | 60 | 上述滑动窗口长度。 | `60` | @@ -109,7 +133,7 @@ cp fail2ban/filter.d/* /etc/fail2ban/filter.d/ | `discover_modules` | bool | false | 启动/reload 时自动从上游拉 module 列表。上游不可达会导致启动失败。与 `modules` 二选一。 | | `use_proxy_protocol` | bool | false | 与上游通信时附加 PROXY protocol 头，便于上游记录真实客户端 IP。需上游 rsyncd 启用 PROXY protocol 支持。 | | `max_active_connections` | int | 0 | 该上游最大并发 relay 连接数，0 表示不限制。 | -| `max_queued_connections` | int | 0 | 达到上限后排队的最大长度，0 表示不排队。 | +| `max_queued_connections` | int | 0 | 达到上限后排队的最大长度，0 表示不限排队长度。 | | `per_ip_max_active_connections` | int | 0 | 覆盖 `[proxy]` 中的同名值；0 表示继承 proxy-wide 默认。 | # 监控 diff --git a/assets/config.example.toml b/assets/config.example.toml index 94bc8a6..b8a9eb1 100644 --- a/assets/config.example.toml +++ b/assets/config.example.toml @@ -11,6 +11,15 @@ error_log = "/var/log/rsync-proxy/error.log" tls_cert_file = "/etc/rsync-proxy/tls/server.crt" tls_key_file = "/etc/rsync-proxy/tls/server.key" +# Shared runtime state for all process generations (Linux). +# Use a local directory writable by the service user, never NFS. +#state_dir = "/run/rsync-proxy" + +# Maximum seconds to drain a retiring generation after a successful upgrade. +# 0 (default) adds no deadline; existing connection protection rules still apply. +# Each retiring generation keeps the timeout assigned when it stopped accepting. +#upgrade_drain_timeout = 0 + motd = "Served by rsync-proxy (https://github.com/ustclug/rsync-proxy)" # Idle timeout (seconds) applied during the bidirectional relay phase @@ -37,8 +46,8 @@ motd = "Served by rsync-proxy (https://github.com/ustclug/rsync-proxy)" # is 120 (2 minutes). #tcp_keepalive = 0 -# Default per-IP per-upstream concurrent active connection cap. -# Limits how many simultaneous active relay connections a single +# Default per-IP per-upstream connection cap (active AND queued). +# Limits how many simultaneous connections a single # client IP may have to any one upstream, preventing a single client # from monopolizing upstream slots. 0 (the default) disables the # limit. A typical value for public mirrors is 4. Each upstream may @@ -75,6 +84,7 @@ motd = "Served by rsync-proxy (https://github.com/ustclug/rsync-proxy)" address = "127.0.0.1:1234" modules = ["foo"] max_active_connections = 60 +# 0 means unlimited queue length. Limits apply across all live generations. max_queued_connections = 60 # Override the proxy-wide per_ip_max_active_connections for this # upstream. 0 means inherit the proxy-wide default. diff --git a/assets/logrotate.conf b/assets/logrotate.conf index 9e5cd17..afed21f 100644 --- a/assets/logrotate.conf +++ b/assets/logrotate.conf @@ -9,6 +9,6 @@ missingok postrotate - systemctl reload rsync-proxy.service + /usr/bin/rsync-proxy --config=/etc/rsync-proxy/config.toml reopen-logs endscript } diff --git a/assets/postinst.sh b/assets/postinst.sh index 134a537..1403ea6 100755 --- a/assets/postinst.sh +++ b/assets/postinst.sh @@ -40,11 +40,15 @@ if [ "$1" = "configure" ] || [ "$1" = "abort-upgrade" ] || [ "$1" = "abort-decon deb-systemd-helper update-state rsync-proxy.service >/dev/null || true fi - # Restart only if it was already started + # Upgrade only a running service. Never turn a failed handoff into a + # disruptive restart; pre-handoff versions need one manual restart. if [ -d /run/systemd/system ]; then systemctl --system daemon-reload >/dev/null || true - if [ -n "$2" ]; then - deb-systemd-invoke try-restart rsync-proxy.service >/dev/null || true + if [ -n "$2" ] && systemctl --system is-active --quiet rsync-proxy.service; then + if ! /usr/bin/rsync-proxy --config=/etc/rsync-proxy/config.toml upgrade; then + echo "rsync-proxy: hot upgrade failed; the existing process was retained. Versions without handoff support need one manual systemctl restart rsync-proxy." >&2 + exit 1 + fi fi fi fi diff --git a/assets/rsync-proxy.service b/assets/rsync-proxy.service index 012b1a3..8a106c8 100644 --- a/assets/rsync-proxy.service +++ b/assets/rsync-proxy.service @@ -2,10 +2,13 @@ Description=Rsync Proxy After=network.target ConditionPathExists=/etc/rsync-proxy/config.toml +StartLimitIntervalSec=10s [Service] +Type=notify +NotifyAccess=all +KillMode=control-group Restart=on-failure -StartLimitInterval=10s ExecStart=/usr/bin/rsync-proxy --config=/etc/rsync-proxy/config.toml ExecReload=/usr/bin/rsync-proxy --config=/etc/rsync-proxy/config.toml reload diff --git a/cmd/cmd.go b/cmd/cmd.go index a119a22..46825da 100644 --- a/cmd/cmd.go +++ b/cmd/cmd.go @@ -142,6 +142,7 @@ func SendConnectionsRequest(addr string, stdout, stderr io.Writer) error { var result struct { Connections []struct { Index int `json:"index"` + Generation string `json:"generation"` RemoteAddr string `json:"remote"` Module string `json:"module"` Upstream string `json:"upstream"` @@ -176,6 +177,7 @@ func SendConnectionsRequest(addr string, stdout, stderr io.Writer) error { }), tablewriter.WithHeaderAutoFormat(tw.Off), tablewriter.WithAlignment(tw.Alignment{ + tw.AlignDefault, // Generation tw.AlignRight, // Index tw.AlignRight, // RemoteAddr tw.AlignDefault, // Module @@ -185,9 +187,10 @@ func SendConnectionsRequest(addr string, stdout, stderr io.Writer) error { tw.AlignRight, // SentBytes }), ) - table.Header("Index", "Remote", "Module", "Upstream", "Connected", "Received", "Sent") + table.Header("Generation", "Index", "Remote", "Module", "Upstream", "Connected", "Received", "Sent") for _, conn := range result.Connections { _ = table.Append([]string{ + conn.Generation, strconv.Itoa(conn.Index), conn.RemoteAddr, conn.Module, @@ -357,10 +360,7 @@ func New() *cobra.Command { return fmt.Errorf("load config: %w", err) } - if err := s.Listen(); err != nil { - return fmt.Errorf("server listen: %w", err) - } - return s.Run() + return runDaemon(s) }, SilenceUsage: true, } @@ -372,6 +372,8 @@ func New() *cobra.Command { c.AddCommand( newConnectionsCmd(), newReloadCmd(), + newUpgradeCmd(s), + newReopenLogsCmd(s), newUpstreamModulesCmd(s), newVersionCmd(), ) diff --git a/cmd/daemon_linux.go b/cmd/daemon_linux.go new file mode 100644 index 0000000..6426320 --- /dev/null +++ b/cmd/daemon_linux.go @@ -0,0 +1,592 @@ +//go:build linux + +package cmd + +import ( + "context" + "crypto/rand" + "encoding/hex" + "encoding/json" + "errors" + "fmt" + "io" + "log" + "net" + "net/http" + "os" + "os/exec" + "os/signal" + "path/filepath" + "strconv" + "strings" + "sync" + "sync/atomic" + "syscall" + "time" + + "golang.org/x/sys/unix" + + "github.com/ustclug/rsync-proxy/pkg/server" + "github.com/ustclug/rsync-proxy/pkg/state" +) + +const handoffEnv = "RSYNC_PROXY_HANDOFF" + +type handoffMessage struct { + Protocol int + Parent string + StateDir string + Listeners [3]string + Stage string + ID string + DrainTimeout time.Duration +} + +type daemon struct { + ready atomic.Bool + s *server.Server + store *state.Store + id, dir, binary string + mu sync.Mutex // serializes reload and upgrade through the commit point + private *http.Server + privateListener net.Listener + lock *os.File + runResult chan error + done chan struct{} + backgroundDone chan struct{} + drainOnce sync.Once + drainTimer *time.Timer + stopAdmission func() +} + +func runDaemon(s *server.Server) error { + dir := s.StateDir + if dir == "" { + dir = "/run/rsync-proxy" + } + dir, err := filepath.Abs(dir) + if err != nil { + return err + } + s.StateDir = dir + if err = os.MkdirAll(dir, 0700); err != nil { + return err + } + startup, err := state.Lock(filepath.Join(dir, "startup.lock")) + if err != nil { + return fmt.Errorf("another startup is in progress: %w", err) + } + defer startup.Close() + store, err := state.Open(dir) + if err != nil { + return err + } + defer store.Close() + binary, err := os.Executable() + if err != nil { + return err + } + idBytes := make([]byte, 12) + if _, err = rand.Read(idBytes); err != nil { + return err + } + d := &daemon{s: s, store: store, dir: dir, id: hex.EncodeToString(idBytes), binary: binary, runResult: make(chan error, 1), done: make(chan struct{}), backgroundDone: make(chan struct{})} + var parent net.Conn + var decoder *json.Decoder + var hello handoffMessage + if os.Getenv(handoffEnv) != "" { + if os.Getenv(handoffEnv) != strconv.Itoa(state.Protocol) { + return errors.New("unsupported handoff protocol") + } + f := os.NewFile(3, "handoff") + parent, err = net.FileConn(f) + _ = f.Close() + if err != nil { + return err + } + defer parent.Close() + decoder = json.NewDecoder(parent) + if err = decoder.Decode(&hello); err != nil { + return err + } + if hello.Protocol != state.Protocol || hello.StateDir != dir { + return errors.New("incompatible handoff protocol or state_dir") + } + } else { + gens, err := store.Generations() + if err != nil { + return err + } + for _, g := range gens { + live, err := state.IsLocked(filepath.Join(dir, g.ID+".lock")) + if err != nil { + return err + } + if live { + return fmt.Errorf("generation %s is still running; use upgrade", g.ID) + } + } + if err = store.Reset(); err != nil { + return err + } + for _, g := range gens { + _ = os.Remove(filepath.Join(dir, g.ID+".lock")) + _ = os.Remove(g.Control) + } + } + d.lock, err = state.Lock(filepath.Join(dir, d.id+".lock")) + if err != nil { + return err + } + defer d.lock.Close() + control := filepath.Join(dir, d.id+".sock") + d.privateListener, err = net.Listen("unix", control) + if err != nil { + return err + } + defer d.privateListener.Close() + if err = os.Chmod(control, 0600); err != nil { + return err + } + if err = store.Register(state.Generation{ID: d.id, PID: os.Getpid(), Version: Version, Control: control}); err != nil { + return err + } + defer func() { + if err := store.Finish(d.id); err != nil { + log.Printf("[ERROR] retire generation: %v", err) + } + }() + s.Generation = d.id + // Configuration was loaded before attaching the store: a prospective child + // must not publish new limits until the parent authorizes the commit. + s.Shared = store + s.ControlHandler = d.control + defer s.CloseLogs() + if parent != nil { + n := 2 + if hello.Listeners[1] != "" { + n++ + } + files := make([]*os.File, n) + for i := range files { + files[i] = os.NewFile(uintptr(4+i), "listener") + } + err = s.ListenInherited(files, hello.Listeners) + for _, f := range files { + _ = f.Close() + } + } else { + err = s.Listen() + } + if err != nil { + return err + } + defer s.Close() + d.private = &http.Server{Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/internal/reopen" { + if err := s.ReopenLogs(); err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + } + return + } + if r.URL.Path == "/internal/stop" { + d.beginDrain(0) + return + } + if !d.control(w, r) { + http.NotFound(w, r) + } + })} + go func() { + if err := d.private.Serve(d.privateListener); err != nil && !errors.Is(err, http.ErrServerClosed) && !errors.Is(err, net.ErrClosed) { + log.Printf("[ERROR] private control: %v", err) + } + }() + defer func() { + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + _ = d.private.Shutdown(ctx) + }() + if parent != nil { + enc := json.NewEncoder(parent) + if err = enc.Encode(handoffMessage{Stage: "ready", ID: d.id, DrainTimeout: s.DrainTimeout()}); err != nil { + return err + } + var commit handoffMessage + if err = decoder.Decode(&commit); err != nil { + return err + } + if commit.Stage != "commit" { + return errors.New("upgrade was not committed") + } + } + if err = store.Activate(d.id, hello.Parent, s.Limits()); err != nil { + return err + } + d.stopAdmission = s.StartAdmission() + defer d.stopAdmission() + go func() { d.runResult <- s.Run() }() + runFinished := false + defer func() { + if !runFinished { + s.Close() + <-d.runResult + } + }() + <-s.Started + if err = s.PublishSnapshot(); err != nil { + return err + } + if parent == nil { + if err = notifySystemd("READY=1\nMAINPID=" + strconv.Itoa(os.Getpid())); err != nil { + return err + } + } else { + if err = json.NewEncoder(parent).Encode(handoffMessage{Stage: "active", ID: d.id}); err != nil { + return err + } + // Keep the parent alive until it has transferred systemd ownership. + var ack handoffMessage + if err = decoder.Decode(&ack); err != nil { + return err + } + if ack.Stage != "owned" { + return errors.New("missing ownership acknowledgement") + } + } + _ = startup.Close() + d.ready.Store(true) + go d.background() + defer func() { close(d.done); <-d.backgroundDone }() + signals := make(chan os.Signal, 2) + signal.Notify(signals, syscall.SIGTERM, syscall.SIGINT) + defer signal.Stop(signals) + var runErr error + select { + case runErr = <-d.runResult: + case <-signals: + d.stopAll() + select { + case runErr = <-d.runResult: + case <-signals: + s.AbortConnections() + runErr = <-d.runResult + } + } + runFinished = true + if d.drainTimer != nil { + d.drainTimer.Stop() + } + // Finish HTTP handlers before closing their database/log dependencies. + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + if s.HTTPServer != nil { + _ = s.HTTPServer.Shutdown(ctx) + } + if err = s.PublishSnapshot(); err != nil { + log.Printf("[ERROR] final snapshot: %v", err) + } + current, _ := store.Current() + if current == d.id { + for _, addr := range s.ConfiguredListeners() { + if strings.HasPrefix(addr, "/") { + _ = os.Remove(addr) + } + } + } + return runErr +} + +func (d *daemon) beginDrain(timeout time.Duration) { + d.drainOnce.Do(func() { + var deadline time.Time + if timeout > 0 { + deadline = time.Now().Add(timeout) + d.drainTimer = time.AfterFunc(timeout, func() { log.Printf("[INFO] generation %s drain timeout reached", d.id); d.s.AbortConnections() }) + } + if err := d.store.Drain(d.id, deadline); err != nil { + log.Printf("[ERROR] record drain: %v", err) + } + log.Printf("[INFO] draining generation %s (timeout %s)", d.id, timeout) + d.s.StopAccepting() + }) +} + +func (d *daemon) background() { + defer close(d.backgroundDone) + ticker := time.NewTicker(time.Second) + defer ticker.Stop() + for { + select { + case <-d.done: + return + case <-ticker.C: + if err := d.s.PublishSnapshot(); err != nil { + log.Printf("[ERROR] publish snapshot: %v", err) + } + gens, err := d.store.Generations() + if err != nil { + continue + } + for _, g := range gens { + if g.ID == d.id || g.State == state.Exited { + continue + } + live, err := state.IsLocked(filepath.Join(d.dir, g.ID+".lock")) + if err == nil && !live { + if err := d.store.Finish(g.ID); err != nil { + log.Printf("[ERROR] reclaim generation: %v", err) + } + } + } + } + } +} + +func (d *daemon) stopAll() { + current, err := d.store.Current() + if err == nil && current == d.id { + gens, err := d.store.Generations() + if err == nil { + for _, g := range gens { + if g.ID != d.id && g.State != state.Exited { + if err := privatePost(g.Control, "/internal/stop"); err != nil { + log.Printf("[ERROR] stop generation %s: %v", g.ID, err) + } + } + } + } + } + d.beginDrain(0) +} + +func privatePost(addr, path string) error { + client := makeHttpClient(addr) + client.Timeout = 3 * time.Second + resp, err := client.Post("http://."+path, "application/json", nil) + 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("%s: %s", resp.Status, b) + } + return nil +} + +func (d *daemon) control(w http.ResponseWriter, r *http.Request) bool { + switch r.URL.Path { + case "/upgrade", "/reload", "/reopen-logs": + default: + return false + } + if !d.ready.Load() { + http.Error(w, "generation is starting", http.StatusServiceUnavailable) + return true + } + if r.Method != http.MethodPost { + http.Error(w, "POST required", http.StatusMethodNotAllowed) + return true + } + if !d.mu.TryLock() { + http.Error(w, "upgrade or reload in progress", http.StatusConflict) + return true + } + defer d.mu.Unlock() + current, err := d.store.Current() + if err != nil { + http.Error(w, err.Error(), http.StatusServiceUnavailable) + return true + } + if current != d.id { + // A keepalive HTTP connection can still belong to an older generation. + gens, err := d.store.Generations() + if err != nil { + http.Error(w, err.Error(), http.StatusServiceUnavailable) + return true + } + for _, g := range gens { + if g.ID == current { + client := makeHttpClient(g.Control) + client.Timeout = 0 + req, _ := http.NewRequestWithContext(r.Context(), http.MethodPost, "http://."+r.URL.RequestURI(), nil) + resp, err := client.Do(req) + if err != nil { + http.Error(w, err.Error(), http.StatusServiceUnavailable) + return true + } + defer resp.Body.Close() + w.WriteHeader(resp.StatusCode) + _, _ = io.Copy(w, resp.Body) + return true + } + } + http.Error(w, "current generation unavailable", http.StatusServiceUnavailable) + return true + } + switch r.URL.Path { + case "/reload": + if err := d.s.ReadConfigFromFile(true); err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + } else { + _, _ = fmt.Fprintln(w, `{"message":"Successfully reloaded"}`) + } + case "/reopen-logs": + gens, err := d.store.Generations() + if err != nil { + http.Error(w, err.Error(), http.StatusServiceUnavailable) + return true + } + for _, g := range gens { + if g.State != state.Exited { + if err := privatePost(g.Control, "/internal/reopen"); err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return true + } + } + } + _, _ = fmt.Fprintln(w, `{"message":"Reopened logs in all generations"}`) + case "/upgrade": + timeout := 30 * time.Second + if raw := r.URL.Query().Get("timeout"); raw != "" { + timeout, err = time.ParseDuration(raw) + if err != nil || timeout <= 0 { + http.Error(w, "invalid timeout", http.StatusBadRequest) + return true + } + } + if err := d.upgrade(timeout); err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + } else { + _, _ = fmt.Fprintln(w, `{"message":"Upgrade completed"}`) + } + } + return true +} + +func (d *daemon) upgrade(timeout time.Duration) (resultErr error) { + files, err := d.s.ListenerFiles() + if err != nil { + return err + } + defer func() { + for _, f := range files { + _ = f.Close() + } + }() + pair, err := unix.Socketpair(unix.AF_UNIX, unix.SOCK_STREAM|unix.SOCK_CLOEXEC, 0) + if err != nil { + return err + } + pf := os.NewFile(uintptr(pair[0]), "parent") + childFile := os.NewFile(uintptr(pair[1]), "child") + defer childFile.Close() + conn, err := net.FileConn(pf) + _ = pf.Close() + if err != nil { + return err + } + defer conn.Close() + _ = conn.SetDeadline(time.Now().Add(timeout)) + env := make([]string, 0, len(os.Environ())+1) + for _, v := range os.Environ() { + if !strings.HasPrefix(v, handoffEnv+"=") { + env = append(env, v) + } + } + env = append(env, handoffEnv+"="+strconv.Itoa(state.Protocol)) + child := &exec.Cmd{Path: d.binary, Args: append([]string{d.binary}, os.Args[1:]...), Env: env, ExtraFiles: append([]*os.File{childFile}, files...), Stdin: os.Stdin, Stdout: os.Stdout, Stderr: os.Stderr} + if err = child.Start(); err != nil { + return err + } + _ = childFile.Close() + wait := make(chan error, 1) + go func() { wait <- child.Wait() }() + success := false + defer func() { + if success { + return + } + _ = child.Process.Kill() + <-wait + // The old listeners have not been closed. Restore admission policy if + // the child failed after committing but before acknowledging ownership. + current, e := d.store.Current() + if e == nil && current != d.id { + if e = d.store.Activate(d.id, current, d.s.Limits()); e != nil { + log.Printf("[ERROR] restore serving generation: %v", e) + } + } + _ = notifySystemd("MAINPID=" + strconv.Itoa(os.Getpid())) + }() + enc := json.NewEncoder(conn) + dec := json.NewDecoder(conn) + if err = enc.Encode(handoffMessage{Protocol: state.Protocol, Parent: d.id, StateDir: d.dir, Listeners: d.s.ConfiguredListeners()}); err != nil { + return err + } + var ready handoffMessage + if err = dec.Decode(&ready); err != nil { + return fmt.Errorf("new generation did not become ready: %w", err) + } + if ready.Stage != "ready" { + return errors.New("invalid readiness acknowledgement") + } + if err = enc.Encode(handoffMessage{Stage: "commit"}); err != nil { + return err + } + var active handoffMessage + if err = dec.Decode(&active); err != nil { + return fmt.Errorf("new generation failed to activate: %w", err) + } + if active.Stage != "active" || active.ID != ready.ID { + return errors.New("invalid activation acknowledgement") + } + if err = notifySystemd("MAINPID=" + strconv.Itoa(child.Process.Pid)); err != nil { + return fmt.Errorf("transfer systemd ownership: %w", err) + } + if err = enc.Encode(handoffMessage{Stage: "owned"}); err != nil { + return err + } + success = true + d.beginDrain(ready.DrainTimeout) + log.Printf("[INFO] upgraded %s to %s (pid %d)", d.id, ready.ID, child.Process.Pid) + return nil +} + +// A barrier makes ownership changes visible to PID 1 before the old main +// process is allowed to exit. Only the current main process sends MAINPID. +func notifySystemd(message string) error { + path := os.Getenv("NOTIFY_SOCKET") + if path == "" { + return nil + } + fd, err := unix.Socket(unix.AF_UNIX, unix.SOCK_DGRAM|unix.SOCK_CLOEXEC, 0) + if err != nil { + return err + } + defer unix.Close(fd) + addr := &unix.SockaddrUnix{Name: path} + if err = unix.Sendto(fd, []byte(message), 0, addr); err != nil { + return err + } + r, w, err := os.Pipe() + if err != nil { + return err + } + defer r.Close() + defer w.Close() + if err = unix.Sendmsg(fd, []byte("BARRIER=1"), unix.UnixRights(int(w.Fd())), addr, 0); err != nil { + return err + } + _ = w.Close() + if err = r.SetReadDeadline(time.Now().Add(3 * time.Second)); err != nil { + return err + } + var b [1]byte + _, err = r.Read(b[:]) + if errors.Is(err, io.EOF) { + return nil + } + return err +} diff --git a/cmd/daemon_linux_test.go b/cmd/daemon_linux_test.go new file mode 100644 index 0000000..c03cf89 --- /dev/null +++ b/cmd/daemon_linux_test.go @@ -0,0 +1,580 @@ +//go:build linux + +package cmd + +import ( + "bufio" + "bytes" + "crypto/ecdsa" + "crypto/elliptic" + "crypto/rand" + "crypto/tls" + "crypto/x509" + "crypto/x509/pkix" + "encoding/json" + "encoding/pem" + "fmt" + "io" + "math/big" + "net" + "net/http" + "os" + "os/exec" + "path/filepath" + "strings" + "sync" + "syscall" + "testing" + "time" + + "github.com/stretchr/testify/require" + + "github.com/ustclug/rsync-proxy/pkg/state" +) + +// Re-exec the test executable so -race also instruments every generation. A +// private copy is atomically replaced during tests, just like an installed bin. +func TestDaemonProcess(t *testing.T) { + config := os.Getenv("RSYNC_DAEMON_TEST_CONFIG") + if config == "" { + return + } + exe, err := os.Executable() + if err != nil { + panic(err) + } + version, err := os.ReadFile(exe + ".version") + if err != nil { + panic(err) + } + Version = string(version) + c := New() + c.SetArgs([]string{"--config", config}) + if err := c.Execute(); err != nil { + fmt.Fprintln(os.Stderr, err) + os.Exit(1) + } + os.Exit(0) +} + +type daemonStatus struct { + Count int `json:"count"` + Connections []struct { + Generation string `json:"generation"` + Index uint32 `json:"index"` + } `json:"connections"` + Generations []state.Generation `json:"generations"` +} + +type daemonFixture struct { + t *testing.T + dir, binary, config, socket, listen, tlsListen, upstream, cert, key string + log *os.File + pid int + max, perIP int + clients []net.Conn + listener net.Listener +} + +func freeAddress(t *testing.T) string { + t.Helper() + l, err := net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + addr := l.Addr().String() + require.NoError(t, l.Close()) + return addr +} + +func newDaemonFixture(t *testing.T, max, perIP int) *daemonFixture { + t.Helper() + dir, err := os.MkdirTemp("", "rp-upgrade-") + require.NoError(t, err) + f := &daemonFixture{t: t, dir: dir, max: max, perIP: perIP, binary: filepath.Join(dir, "proxy"), config: filepath.Join(dir, "config.toml"), socket: filepath.Join(dir, "http.sock"), listen: freeAddress(t), tlsListen: freeAddress(t)} + f.cert, f.key = testCertificate(t, dir) + f.listener, err = net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + f.upstream = f.listener.Addr().String() + go func() { + for { + c, err := f.listener.Accept() + if err != nil { + return + } + go echoRsync(c) + } + }() + f.install("A") + f.writeConfig(0) + f.log, err = os.Create(filepath.Join(dir, "process.log")) + require.NoError(t, err) + c := exec.Command(f.binary, "-test.run=^TestDaemonProcess$") + c.Env = append(os.Environ(), "RSYNC_DAEMON_TEST_CONFIG="+f.config, "NOTIFY_SOCKET=") + c.Stdout = f.log + c.Stderr = f.log + require.NoError(t, c.Start()) + f.pid = c.Process.Pid + wait := make(chan error, 1) + go func() { wait <- c.Wait() }() + t.Cleanup(func() { + for _, conn := range f.clients { + _ = conn.Close() + } + _ = f.listener.Close() + store, err := state.Open(dir) + if err == nil { + gens, _ := store.Generations() + for _, g := range gens { + if g.State != "exited" { + _ = syscall.Kill(g.PID, syscall.SIGKILL) + } + } + _ = store.Close() + } + _ = c.Process.Kill() + select { + case <-wait: + case <-time.After(3 * time.Second): + } + _ = f.log.Close() + b, _ := os.ReadFile(filepath.Join(dir, "process.log")) + if bytes.Contains(b, []byte("DATA RACE")) { + t.Errorf("race in daemon subprocess:\n%s", b) + } + if t.Failed() { + t.Logf("daemon log:\n%s", b) + } + _ = os.RemoveAll(dir) + }) + require.Eventually(t, func() bool { st, err := f.status(); return err == nil && len(st.Generations) == 1 }, 10*time.Second, 20*time.Millisecond) + return f +} + +func echoRsync(c net.Conn) { + defer c.Close() + r := bufio.NewReader(c) + if _, err := r.ReadString('\n'); err != nil { + return + } + if _, err := io.WriteString(c, "@RSYNCD: 32.0 sha512 sha256 sha1 md5 md4\n"); err != nil { + return + } + if _, err := r.ReadString('\n'); err != nil { + return + } + _, _ = io.WriteString(c, "READY\n") + _, _ = io.Copy(c, r) +} + +func testCertificate(t *testing.T, dir string) (string, string) { + t.Helper() + key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader) + require.NoError(t, err) + tmpl := &x509.Certificate{SerialNumber: big.NewInt(1), Subject: pkix.Name{CommonName: "test"}, NotBefore: time.Now().Add(-time.Hour), NotAfter: time.Now().Add(time.Hour), KeyUsage: x509.KeyUsageDigitalSignature, ExtKeyUsage: []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth}} + der, err := x509.CreateCertificate(rand.Reader, tmpl, tmpl, &key.PublicKey, key) + require.NoError(t, err) + keyDER, err := x509.MarshalPKCS8PrivateKey(key) + require.NoError(t, err) + certPath, keyPath := filepath.Join(dir, "cert.pem"), filepath.Join(dir, "key.pem") + require.NoError(t, os.WriteFile(certPath, pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: der}), 0600)) + require.NoError(t, os.WriteFile(keyPath, pem.EncodeToMemory(&pem.Block{Type: "PRIVATE KEY", Bytes: keyDER}), 0600)) + return certPath, keyPath +} + +func (f *daemonFixture) install(version string) { + f.t.Helper() + exe, err := os.Executable() + require.NoError(f.t, err) + in, err := os.Open(exe) + require.NoError(f.t, err) + defer in.Close() + out, err := os.OpenFile(f.binary+".new", os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0700) + require.NoError(f.t, err) + _, err = io.Copy(out, in) + require.NoError(f.t, err) + require.NoError(f.t, out.Close()) + require.NoError(f.t, os.Rename(f.binary+".new", f.binary)) + require.NoError(f.t, os.WriteFile(f.binary+".version", []byte(version), 0600)) +} + +func (f *daemonFixture) writeConfig(timeout int) { + f.t.Helper() + config := fmt.Sprintf(`[proxy] +listen = %q +listen_tls = %q +listen_http = %q +state_dir = %q +tls_cert_file = %q +tls_key_file = %q +access_log = %q +upgrade_drain_timeout = %d +[upstreams.u] +address = %q +modules = ["foo"] +max_active_connections = %d +max_queued_connections = 100 +per_ip_max_active_connections = %d +`, f.listen, f.tlsListen, f.socket, f.dir, f.cert, f.key, filepath.Join(f.dir, "access.log"), timeout, f.upstream, f.max, f.perIP) + require.NoError(f.t, os.WriteFile(f.config, []byte(config), 0600)) +} + +func (f *daemonFixture) status() (daemonStatus, error) { + var st daemonStatus + resp, err := makeHttpClient(f.socket).Get("http://./status") + if err != nil { + return st, err + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + return st, fmt.Errorf("status: %s", resp.Status) + } + err = json.NewDecoder(resp.Body).Decode(&st) + return st, err +} + +func (f *daemonFixture) upgrade() { + f.t.Helper() + var out bytes.Buffer + require.NoError(f.t, sendControl(f.socket, "/upgrade?timeout=10s", 10*time.Second, &out), out.String()) +} + +type relayReader struct { + *bufio.Reader + ready bool +} + +func (f *daemonFixture) connect(useTLS bool) (net.Conn, *relayReader) { + f.t.Helper() + var c net.Conn + var err error + if useTLS { + c, err = tls.Dial("tcp", f.tlsListen, &tls.Config{InsecureSkipVerify: true}) + } else { + c, err = net.Dial("tcp", f.listen) + } + require.NoError(f.t, err) + f.clients = append(f.clients, c) + require.NoError(f.t, c.SetDeadline(time.Now().Add(5*time.Second))) + r := bufio.NewReader(c) + _, err = io.WriteString(c, "@RSYNCD: 32.0\n") + require.NoError(f.t, err) + line, err := r.ReadString('\n') + require.NoError(f.t, err) + require.Contains(f.t, line, "@RSYNCD:") + _, err = io.WriteString(c, "foo\n") + require.NoError(f.t, err) + return c, &relayReader{Reader: r} +} + +func checkEcho(t *testing.T, c net.Conn, r *relayReader) { + t.Helper() + require.NoError(t, c.SetDeadline(time.Now().Add(5*time.Second))) + if !r.ready { + line, err := r.ReadString('\n') + require.NoError(t, err) + require.Equal(t, "READY\n", line) + r.ready = true + } + _, err := io.WriteString(c, "ping\n") + require.NoError(t, err) + line, err := r.ReadString('\n') + require.NoError(t, err) + require.Equal(t, "ping\n", line) +} + +func TestMultiGenerationUpgrade(t *testing.T) { + f := newDaemonFixture(t, 3, 10) + a, ar := f.connect(false) + checkEcho(t, a, ar) + // Accept a TLS connection before upgrade, but defer its handshake until B. + raw, err := net.Dial("tcp", f.tlsListen) + require.NoError(t, err) + f.clients = append(f.clients, raw) + require.Eventually(t, func() bool { st, err := f.status(); return err == nil && st.Count == 2 }, 3*time.Second, 10*time.Millisecond) + f.install("B") + f.upgrade() + b, br := f.connect(true) + checkEcho(t, b, br) + f.install("C") + f.upgrade() + checkEcho(t, a, ar) + checkEcho(t, b, br) + lazy := tls.Client(raw, &tls.Config{InsecureSkipVerify: true}) + require.NoError(t, lazy.SetDeadline(time.Now().Add(5*time.Second))) + require.NoError(t, lazy.Handshake()) + lr := &relayReader{Reader: bufio.NewReader(lazy)} + _, err = io.WriteString(lazy, "@RSYNCD: 32.0\n") + require.NoError(t, err) + _, err = lr.ReadString('\n') + require.NoError(t, err) + _, err = io.WriteString(lazy, "foo\n") + require.NoError(t, err) + checkEcho(t, lazy, lr) + c, cr := f.connect(false) + line, err := cr.ReadString('\n') + require.NoError(t, err) + require.Contains(t, line, "queued") + _, err = cr.ReadString('\n') + require.NoError(t, err) + require.Eventually(t, func() bool { + st, err := f.status() + if err != nil { + return false + } + n := 0 + for _, g := range st.Generations { + if g.State != "exited" { + n++ + } + } + return n == 3 && st.Count == 4 + }, 4*time.Second, 20*time.Millisecond) + // Releasing an A slot must promote C without exceeding the global cap. + require.NoError(t, a.Close()) + checkEcho(t, c, cr) + resp, err := makeHttpClient(f.socket).Get("http://./metrics") + require.NoError(t, err) + data, err := io.ReadAll(resp.Body) + _ = resp.Body.Close() + require.NoError(t, err) + require.Contains(t, string(data), "generation=\"") + require.Contains(t, string(data), `version="A"`) + require.Contains(t, string(data), `version="C"`) + require.NoError(t, os.Rename(filepath.Join(f.dir, "access.log"), filepath.Join(f.dir, "access.log.old"))) + var out bytes.Buffer + require.NoError(t, sendControl(f.socket, "/reopen-logs", 5*time.Second, &out)) + require.NoError(t, lazy.Close()) + require.NoError(t, b.Close()) + require.NoError(t, c.Close()) + require.Eventually(t, func() bool { + st, err := f.status() + if err != nil { + return false + } + n := 0 + for _, g := range st.Generations { + if g.State != "exited" { + n++ + } + } + return n == 1 && st.Count == 0 + }, 5*time.Second, 20*time.Millisecond) + _, err = os.Stat(f.socket) + require.NoError(t, err) + logs, err := os.ReadFile(filepath.Join(f.dir, "access.log")) + require.NoError(t, err) + require.Contains(t, string(logs), "finishes module foo") +} + +func TestUpgradeFailureAndDrainTimeout(t *testing.T) { + f := newDaemonFixture(t, 1, 2) + a, ar := f.connect(false) + checkEcho(t, a, ar) + queued, qr := f.connect(false) + line, err := qr.ReadString('\n') + require.NoError(t, err) + require.Contains(t, line, "queued") + f.install("B") + original, err := os.ReadFile(f.config) + require.NoError(t, err) + require.NoError(t, os.WriteFile(f.config, []byte("invalid toml ["), 0600)) + var out bytes.Buffer + require.Error(t, sendControl(f.socket, "/upgrade?timeout=2s", 2*time.Second, &out)) + checkEcho(t, a, ar) + require.NoError(t, os.WriteFile(f.config, original, 0600)) + // A silent executable never acknowledges readiness; rollback must keep A. + require.NoError(t, os.WriteFile(f.binary+".new", []byte("#!/bin/sh\nexec sleep 10\n"), 0700)) + require.NoError(t, os.Rename(f.binary+".new", f.binary)) + require.Error(t, sendControl(f.socket, "/upgrade?timeout=100ms", time.Second, &out)) + checkEcho(t, a, ar) + f.install("B") + f.writeConfig(2) + f.upgrade() + // Per-IP cap includes A's active and queued sessions after handoff. + c, cr := f.connect(false) + line, err = cr.ReadString('\n') + require.NoError(t, err) + require.Contains(t, line, "per-IP connection limit") + _ = c.Close() + f.install("C") + f.writeConfig(0) + f.upgrade() + // C's unlimited drain setting must not reset A's previously assigned timer. + require.NoError(t, a.SetReadDeadline(time.Now().Add(5*time.Second))) + _, err = ar.ReadString('\n') + require.Error(t, err) + require.NoError(t, queued.SetReadDeadline(time.Now().Add(5*time.Second))) + _, _ = io.Copy(io.Discard, qr) + require.Eventually(t, func() bool { + st, err := f.status() + if err != nil { + return false + } + for _, g := range st.Generations { + if g.Version == "A" { + return g.State == "exited" + } + } + return false + }, 5*time.Second, 20*time.Millisecond) + n, nr := f.connect(false) + checkEcho(t, n, nr) +} + +func TestDeadGenerationReleasesCapacity(t *testing.T) { + f := newDaemonFixture(t, 1, 0) + a, ar := f.connect(false) + checkEcho(t, a, ar) + f.install("B") + f.upgrade() + b, br := f.connect(false) + line, err := br.ReadString('\n') + require.NoError(t, err) + require.Contains(t, line, "queued") + _, err = br.ReadString('\n') + require.NoError(t, err) + require.NoError(t, syscall.Kill(f.pid, syscall.SIGKILL)) + checkEcho(t, b, br) +} + +func TestConnectionsDuringUpgrade(t *testing.T) { + f := newDaemonFixture(t, 0, 0) + var wg sync.WaitGroup + errorsC := make(chan error, 200) + wg.Go(func() { + for range 200 { + c, err := net.DialTimeout("tcp", f.listen, time.Second) + if err != nil { + errorsC <- err + continue + } + _ = c.SetDeadline(time.Now().Add(time.Second)) + r := bufio.NewReader(c) + _, err = io.WriteString(c, "@RSYNCD: 32.0\n") + if err == nil { + var line string + line, err = r.ReadString('\n') + if err == nil && !strings.HasPrefix(line, "@RSYNCD:") { + err = fmt.Errorf("bad greeting %q", line) + } + } + if err == nil { + _, err = io.WriteString(c, "foo\n") + } + if err == nil { + var line string + line, err = r.ReadString('\n') + if err == nil && line != "READY\n" { + err = fmt.Errorf("bad upstream response %q", line) + } + } + _ = c.Close() + if err != nil { + errorsC <- err + } + time.Sleep(10 * time.Millisecond) + } + }) + f.install("B") + f.upgrade() + f.install("C") + f.upgrade() + wg.Wait() + close(errorsC) + for err := range errorsC { + require.NoError(t, err) + } +} + +func TestUpgradeRejectsIncompatibleConfig(t *testing.T) { + f := newDaemonFixture(t, 0, 0) + c, r := f.connect(false) + checkEcho(t, c, r) + original, err := os.ReadFile(f.config) + require.NoError(t, err) + for _, tc := range []struct{ name, old, replacement string }{ + {"listener", f.listen, freeAddress(t)}, + {"certificate", f.cert, filepath.Join(f.dir, "missing.pem")}, + {"state directory", "state_dir = " + fmt.Sprintf("%q", f.dir), "state_dir = " + fmt.Sprintf("%q", filepath.Join(f.dir, "other"))}, + } { + t.Run(tc.name, func(t *testing.T) { + require.NoError(t, os.WriteFile(f.config, bytes.Replace(original, []byte(tc.old), []byte(tc.replacement), 1), 0600)) + var out bytes.Buffer + require.Error(t, sendControl(f.socket, "/upgrade?timeout=2s", 2*time.Second, &out)) + checkEcho(t, c, r) + require.NoError(t, os.WriteFile(f.config, original, 0600)) + }) + } + f.install("B") + f.upgrade() + checkEcho(t, c, r) +} + +func TestRealRsyncTransferAcrossUpgrades(t *testing.T) { + if _, err := exec.LookPath("rsync"); err != nil { + t.Skip("rsync is not installed") + } + f := newDaemonFixture(t, 2, 0) + source := filepath.Join(f.dir, "source") + destination := filepath.Join(f.dir, "destination") + require.NoError(t, os.Mkdir(source, 0755)) + require.NoError(t, os.Mkdir(destination, 0755)) + data := make([]byte, 1024*1024) + _, err := rand.Read(data) + require.NoError(t, err) + require.NoError(t, os.WriteFile(filepath.Join(source, "payload"), data, 0644)) + address := freeAddress(t) + _, port, err := net.SplitHostPort(address) + require.NoError(t, err) + rsyncConfig := filepath.Join(f.dir, "rsyncd.conf") + require.NoError(t, os.WriteFile(rsyncConfig, []byte(fmt.Sprintf("use chroot = no\n[foo]\npath = %s\nread only = yes\n", source)), 0600)) + daemon := exec.Command("rsync", "--daemon", "--no-detach", "--address=127.0.0.1", "--port="+port, "--config="+rsyncConfig) + daemon.Stdout = f.log + daemon.Stderr = f.log + require.NoError(t, daemon.Start()) + t.Cleanup(func() { _ = daemon.Process.Kill(); _ = daemon.Wait() }) + require.Eventually(t, func() bool { + conn, err := net.DialTimeout("tcp", address, time.Second) + if err != nil { + return false + } + _ = conn.Close() + return true + }, 3*time.Second, 10*time.Millisecond) + f.upstream = address + f.writeConfig(0) + var reload bytes.Buffer + require.NoError(t, sendControl(f.socket, "/reload", 5*time.Second, &reload)) + client := exec.Command("rsync", "--bwlimit=256", "rsync://"+f.listen+"/foo/payload", destination+"/") + var output bytes.Buffer + client.Stdout = &output + client.Stderr = &output + require.NoError(t, client.Start()) + finished := make(chan error, 1) + go func() { finished <- client.Wait() }() + t.Cleanup(func() { _ = client.Process.Kill() }) + require.Eventually(t, func() bool { + select { + case err := <-finished: + t.Fatalf("rsync exited before upgrade: %v: %s", err, output.String()) + default: + } + st, err := f.status() + return err == nil && st.Count > 0 + }, 3*time.Second, 10*time.Millisecond) + f.install("B") + f.upgrade() + f.install("C") + f.upgrade() + listing, err := exec.Command("rsync", "--list-only", "rsync://"+f.listen+"/foo/").CombinedOutput() + require.NoError(t, err, string(listing)) + require.Contains(t, string(listing), "payload") + select { + case err := <-finished: + require.NoError(t, err, output.String()) + case <-time.After(15 * time.Second): + t.Fatal("rsync transfer did not finish") + } + actual, err := os.ReadFile(filepath.Join(destination, "payload")) + require.NoError(t, err) + require.Equal(t, data, actual) +} diff --git a/cmd/daemon_other.go b/cmd/daemon_other.go new file mode 100644 index 0000000..c449dd2 --- /dev/null +++ b/cmd/daemon_other.go @@ -0,0 +1,12 @@ +//go:build !linux + +package cmd + +import "github.com/ustclug/rsync-proxy/pkg/server" + +func runDaemon(s *server.Server) error { + if err := s.Listen(); err != nil { + return err + } + return s.Run() +} diff --git a/cmd/upgrade.go b/cmd/upgrade.go new file mode 100644 index 0000000..3f9b977 --- /dev/null +++ b/cmd/upgrade.go @@ -0,0 +1,82 @@ +package cmd + +import ( + "fmt" + "io" + "net/http" + "net/url" + "os" + "time" + + "github.com/pelletier/go-toml" + "github.com/spf13/cobra" + + "github.com/ustclug/rsync-proxy/pkg/server" +) + +func sendControl(addr, path string, timeout time.Duration, out io.Writer) error { + client := makeHttpClient(addr) + client.Timeout = timeout + 5*time.Second + resp, err := client.Post("http://."+path, "application/json", nil) + if err != nil { + return err + } + defer resp.Body.Close() + _, _ = io.Copy(out, resp.Body) + if resp.StatusCode != http.StatusOK { + return fmt.Errorf("control request failed: %s", resp.Status) + } + return nil +} + +func newUpgradeCmd(s *server.Server) *cobra.Command { + var timeout time.Duration + c := &cobra.Command{Use: "upgrade", Short: "Start the installed binary and drain the previous generation", Args: cobra.NoArgs, + RunE: func(c *cobra.Command, _ []string) error { + if timeout <= 0 { + return fmt.Errorf("timeout must be positive") + } + addr, err := controlAddress(c, s) + if err != nil { + return err + } + return sendControl(addr, "/upgrade?timeout="+url.QueryEscape(timeout.String()), timeout, c.OutOrStdout()) + }} + c.Flags().DurationVar(&timeout, "timeout", 30*time.Second, "New generation readiness timeout") + return c +} + +func newReopenLogsCmd(s *server.Server) *cobra.Command { + return &cobra.Command{Use: "reopen-logs", Short: "Reopen log files in all live generations", Args: cobra.NoArgs, + RunE: func(c *cobra.Command, _ []string) error { + addr, err := controlAddress(c, s) + if err != nil { + return err + } + return sendControl(addr, "/reopen-logs", 30*time.Second, c.OutOrStdout()) + }} +} + +// Lifecycle commands honor --config without opening logs, discovering modules, +// or otherwise initializing a second server. Explicit --host takes precedence. +func controlAddress(c *cobra.Command, s *server.Server) (string, error) { + if c.Flags().Changed("host") || c.InheritedFlags().Changed("host") { + return daemonSocket, nil + } + f, err := os.Open(s.ConfigPath) + if os.IsNotExist(err) { + return daemonSocket, nil + } + if err != nil { + return "", err + } + defer f.Close() + var config server.Config + if err = toml.NewDecoder(f).Decode(&config); err != nil { + return "", err + } + if config.Proxy.ListenHTTP != "" { + return config.Proxy.ListenHTTP, nil + } + return daemonSocket, nil +} diff --git a/contrib/grafana/dashboard.json b/contrib/grafana/dashboard.json index 76a40e4..c4c1c7e 100644 --- a/contrib/grafana/dashboard.json +++ b/contrib/grafana/dashboard.json @@ -1248,7 +1248,7 @@ { "id": 16, "title": "Live Connections (per-connection snapshot)", - "description": "One row per currently-active connection. Bytes are cumulative for the session and reset when the connection ends or rsync-proxy restarts.", + "description": "One row per currently-active connection, keyed by generation and index across hot upgrades. Bytes are cumulative for the session.", "type": "table", "gridPos": { "h": 9, @@ -1331,7 +1331,7 @@ "type": "prometheus", "uid": "${DS_PROMETHEUS}" }, - "expr": "rsync_proxy_connection_duration_seconds{job=\"$job\",instance=\"$instance\"}", + "expr": "label_join(rsync_proxy_connection_duration_seconds{job=\"$job\",instance=\"$instance\"}, \"connection\", \":\", \"generation\", \"index\")", "format": "table", "instant": true, "refId": "A" @@ -1341,7 +1341,7 @@ "type": "prometheus", "uid": "${DS_PROMETHEUS}" }, - "expr": "rsync_proxy_connection_sent_bytes{job=\"$job\",instance=\"$instance\"}", + "expr": "label_join(rsync_proxy_connection_sent_bytes{job=\"$job\",instance=\"$instance\"}, \"connection\", \":\", \"generation\", \"index\")", "format": "table", "instant": true, "refId": "B" @@ -1351,7 +1351,7 @@ "type": "prometheus", "uid": "${DS_PROMETHEUS}" }, - "expr": "rsync_proxy_connection_received_bytes{job=\"$job\",instance=\"$instance\"}", + "expr": "label_join(rsync_proxy_connection_received_bytes{job=\"$job\",instance=\"$instance\"}, \"connection\", \":\", \"generation\", \"index\")", "format": "table", "instant": true, "refId": "C" @@ -1361,7 +1361,7 @@ { "id": "joinByField", "options": { - "byField": "index", + "byField": "connection", "mode": "outer" } }, @@ -1369,6 +1369,12 @@ "id": "organize", "options": { "excludeByName": { + "index 1": true, + "index 2": true, + "index 3": true, + "generation 1": true, + "generation 2": true, + "generation 3": true, "Time 1": true, "Time 2": true, "Time 3": true, @@ -1387,7 +1393,7 @@ "upstream 3": true }, "indexByName": { - "index": 0, + "connection": 0, "module 1": 1, "upstream 1": 2, "Value #A": 3, diff --git a/go.mod b/go.mod index 6a0f54e..31bf013 100644 --- a/go.mod +++ b/go.mod @@ -7,8 +7,12 @@ require ( github.com/olekukonko/tablewriter v1.1.4 github.com/pelletier/go-toml v1.9.5 github.com/prometheus/client_golang v1.23.2 + github.com/prometheus/client_model v0.6.2 + github.com/prometheus/common v0.66.1 github.com/spf13/cobra v1.10.2 github.com/stretchr/testify v1.11.1 + golang.org/x/sys v0.37.0 + modernc.org/sqlite v1.46.1 ) require ( @@ -17,22 +21,27 @@ require ( github.com/clipperhouse/displaywidth v0.10.0 // indirect github.com/clipperhouse/uax29/v2 v2.6.0 // indirect github.com/davecgh/go-spew v1.1.1 // indirect + github.com/dustin/go-humanize v1.0.1 // indirect + github.com/google/uuid v1.6.0 // indirect github.com/inconshreveable/mousetrap v1.1.0 // indirect github.com/kr/text v0.2.0 // indirect github.com/mattn/go-colorable v0.1.14 // indirect github.com/mattn/go-isatty v0.0.20 // indirect github.com/mattn/go-runewidth v0.0.19 // indirect github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect + github.com/ncruces/go-strftime v1.0.0 // indirect github.com/olekukonko/cat v0.0.0-20250911104152-50322a0618f6 // indirect github.com/olekukonko/errors v1.2.0 // indirect github.com/olekukonko/ll v0.1.6 // indirect github.com/pmezard/go-difflib v1.0.0 // indirect - github.com/prometheus/client_model v0.6.2 // indirect - github.com/prometheus/common v0.66.1 // indirect github.com/prometheus/procfs v0.16.1 // indirect + github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect github.com/spf13/pflag v1.0.10 // indirect go.yaml.in/yaml/v2 v2.4.2 // indirect - golang.org/x/sys v0.35.0 // indirect + golang.org/x/exp v0.0.0-20251023183803-a4bb9ffd2546 // indirect google.golang.org/protobuf v1.36.8 // indirect gopkg.in/yaml.v3 v3.0.1 // indirect + modernc.org/libc v1.67.6 // indirect + modernc.org/mathutil v1.7.1 // indirect + modernc.org/memory v1.11.0 // indirect ) diff --git a/go.sum b/go.sum index d09b71d..c2ca8c7 100644 --- a/go.sum +++ b/go.sum @@ -10,10 +10,18 @@ github.com/cpuguy83/go-md2man/v2 v2.0.6/go.mod h1:oOW0eioCTA6cOiMLiUPZOpcVxMig6N github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E= 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/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY= +github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto= github.com/fatih/color v1.18.0 h1:S8gINlzdQ840/4pfAwic/ZE0djQEH3wM94VfqLTZcOM= github.com/fatih/color v1.18.0/go.mod h1:4FelSpRwEGDpQ12mAdzqdOukCy4u8WUtOY6lkT/6HfU= 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/pprof v0.0.0-20250317173921-a4b03ec1a45e h1:ijClszYn+mADRFY17kjQEVQ1XRhq2/JR1M3sGqeJoxs= +github.com/google/pprof v0.0.0-20250317173921-a4b03ec1a45e/go.mod h1:boTsfXsheKC2y+lKOCMpSfarhxDeIzfZG1jqGcPl3cA= +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/hashicorp/golang-lru/v2 v2.0.7 h1:a+bsQ5rvGLjzHuww6tVxozPZFVghXaHOwFs4luLUK2k= +github.com/hashicorp/golang-lru/v2 v2.0.7/go.mod h1:QeFd9opnmA6QUJc5vARoKUSoFhyfM2/ZepoAG6RGpeM= github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8= github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw= github.com/klauspost/compress v1.18.0 h1:c/Cqfb0r+Yi+JtIEq73FWXVkRonBlf0CRNYc8Zttxdo= @@ -32,6 +40,8 @@ github.com/mattn/go-runewidth v0.0.19 h1:v++JhqYnZuu5jSKrk9RbgF5v4CGUjqRfBm05byF github.com/mattn/go-runewidth v0.0.19/go.mod h1:XBkDxAl56ILZc9knddidhrOlY5R/pDhgLpndooCuJAs= 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/ncruces/go-strftime v1.0.0 h1:HMFp8mLCTPp341M/ZnA4qaf7ZlsbTc+miZjCLOFAw7w= +github.com/ncruces/go-strftime v1.0.0/go.mod h1:Fwc5htZGVVkseilnfgOVb9mKy6w1naJmn9CehxcKcls= github.com/olekukonko/cat v0.0.0-20250911104152-50322a0618f6 h1:zrbMGy9YXpIeTnGj4EljqMiZsIcE09mmF8XsD5AYOJc= github.com/olekukonko/cat v0.0.0-20250911104152-50322a0618f6/go.mod h1:rEKTHC9roVVicUIfZK7DYrdIoM0EOr8mK1Hj5s3JjH0= github.com/olekukonko/errors v1.2.0 h1:10Zcn4GeV59t/EGqJc8fUjtFT/FuUh5bTMzZ1XwmCRo= @@ -52,6 +62,8 @@ github.com/prometheus/common v0.66.1 h1:h5E0h5/Y8niHc5DlaLlWLArTQI7tMrsfQjHV+d9Z github.com/prometheus/common v0.66.1/go.mod h1:gcaUsgf3KfRSwHY4dIMXLPV0K/Wg1oZ8+SbZk/HH/dA= github.com/prometheus/procfs v0.16.1 h1:hZ15bTNuirocR6u0JZ6BAHHmwS1p8B4P6MRqxtzMyRg= github.com/prometheus/procfs v0.16.1/go.mod h1:teAbpZRB1iIAJYREa1LsoWUXykVXA1KlTmWl8x/U+Is= +github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec h1:W09IVJc94icq4NjY3clb7Lk8O1qJ8BdBEF8z0ibU0rE= +github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec/go.mod h1:qqbHyh8v60DhA7CoWK5oRCqLrMHRGoxYCSS9EjAz6Eo= 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/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM= @@ -67,9 +79,17 @@ go.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE= go.yaml.in/yaml/v2 v2.4.2 h1:DzmwEr2rDGHl7lsFgAHxmNz/1NlQ7xLIrlN2h5d1eGI= go.yaml.in/yaml/v2 v2.4.2/go.mod h1:081UH+NErpNdqlCXm3TtEran0rJZGxAYx9hb/ELlsPU= go.yaml.in/yaml/v3 v3.0.4/go.mod h1:DhzuOOF2ATzADvBadXxruRBLzYTpT36CKvDb3+aBEFg= +golang.org/x/exp v0.0.0-20251023183803-a4bb9ffd2546 h1:mgKeJMpvi0yx/sU5GsxQ7p6s2wtOnGAHZWCHUM4KGzY= +golang.org/x/exp v0.0.0-20251023183803-a4bb9ffd2546/go.mod h1:j/pmGrbnkbPtQfxEe5D0VQhZC6qKbfKifgD0oM7sR70= +golang.org/x/mod v0.29.0 h1:HV8lRxZC4l2cr3Zq1LvtOsi/ThTgWnUk/y64QSs8GwA= +golang.org/x/mod v0.29.0/go.mod h1:NyhrlYXJ2H4eJiRy/WDBO6HMqZQ6q9nk4JzS3NuCK+w= +golang.org/x/sync v0.17.0 h1:l60nONMj9l5drqw6jlhIELNv9I0A4OFgRsG9k2oT9Ug= +golang.org/x/sync v0.17.0/go.mod h1:9KTHXmSnoGruLpwFjVSX0lNNA75CykiMECbovNTZqGI= golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= -golang.org/x/sys v0.35.0 h1:vz1N37gP5bs89s7He8XuIYXpyY0+QlsKmzipCbUtyxI= -golang.org/x/sys v0.35.0/go.mod h1:BJP2sWEmIv4KK5OTEluFJCKSidICx8ciO85XgH3Ak8k= +golang.org/x/sys v0.37.0 h1:fdNQudmxPjkdUTPnLn5mdQv7Zwvbvpaxqs831goi9kQ= +golang.org/x/sys v0.37.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks= +golang.org/x/tools v0.38.0 h1:Hx2Xv8hISq8Lm16jvBZ2VQf+RLmbd7wVUsALibYI/IQ= +golang.org/x/tools v0.38.0/go.mod h1:yEsQ/d/YK8cjh0L6rZlY8tgtlKiBNTL14pGDJPJpYQs= google.golang.org/protobuf v1.36.8 h1:xHScyCOEuuwZEc6UtSOvPbAT4zRh0xcNRYekJwfqyMc= google.golang.org/protobuf v1.36.8/go.mod h1:fuxRtAxBytpl4zzqUh6/eyUujkJdNiuEkXntxiD/uRU= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= @@ -77,3 +97,31 @@ gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntN gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q= gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= +modernc.org/cc/v4 v4.27.1 h1:9W30zRlYrefrDV2JE2O8VDtJ1yPGownxciz5rrbQZis= +modernc.org/cc/v4 v4.27.1/go.mod h1:uVtb5OGqUKpoLWhqwNQo/8LwvoiEBLvZXIQ/SmO6mL0= +modernc.org/ccgo/v4 v4.30.1 h1:4r4U1J6Fhj98NKfSjnPUN7Ze2c6MnAdL0hWw6+LrJpc= +modernc.org/ccgo/v4 v4.30.1/go.mod h1:bIOeI1JL54Utlxn+LwrFyjCx2n2RDiYEaJVSrgdrRfM= +modernc.org/fileutil v1.3.40 h1:ZGMswMNc9JOCrcrakF1HrvmergNLAmxOPjizirpfqBA= +modernc.org/fileutil v1.3.40/go.mod h1:HxmghZSZVAz/LXcMNwZPA/DRrQZEVP9VX0V4LQGQFOc= +modernc.org/gc/v2 v2.6.5 h1:nyqdV8q46KvTpZlsw66kWqwXRHdjIlJOhG6kxiV/9xI= +modernc.org/gc/v2 v2.6.5/go.mod h1:YgIahr1ypgfe7chRuJi2gD7DBQiKSLMPgBQe9oIiito= +modernc.org/gc/v3 v3.1.1 h1:k8T3gkXWY9sEiytKhcgyiZ2L0DTyCQ/nvX+LoCljoRE= +modernc.org/gc/v3 v3.1.1/go.mod h1:HFK/6AGESC7Ex+EZJhJ2Gni6cTaYpSMmU/cT9RmlfYY= +modernc.org/goabi0 v0.2.0 h1:HvEowk7LxcPd0eq6mVOAEMai46V+i7Jrj13t4AzuNks= +modernc.org/goabi0 v0.2.0/go.mod h1:CEFRnnJhKvWT1c1JTI3Avm+tgOWbkOu5oPA8eH8LnMI= +modernc.org/libc v1.67.6 h1:eVOQvpModVLKOdT+LvBPjdQqfrZq+pC39BygcT+E7OI= +modernc.org/libc v1.67.6/go.mod h1:JAhxUVlolfYDErnwiqaLvUqc8nfb2r6S6slAgZOnaiE= +modernc.org/mathutil v1.7.1 h1:GCZVGXdaN8gTqB1Mf/usp1Y/hSqgI2vAGGP4jZMCxOU= +modernc.org/mathutil v1.7.1/go.mod h1:4p5IwJITfppl0G4sUEDtCr4DthTaT47/N3aT6MhfgJg= +modernc.org/memory v1.11.0 h1:o4QC8aMQzmcwCK3t3Ux/ZHmwFPzE6hf2Y5LbkRs+hbI= +modernc.org/memory v1.11.0/go.mod h1:/JP4VbVC+K5sU2wZi9bHoq2MAkCnrt2r98UGeSK7Mjw= +modernc.org/opt v0.1.4 h1:2kNGMRiUjrp4LcaPuLY2PzUfqM/w9N23quVwhKt5Qm8= +modernc.org/opt v0.1.4/go.mod h1:03fq9lsNfvkYSfxrfUhZCWPk1lm4cq4N+Bh//bEtgns= +modernc.org/sortutil v1.2.1 h1:+xyoGf15mM3NMlPDnFqrteY07klSFxLElE2PVuWIJ7w= +modernc.org/sortutil v1.2.1/go.mod h1:7ZI3a3REbai7gzCLcotuw9AC4VZVpYMjDzETGsSMqJE= +modernc.org/sqlite v1.46.1 h1:eFJ2ShBLIEnUWlLy12raN0Z1plqmFX9Qe3rjQTKt6sU= +modernc.org/sqlite v1.46.1/go.mod h1:CzbrU2lSB1DKUusvwGz7rqEKIq+NUd8GWuBBZDs9/nA= +modernc.org/strutil v1.2.1 h1:UneZBkQA+DX2Rp35KcM69cSsNES9ly8mQWD71HKlOA0= +modernc.org/strutil v1.2.1/go.mod h1:EHkiggD70koQxjVdSBM3JKM7k6L0FbGE5eymy9i3B9A= +modernc.org/token v1.1.0 h1:Xl7Ap9dKaEs5kLoOQeQmPWevfnk/DM5qcLcYlA8ys6Y= +modernc.org/token v1.1.0/go.mod h1:UGzOrNV1mAFSEB63lOFHIpNRUVMvYTc6yu1SMY/XTDM= diff --git a/pkg/server/admission.go b/pkg/server/admission.go new file mode 100644 index 0000000..d613860 --- /dev/null +++ b/pkg/server/admission.go @@ -0,0 +1,157 @@ +package server + +import ( + "context" + "errors" + "fmt" + "log" + "sync" + "time" + + "github.com/ustclug/rsync-proxy/pkg/queue" + "github.com/ustclug/rsync-proxy/pkg/state" +) + +func sharedLimits(upstreams []upstreamConfig) []state.Limit { + result := make([]state.Limit, 0, len(upstreams)) + for _, u := range upstreams { + result = append(result, state.Limit{Name: u.Name, Active: u.MaxActiveConns, Queued: u.MaxQueuedConns, PerIP: u.PerIPMaxActiveConns}) + } + return result +} + +func (s *Server) Limits() []state.Limit { + s.reloadLock.RLock() + defer s.reloadLock.RUnlock() + return sharedLimits(s.upstreams) +} + +type admissionHandle struct { + c <-chan queue.Status + release func() +} +type admissionWait struct { + c chan queue.Status + last queue.Status +} + +type admission struct { + s *Server + mu sync.Mutex + waiting map[int64]*admissionWait + releases map[int64]bool + done chan struct{} +} + +func (s *Server) StartAdmission() func() { + a := &admission{s: s, waiting: make(map[int64]*admissionWait), releases: make(map[int64]bool), done: make(chan struct{})} + s.admission = a + ctx, cancel := context.WithCancel(context.Background()) + go func() { + defer close(a.done) + ticker := time.NewTicker(100 * time.Millisecond) + defer ticker.Stop() + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + a.poll() + } + } + }() + return func() { cancel(); <-a.done } +} + +func offerStatus(c chan queue.Status, status queue.Status) { + select { + case <-c: + default: + } + c <- status +} + +func (a *admission) poll() { + a.mu.Lock() + ids := make([]int64, 0, len(a.releases)) + for id := range a.releases { + ids = append(ids, id) + } + hasWaiting := len(a.waiting) > 0 + a.mu.Unlock() + for _, id := range ids { + a.release(id) + } + if !hasWaiting { + return + } + tickets, err := a.s.Shared.Poll(a.s.Generation) + if err != nil { + log.Printf("[ERROR] poll admission: %v", err) + return + } + a.mu.Lock() + defer a.mu.Unlock() + for id, c := range a.waiting { + t, ok := tickets[id] + if !ok { + continue + } + status := queue.Status{Ok: t.Active, Index: t.Index, Max: t.Queued} + if status != c.last { + offerStatus(c.c, status) + c.last = status + } + if t.Active { + delete(a.waiting, id) + } + } +} + +func (a *admission) release(id int64) { + a.mu.Lock() + delete(a.waiting, id) + a.mu.Unlock() + ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + defer cancel() + err := a.s.Shared.Release(ctx, a.s.Generation, id) + a.mu.Lock() + defer a.mu.Unlock() + if err != nil { + a.releases[id] = true + log.Printf("[ERROR] release admission (will retry): %v", err) + } else { + delete(a.releases, id) + } +} + +func (s *Server) acquire(ctx context.Context, upstream, ip string, q *queue.Queue) (admissionHandle, error) { + if s.Shared == nil { + counter := s.getPerIPCounter(upstream, ip) + limit := s.getPerIPLimitForUpstream(upstream) + if n := counter.Add(1); limit > 0 && n > int64(limit) { + counter.Add(-1) + return admissionHandle{}, fmt.Errorf("per-IP cap of %d reached: %w", limit, state.ErrPerIP) + } + h := q.Acquire() + return admissionHandle{c: h.C, release: func() { h.Release(); counter.Add(-1) }}, nil + } + t, err := s.Shared.Acquire(ctx, s.Generation, upstream, ip) + if err != nil && !errors.Is(err, state.ErrFull) { + return admissionHandle{}, err + } + c := make(chan queue.Status, 1) + if errors.Is(err, state.ErrFull) { + c <- queue.Status{Full: true} + return admissionHandle{c: c, release: func() {}}, nil + } + a := s.admission + a.mu.Lock() + if !t.Active { + a.waiting[t.ID] = &admissionWait{c: c, last: queue.Status{Ok: t.Active, Index: t.Index, Max: t.Queued}} + } + c <- queue.Status{Ok: t.Active, Index: t.Index, Max: t.Queued} + a.mu.Unlock() + var once sync.Once + return admissionHandle{c: c, release: func() { once.Do(func() { a.release(t.ID) }) }}, nil +} diff --git a/pkg/server/config.go b/pkg/server/config.go index de10a50..faaecef 100644 --- a/pkg/server/config.go +++ b/pkg/server/config.go @@ -23,15 +23,17 @@ type Upstream struct { } type ProxySettings struct { - Listen string `toml:"listen"` - ListenTLS string `toml:"listen_tls"` - ListenHTTP string `toml:"listen_http"` - Motd string `toml:"motd"` - AccessLog string `toml:"access_log"` - AccessJSONLog string `toml:"access_json_log"` - ErrorLog string `toml:"error_log"` - TLSCertFile string `toml:"tls_cert_file"` - TLSKeyFile string `toml:"tls_key_file"` + StateDir string `toml:"state_dir"` + UpgradeDrainTimeoutSecs int `toml:"upgrade_drain_timeout"` + Listen string `toml:"listen"` + ListenTLS string `toml:"listen_tls"` + ListenHTTP string `toml:"listen_http"` + Motd string `toml:"motd"` + AccessLog string `toml:"access_log"` + AccessJSONLog string `toml:"access_json_log"` + ErrorLog string `toml:"error_log"` + TLSCertFile string `toml:"tls_cert_file"` + TLSKeyFile string `toml:"tls_key_file"` // RelayIdleTimeoutSecs is the idle timeout (in seconds) applied // during the bidirectional relay phase of a connection. If no data // flows in either direction for this duration, the connection is diff --git a/pkg/server/handoff.go b/pkg/server/handoff.go new file mode 100644 index 0000000..30337e5 --- /dev/null +++ b/pkg/server/handoff.go @@ -0,0 +1,95 @@ +package server + +import ( + "fmt" + "net" + "os" + "syscall" + "time" + + "golang.org/x/sys/unix" +) + +func (s *Server) ConfiguredListeners() [3]string { return s.configuredListeners } + +// ListenInherited imports raw listeners, never TLS connection state. The +// descriptor order is plain rsync, optional TLS rsync, then management HTTP. +func (s *Server) ListenInherited(files []*os.File, expected [3]string) error { + if s.configuredListeners != expected { + return fmt.Errorf("listener addresses changed; hot upgrade requires unchanged listen/listen_tls/listen_http") + } + want := 2 + if s.TLSListenAddr != "" { + want++ + } + if len(files) != want { + return fmt.Errorf("expected %d listener descriptors, got %d", want, len(files)) + } + i := 0 + s.listenFactory = func(_ string) (net.Listener, error) { + f := files[i] + i++ + l, err := net.FileListener(f) + if u, ok := l.(*net.UnixListener); ok { + u.SetUnlinkOnClose(false) + } + return l, err + } + defer func() { s.listenFactory = nil }() + return s.Listen() +} + +func (s *Server) ListenerFiles() ([]*os.File, error) { + var files []*os.File + for _, l := range []net.Listener{s.TCPListener, s.rawTLSListener, s.HTTPListener} { + if l == nil { + continue + } + raw, ok := l.(syscall.Conn) + if !ok { + for _, f := range files { + _ = f.Close() + } + return nil, fmt.Errorf("listener %T does not support handoff", l) + } + rc, err := raw.SyscallConn() + fd := -1 + if err == nil { + var dupErr error + err = rc.Control(func(original uintptr) { fd, dupErr = unix.FcntlInt(original, unix.F_DUPFD_CLOEXEC, 0) }) + if err == nil { + err = dupErr + } + } + if err != nil { + for _, f := range files { + _ = f.Close() + } + return nil, err + } + // net.Listener.File followed by exec's File.Fd makes the shared socket + // blocking. An unsuccessful/non-Go child would then wedge the old Accept. + // NewFile preserves O_NONBLOCK when exec obtains the inherited descriptor. + f := os.NewFile(uintptr(fd), "listener") + files = append(files, f) + } + return files, nil +} + +func (s *Server) DrainTimeout() time.Duration { + s.reloadLock.RLock() + defer s.reloadLock.RUnlock() + return s.UpgradeDrainTimeout +} + +func (s *Server) ReopenLogs() error { + if err := s.accessLog.Reopen(); err != nil { + return err + } + if err := s.accessJSONLog.Reopen(); err != nil { + return err + } + return s.errorLog.Reopen() +} + +func (s *Server) CloseLogs() { s.accessLog.Close(); s.accessJSONLog.Close(); s.errorLog.Close() } diff --git a/pkg/server/metrics.go b/pkg/server/metrics.go index d8c6a01..0ad5b9a 100644 --- a/pkg/server/metrics.go +++ b/pkg/server/metrics.go @@ -24,8 +24,12 @@ func prometheusLabelValueOrUnknown(s string) string { return s } -func prometheusLabels(index uint32, module, upstream string) string { - return fmt.Sprintf( +func prometheusLabels(index uint32, module, upstream string, generation ...string) string { + prefix := "" + if len(generation) > 0 && generation[0] != "" { + prefix = fmt.Sprintf(`generation="%s",`, prometheusEscapeLabelValue(generation[0])) + } + return prefix + fmt.Sprintf( `index="%d",module="%s",upstream="%s"`, index, prometheusEscapeLabelValue(prometheusLabelValueOrUnknown(module)), @@ -170,21 +174,21 @@ func (s *Server) writePrometheusMetrics(w io.Writer, now time.Time) { _, _ = fmt.Fprintln(w, "# TYPE rsync_proxy_connection_sent_bytes gauge") for _, conn := range connections { snapshot := conn.snapshot() - _, _ = fmt.Fprintf(w, "rsync_proxy_connection_sent_bytes{%s} %d\n", prometheusLabels(snapshot.Index, snapshot.Module, snapshot.Upstream), snapshot.SentBytes) + _, _ = fmt.Fprintf(w, "rsync_proxy_connection_sent_bytes{%s} %d\n", prometheusLabels(snapshot.Index, snapshot.Module, snapshot.Upstream, snapshot.Generation), snapshot.SentBytes) } _, _ = fmt.Fprintln(w, "# HELP rsync_proxy_connection_received_bytes Bytes received from clients for active connections.") _, _ = fmt.Fprintln(w, "# TYPE rsync_proxy_connection_received_bytes gauge") for _, conn := range connections { snapshot := conn.snapshot() - _, _ = fmt.Fprintf(w, "rsync_proxy_connection_received_bytes{%s} %d\n", prometheusLabels(snapshot.Index, snapshot.Module, snapshot.Upstream), snapshot.ReceivedBytes) + _, _ = fmt.Fprintf(w, "rsync_proxy_connection_received_bytes{%s} %d\n", prometheusLabels(snapshot.Index, snapshot.Module, snapshot.Upstream, snapshot.Generation), snapshot.ReceivedBytes) } _, _ = fmt.Fprintln(w, "# HELP rsync_proxy_connection_connected_timestamp_seconds Unix timestamp when active connections were established.") _, _ = fmt.Fprintln(w, "# TYPE rsync_proxy_connection_connected_timestamp_seconds gauge") for _, conn := range connections { snapshot := conn.snapshot() - _, _ = fmt.Fprintf(w, "rsync_proxy_connection_connected_timestamp_seconds{%s} %d\n", prometheusLabels(snapshot.Index, snapshot.Module, snapshot.Upstream), snapshot.ConnectedAt.Unix()) + _, _ = fmt.Fprintf(w, "rsync_proxy_connection_connected_timestamp_seconds{%s} %d\n", prometheusLabels(snapshot.Index, snapshot.Module, snapshot.Upstream, snapshot.Generation), snapshot.ConnectedAt.Unix()) } _, _ = fmt.Fprintln(w, "# HELP rsync_proxy_connection_duration_seconds Current duration of active connections.") @@ -195,7 +199,7 @@ func (s *Server) writePrometheusMetrics(w io.Writer, now time.Time) { if duration < 0 { duration = 0 } - _, _ = fmt.Fprintf(w, "rsync_proxy_connection_duration_seconds{%s} %.3f\n", prometheusLabels(snapshot.Index, snapshot.Module, snapshot.Upstream), duration) + _, _ = fmt.Fprintf(w, "rsync_proxy_connection_duration_seconds{%s} %.3f\n", prometheusLabels(snapshot.Index, snapshot.Module, snapshot.Upstream, snapshot.Generation), duration) } _, _ = fmt.Fprintln(w, "# HELP rsync_proxy_accepted_connections_total Total accepted connections since start.") diff --git a/pkg/server/server.go b/pkg/server/server.go index a2e454b..fe6495f 100644 --- a/pkg/server/server.go +++ b/pkg/server/server.go @@ -14,6 +14,7 @@ import ( "net" "net/http" "os" + "path/filepath" "slices" "sort" "strconv" @@ -27,6 +28,7 @@ import ( "github.com/ustclug/rsync-proxy/pkg/logging" "github.com/ustclug/rsync-proxy/pkg/queue" + "github.com/ustclug/rsync-proxy/pkg/state" ) const ( @@ -55,6 +57,7 @@ const lineFeed = '\n' type ConnInfo struct { mu sync.RWMutex Index uint32 + Generation string LocalAddr string RemoteAddr string ConnectedAt time.Time @@ -66,6 +69,7 @@ type ConnInfo struct { type connInfoSnapshot struct { Index uint32 `json:"index"` + Generation string `json:"generation,omitempty"` LocalAddr string `json:"local"` RemoteAddr string `json:"remote"` ConnectedAt time.Time `json:"connected"` @@ -108,6 +112,7 @@ func (c *ConnInfo) snapshot() connInfoSnapshot { defer c.mu.RUnlock() return connInfoSnapshot{ Index: c.Index, + Generation: c.Generation, LocalAddr: c.LocalAddr, RemoteAddr: c.RemoteAddr, ConnectedAt: c.ConnectedAt, @@ -176,12 +181,31 @@ type moduleCounters struct { } type Server struct { + Started chan struct{} // --- Options section // Listen Address - ListenAddr string - TLSListenAddr string - HTTPListenAddr string - ConfigPath string + ListenAddr string + TLSListenAddr string + HTTPListenAddr string + ConfigPath string + StateDir string + UpgradeDrainTimeout time.Duration + Generation string + Shared *state.Store + // ControlHandler handles daemon lifecycle requests before the normal API. + ControlHandler func(http.ResponseWriter, *http.Request) bool + listenFactory func(string) (net.Listener, error) + rawTLSListener net.Listener + configuredListeners [3]string + lifecycleMu sync.Mutex + forceClosed bool + acceptWG sync.WaitGroup + sessions map[net.Conn]context.CancelFunc + sessionsWG sync.WaitGroup + stopOnce sync.Once + stopC chan struct{} + admission *admission + HTTPServer *http.Server ReadTimeout time.Duration WriteTimeout time.Duration @@ -288,11 +312,14 @@ func New() *Server { accessJSONLog.SetFlags(0) errorLog, _ := logging.NewFileLogger("") s := &Server{ + Started: make(chan struct{}), dialer: net.Dialer{}, // customize keep alive interval? accessLog: accessLog, accessJSONLog: accessJSONLog, errorLog: errorLog, upstreamQueues: make(map[string]*queue.Queue), + sessions: make(map[net.Conn]context.CancelFunc), + stopC: make(chan struct{}), } return s } @@ -301,6 +328,19 @@ func (s *Server) loadConfig(c *Config, openLog bool) error { var tlsCertificate *tls.Certificate serverStarted := s.TCPListener != nil || s.HTTPListener != nil || s.TLSListener != nil + if c.Proxy.UpgradeDrainTimeoutSecs < 0 || c.Proxy.RelayIdleTimeoutSecs < 0 { + return fmt.Errorf("timeouts must be non-negative") + } + if c.Proxy.StateDir != "" { + var err error + c.Proxy.StateDir, err = filepath.Abs(c.Proxy.StateDir) + if err != nil { + return err + } + } + if serverStarted && c.Proxy.StateDir != "" && c.Proxy.StateDir != s.StateDir { + return fmt.Errorf("state_dir cannot change on reload") + } if len(c.Upstreams) == 0 { return fmt.Errorf("no upstream found") } @@ -360,6 +400,9 @@ func (s *Server) loadConfig(c *Config, openLog bool) error { if len(v.Modules) == 0 && !v.DiscoverModules { return fmt.Errorf("upstream=%s must set modules or discover_modules", upstreamName) } + if v.MaxActiveConns < 0 || v.MaxQueuedConns < 0 { + return fmt.Errorf("upstream=%s: connection limits must be non-negative", upstreamName) + } if v.PerIPMaxActiveConns < 0 { return fmt.Errorf("upstream=%s: per_ip_max_active_connections must be non-negative, got %d", upstreamName, v.PerIPMaxActiveConns) } @@ -417,10 +460,17 @@ func (s *Server) loadConfig(c *Config, openLog bool) error { return err } } - s.Motd = c.Proxy.Motd - if c.Proxy.RelayIdleTimeoutSecs < 0 { - return fmt.Errorf("relay_idle_timeout must be non-negative, got %d", c.Proxy.RelayIdleTimeoutSecs) + if s.Shared != nil { + if err := s.Shared.SetLimits(s.Generation, sharedLimits(resolvedUpstreams)); err != nil { + return err + } } + if !serverStarted { + s.configuredListeners = [3]string{c.Proxy.Listen, c.Proxy.ListenTLS, c.Proxy.ListenHTTP} + s.StateDir = c.Proxy.StateDir + } + s.UpgradeDrainTimeout = time.Duration(c.Proxy.UpgradeDrainTimeoutSecs) * time.Second + s.Motd = c.Proxy.Motd s.RelayIdleTimeout = time.Duration(c.Proxy.RelayIdleTimeoutSecs) * time.Second s.RelayMaxDuration = time.Duration(c.Proxy.RelayMaxDurationSecs) * time.Second s.TCPKeepAlive = time.Duration(c.Proxy.TCPKeepAliveSecs) * time.Second @@ -806,6 +856,7 @@ func (s *Server) relay(ctx context.Context, index uint32, downConn net.Conn) (re info := ConnInfo{ Index: index, + Generation: s.Generation, LocalAddr: downConn.LocalAddr().String(), RemoteAddr: downConn.RemoteAddr().String(), ConnectedAt: time.Now().Truncate(time.Second), @@ -917,27 +968,24 @@ func (s *Server) relay(ctx context.Context, index uint32, downConn net.Conn) (re return fmt.Errorf("no queue configured for upstream %s", target.Upstream) } - // Per-IP per-upstream concurrency cap. A positive cap means the - // same client IP may not have more than N simultaneous active - // relay connections to this upstream. Counted before queue - // admission so the cap also bounds queueing, preventing a single - // IP from monopolizing both the active slots and the queue. - if perIPLimit := s.getPerIPLimitForUpstream(target.Upstream); perIPLimit > 0 { - counter := s.getPerIPCounter(target.Upstream, ip) - if n := counter.Add(1); n > int64(perIPLimit) { - counter.Add(-1) - s.getUpstreamCounters(target.Upstream).perIPRejected.Add(1) - record.Status = http.StatusTooManyRequests - s.accessLog.F("client %s rejected for upstream %s module %s: per-IP cap of %d reached", ip, target.Upstream, moduleName, perIPLimit) - _, _ = writeWithTimeout(downConn, fmt.Appendf(nil, "@ERROR: per-IP connection limit of %d for upstream %s reached, retry later\n", perIPLimit, target.Upstream), writeTimeout) - return nil - } - defer counter.Add(-1) + handle, err := s.acquire(ctx, target.Upstream, ip, upstreamQueue) + if errors.Is(err, state.ErrPerIP) { + s.getUpstreamCounters(target.Upstream).perIPRejected.Add(1) + record.Status = http.StatusTooManyRequests + s.accessLog.F("client %s rejected for upstream %s module %s: %v", ip, target.Upstream, moduleName, err) + _, _ = writeWithTimeout(downConn, fmt.Appendf(nil, "@ERROR: per-IP connection limit for upstream %s reached, retry later\n", target.Upstream), writeTimeout) + return nil + } + if err != nil { + return fmt.Errorf("admit connection: %w", err) + } + defer handle.release() + var status queue.Status + select { + case status = <-handle.c: + case <-ctx.Done(): + return ctx.Err() } - - handle := upstreamQueue.Acquire() - defer handle.Release() - status := <-handle.C if status.Full { s.getUpstreamCounters(target.Upstream).queueFull.Add(1) record.Status = http.StatusServiceUnavailable @@ -959,7 +1007,9 @@ func (s *Server) relay(ctx context.Context, index uint32, downConn net.Conn) (re queuing: for !status.Ok { select { - case status = <-handle.C: + case <-ctx.Done(): + return ctx.Err() + case status = <-handle.c: if status.Ok { break queuing } @@ -981,6 +1031,8 @@ func (s *Server) relay(ctx context.Context, index uint32, downConn net.Conn) (re return fmt.Errorf("dial to upstream: %s: %w", upstreamAddr, err) } defer upConn.Close() + stopUpstream := context.AfterFunc(ctx, func() { _ = upConn.Close() }) + defer stopUpstream() record.ServerEndpoint = upConn.RemoteAddr().String() // Enable TCP keepalive on the upstream-side connection so that a // dead/half-open peer is detected within the configured period @@ -1325,13 +1377,12 @@ func (s *Server) runHTTPServer() error { return } - var status struct { - Count int `json:"count"` - Connections []*ConnInfo `json:"connections"` + status, err := s.Status() + if err != nil { + http.Error(w, err.Error(), http.StatusServiceUnavailable) + return } - status.Connections = s.ListConnectionInfo() - status.Count = len(status.Connections) - _ = json.NewEncoder(w).Encode(&status) + _ = json.NewEncoder(w).Encode(status) }) mux.HandleFunc("/telegraf", func(w http.ResponseWriter, r *http.Request) { @@ -1342,6 +1393,14 @@ func (s *Server) runHTTPServer() error { timestamp := time.Now().Truncate(time.Second).UnixNano() count := s.GetActiveConnectionCount() + if s.Shared != nil { + status, err := s.Status() + if err != nil { + http.Error(w, err.Error(), http.StatusServiceUnavailable) + return + } + count = int64(status.Count) + } // https://docs.influxdata.com/influxdb/latest/reference/syntax/line-protocol/ _, _ = fmt.Fprintf(w, "rsync-proxy,host=%s count=%d %d\n", hostname, count, timestamp) }) @@ -1352,20 +1411,46 @@ func (s *Server) runHTTPServer() error { return } - // promhttp.HandlerFor sets the Content-Type itself based on - // content negotiation; do not pre-set it here. - // EnableOpenMetrics is disabled so that no "# EOF" terminator is - // emitted, allowing us to append our own legacy text-format - // metrics after the runtime/process metrics. + var metrics bytes.Buffer + if s.Shared != nil { + if err := s.AggregateMetrics(&metrics); err != nil { + http.Error(w, err.Error(), http.StatusServiceUnavailable) + return + } + } else { + s.writePrometheusMetrics(&metrics, time.Now()) + } promhttp.HandlerFor(prometheus.DefaultGatherer, promhttp.HandlerOpts{DisableCompression: true, EnableOpenMetrics: false}).ServeHTTP(w, r) - s.writePrometheusMetrics(w, time.Now()) + _, _ = w.Write(metrics.Bytes()) }) - return http.Serve(s.HTTPListener, &mux) + s.HTTPServer = &http.Server{Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if s.ControlHandler != nil && s.ControlHandler(w, r) { + return + } + mux.ServeHTTP(w, r) + })} + close(s.Started) + return s.HTTPServer.Serve(s.HTTPListener) } func (s *Server) Listen() error { - l1, err := listenTCPOrUnix(s.ListenAddr) + baseListen := s.listenFactory + if baseListen == nil { + baseListen = listenTCPOrUnix + } + listen := func(addr string) (net.Listener, error) { + l, err := baseListen(addr) + // A shared listener's path belongs to the serving daemon, not an individual + // generation. Set this before starting Accept/Close goroutines. + if s.Shared != nil { + if u, ok := l.(*net.UnixListener); ok { + u.SetUnlinkOnClose(false) + } + } + return l, err + } + l1, err := listen(s.ListenAddr) if err != nil { return fmt.Errorf("create tcp listener: %w", err) } @@ -1374,17 +1459,18 @@ func (s *Server) Listen() error { var lTLS net.Listener if s.TLSListenAddr != "" { - lTLS, err = listenTCPOrUnix(s.TLSListenAddr) + lTLS, err = listen(s.TLSListenAddr) if err != nil { _ = l1.Close() return fmt.Errorf("create tls listener: %w", err) } s.TLSListenAddr = lTLS.Addr().String() log.Printf("[INFO] Rsync TLS proxy listening on %s", s.TLSListenAddr) + s.rawTLSListener = lTLS lTLS = tls.NewListener(lTLS, &tls.Config{GetCertificate: s.getTLSCertificate}) } - l2, err := listenTCPOrUnix(s.HTTPListenAddr) + l2, err := listen(s.HTTPListenAddr) if err != nil { _ = l1.Close() if lTLS != nil { @@ -1401,20 +1487,32 @@ func (s *Server) Listen() error { return nil } -func (s *Server) Close() { - if s.TCPListener != nil { - _ = s.TCPListener.Close() - } - if s.TLSListener != nil { - _ = s.TLSListener.Close() - } - if s.HTTPListener != nil { - _ = s.HTTPListener.Close() +// StopAccepting starts a drain without cancelling accepted sessions. +func (s *Server) StopAccepting() { + s.stopOnce.Do(func() { + for _, l := range []net.Listener{s.TCPListener, s.TLSListener, s.HTTPListener} { + if l != nil { + _ = l.Close() + } + } + close(s.stopC) + }) +} + +// Close stops admission and aborts all accepted sessions, including queue waits. +func (s *Server) Close() { s.StopAccepting(); s.AbortConnections() } + +func (s *Server) AbortConnections() { + s.lifecycleMu.Lock() + defer s.lifecycleMu.Unlock() + s.forceClosed = true + for conn, cancel := range s.sessions { + cancel() + _ = conn.Close() } } func (s *Server) handleConn(ctx context.Context, conn net.Conn) { - s.activeConnCount.Add(1) defer s.activeConnCount.Add(-1) s.acceptedConnCount.Add(1) connIndex := s.connIndex.Add(1) @@ -1442,21 +1540,42 @@ func (s *Server) runRsyncServer(ctx context.Context, listener net.Listener, acce for { conn, err := listener.Accept() if err != nil { - if errors.Is(err, net.ErrClosed) { + if errors.Is(err, net.ErrClosed) || errors.Is(err, http.ErrServerClosed) { return nil } return fmt.Errorf("%s: %w", acceptErr, err) } - go s.handleConn(ctx, conn) + s.lifecycleMu.Lock() + if s.forceClosed { + s.lifecycleMu.Unlock() + _ = conn.Close() + continue + } + connCtx, cancel := context.WithCancel(ctx) + s.sessions[conn] = cancel + s.sessionsWG.Add(1) + s.activeConnCount.Add(1) + s.lifecycleMu.Unlock() + go func() { + defer s.sessionsWG.Done() + defer func() { + cancel() + _ = conn.Close() + s.lifecycleMu.Lock() + delete(s.sessions, conn) + s.lifecycleMu.Unlock() + }() + s.handleConn(connCtx, conn) + }() } } func (s *Server) Run() error { - errC := make(chan error, 1) + errC := make(chan error, 3) go func() { err := s.runHTTPServer() if err != nil { - if errors.Is(err, net.ErrClosed) { + if errors.Is(err, net.ErrClosed) || errors.Is(err, http.ErrServerClosed) { return } errC <- fmt.Errorf("start http server: %w", err) @@ -1465,14 +1584,18 @@ func (s *Server) Run() error { ctx, cancel := context.WithCancel(context.Background()) defer cancel() + s.acceptWG.Add(1) go func() { + defer s.acceptWG.Done() err := s.runRsyncServer(ctx, s.TCPListener, "accept rsync connection") if err != nil { errC <- err } }() if s.TLSListener != nil { + s.acceptWG.Add(1) go func() { + defer s.acceptWG.Done() err := s.runRsyncServer(ctx, s.TLSListener, "accept tls rsync connection") if err != nil { errC <- err @@ -1480,10 +1603,15 @@ func (s *Server) Run() error { }() } - for { - err := <-errC - if err != nil { - return err - } + select { + case err := <-errC: + s.Close() + s.acceptWG.Wait() + s.sessionsWG.Wait() + return err + case <-s.stopC: + s.acceptWG.Wait() + s.sessionsWG.Wait() + return nil } } diff --git a/pkg/server/snapshot.go b/pkg/server/snapshot.go new file mode 100644 index 0000000..336b6e0 --- /dev/null +++ b/pkg/server/snapshot.go @@ -0,0 +1,168 @@ +package server + +import ( + "bytes" + "encoding/json" + "fmt" + "io" + "sort" + "strings" + "time" + + dto "github.com/prometheus/client_model/go" + "github.com/prometheus/common/expfmt" + "github.com/prometheus/common/model" + + "github.com/ustclug/rsync-proxy/pkg/state" +) + +type generationSnapshot struct { + Updated time.Time `json:"updated"` + Connections []connInfoSnapshot `json:"connections"` + Count int64 `json:"count"` + Metrics string `json:"metrics"` +} + +type Status struct { + Count int `json:"count"` + Connections []connInfoSnapshot `json:"connections"` + Generations []state.Generation `json:"generations,omitempty"` +} + +func (s *Server) snapshot() generationSnapshot { + result := generationSnapshot{Updated: time.Now(), Count: s.GetActiveConnectionCount(), Connections: make([]connInfoSnapshot, 0)} + for _, c := range s.ListConnectionInfo() { + result.Connections = append(result.Connections, c.snapshot()) + } + var b bytes.Buffer + s.writePrometheusMetrics(&b, result.Updated) + result.Metrics = b.String() + return result +} + +func (s *Server) PublishSnapshot() error { + b, err := json.Marshal(s.snapshot()) + if err != nil { + return err + } + return s.Shared.Publish(s.Generation, b) +} + +func (s *Server) Status() (Status, error) { + result := Status{Connections: make([]connInfoSnapshot, 0)} + if s.Shared == nil { + result.Connections = s.snapshot().Connections + } else { + generations, err := s.Shared.Generations() + if err != nil { + return result, err + } + result.Generations = generations + for _, g := range generations { + if g.State == state.Exited || g.State == state.Starting { + continue + } + var snap generationSnapshot + if g.ID == s.Generation { + snap = s.snapshot() + } else if len(g.Snapshot) > 0 { + if err := json.Unmarshal(g.Snapshot, &snap); err != nil { + return result, err + } + } + result.Connections = append(result.Connections, snap.Connections...) + } + } + result.Count = len(result.Connections) + return result, nil +} + +// Aggregate only proxy metrics. Runtime/process metrics belong to the serving +// process. Retired generations contribute counters, never stale live gauges. +func (s *Server) AggregateMetrics(w io.Writer) error { + gens, err := s.Shared.Generations() + if err != nil { + return err + } + families := make(map[string]*dto.MetricFamily) + samples := make(map[string]*dto.Metric) + skip := map[string]bool{"rsync_proxy_queued_connections": true, "rsync_proxy_queue_active_max": true, "rsync_proxy_queue_queued_max": true} + for _, g := range gens { + var snap generationSnapshot + if g.ID == s.Generation { + snap = s.snapshot() + } else if len(g.Snapshot) > 0 { + if err := json.Unmarshal(g.Snapshot, &snap); err != nil { + return err + } + } + parser := expfmt.NewTextParser(model.LegacyValidation) + parsed, err := parser.TextToMetricFamilies(strings.NewReader(snap.Metrics)) + if err != nil { + return err + } + for name, f := range parsed { + if skip[name] || (g.State == state.Exited && f.GetType() != dto.MetricType_COUNTER) { + continue + } + dst := families[name] + if dst == nil { + dst = &dto.MetricFamily{Name: f.Name, Help: f.Help, Type: f.Type} + families[name] = dst + } + for _, m := range f.Metric { + sort.Slice(m.Label, func(i, j int) bool { return m.Label[i].GetName() < m.Label[j].GetName() }) + labels, _ := json.Marshal(m.Label) + key := name + string(labels) + if old := samples[key]; old != nil { + if m.Counter != nil { + v := old.Counter.GetValue() + m.Counter.GetValue() + old.Counter.Value = &v + } + if m.Gauge != nil { + v := old.Gauge.GetValue() + m.Gauge.GetValue() + old.Gauge.Value = &v + } + } else { + samples[key] = m + dst.Metric = append(dst.Metric, m) + } + } + } + } + names := make([]string, 0, len(families)) + for name := range families { + names = append(names, name) + } + sort.Strings(names) + for _, name := range names { + if _, err := expfmt.MetricFamilyToText(w, families[name]); err != nil { + return err + } + } + queues, err := s.Shared.Queues() + if err != nil { + return err + } + for _, metric := range []struct { + name, help string + value func(state.QueueInfo) int + }{ + {"rsync_proxy_queued_connections", "Current queued rsync proxy connections per upstream.", func(q state.QueueInfo) int { return q.QueuedCount }}, + {"rsync_proxy_queue_active_max", "Configured max active connections per upstream.", func(q state.QueueInfo) int { return q.Active }}, + {"rsync_proxy_queue_queued_max", "Configured max queued connections per upstream.", func(q state.QueueInfo) int { return q.Queued }}, + } { + _, _ = fmt.Fprintf(w, "# HELP %s %s\n# TYPE %s gauge\n", metric.name, metric.help, metric.name) + for _, q := range queues { + _, _ = fmt.Fprintf(w, "%s{upstream=\"%s\"} %d\n", metric.name, prometheusEscapeLabelValue(q.Name), metric.value(q)) + } + } + _, _ = fmt.Fprintln(w, "# HELP rsync_proxy_generation_info Process generation metadata.\n# TYPE rsync_proxy_generation_info gauge") + for _, g := range gens { + if g.State == state.Exited { + continue + } + _, _ = fmt.Fprintf(w, "rsync_proxy_generation_info{generation=\"%s\",version=\"%s\",state=\"%s\",pid=\"%d\"} 1\n", prometheusEscapeLabelValue(g.ID), prometheusEscapeLabelValue(g.Version), g.State, g.PID) + } + return nil +} diff --git a/pkg/server/snapshot_test.go b/pkg/server/snapshot_test.go new file mode 100644 index 0000000..007c27a --- /dev/null +++ b/pkg/server/snapshot_test.go @@ -0,0 +1,59 @@ +package server + +import ( + "bytes" + "testing" + + "github.com/prometheus/common/expfmt" + "github.com/prometheus/common/model" + "github.com/stretchr/testify/require" + + "github.com/ustclug/rsync-proxy/pkg/queue" + "github.com/ustclug/rsync-proxy/pkg/state" +) + +func TestAggregateRetainsCountersWithoutRetiredGauges(t *testing.T) { + store, err := state.Open(t.TempDir()) + require.NoError(t, err) + defer store.Close() + makeGeneration := func(id string) *Server { + s := New() + s.Shared, s.Generation = store, id + s.upstreams = []upstreamConfig{{Name: "u"}} + s.upstreamQueues["u"] = queue.New(5, 2) + s.connInfo.Store(uint32(1), &ConnInfo{Index: 1, Generation: id, Module: "foo", Upstream: "u"}) + s.activeConnCount.Store(1) + s.acceptedConnCount.Store(3) + s.completedConnCount.Store(2) + s.sentBytesTotal.Store(42) + require.NoError(t, store.Register(state.Generation{ID: id, Version: id})) + return s + } + a, b := makeGeneration("a"), makeGeneration("b") + limits := []state.Limit{{Name: "u", Active: 5, Queued: 2}} + require.NoError(t, store.Activate("a", "", limits)) + require.NoError(t, a.PublishSnapshot()) + require.NoError(t, store.Activate("b", "a", limits)) + for _, retired := range []bool{false, true} { + if retired { + require.NoError(t, store.Finish("a")) + } + var output bytes.Buffer + require.NoError(t, b.AggregateMetrics(&output)) + parser := expfmt.NewTextParser(model.LegacyValidation) + metrics, err := parser.TextToMetricFamilies(&output) + require.NoError(t, err) + require.InDelta(t, float64(6), metrics["rsync_proxy_accepted_connections_total"].Metric[0].Counter.GetValue(), 0.001) + require.InDelta(t, float64(84), metrics["rsync_proxy_sent_bytes_total"].Metric[0].Counter.GetValue(), 0.001) + require.InDelta(t, float64(5), metrics["rsync_proxy_queue_active_max"].Metric[0].Gauge.GetValue(), 0.001) + want := 2 + if retired { + want = 1 + } + require.InDelta(t, float64(want), metrics["rsync_proxy_active_connections"].Metric[0].Gauge.GetValue(), 0.001) + require.Len(t, metrics["rsync_proxy_connection_sent_bytes"].Metric, want) + status, err := b.Status() + require.NoError(t, err) + require.Equal(t, want, status.Count) + } +} diff --git a/pkg/server/utils.go b/pkg/server/utils.go index c5381ae..026e971 100644 --- a/pkg/server/utils.go +++ b/pkg/server/utils.go @@ -2,10 +2,12 @@ package server import ( "context" + "errors" "fmt" "net" "os" "strings" + "syscall" "time" ) @@ -91,13 +93,35 @@ func readLine(conn net.Conn, buf []byte, timeout time.Duration) (n int, err erro func listenTCPOrUnix(addr string) (net.Listener, error) { if strings.HasPrefix(addr, "/") { - os.Remove(addr) l, err := net.Listen("unix", addr) + if errors.Is(err, syscall.EADDRINUSE) { + // Only reclaim a stale socket. Never unlink an active listener, a symlink, + // or an ordinary file belonging to another instance. + info, statErr := os.Lstat(addr) + if statErr == nil && info.Mode()&os.ModeSocket != 0 { + conn, dialErr := net.DialTimeout("unix", addr, 200*time.Millisecond) + if dialErr == nil { + _ = conn.Close() + } else if errors.Is(dialErr, syscall.ECONNREFUSED) { + current, statErr := os.Lstat(addr) + if statErr == nil && os.SameFile(info, current) { + if removeErr := os.Remove(addr); removeErr != nil { + return nil, removeErr + } + l, err = net.Listen("unix", addr) + } + } + } + } if err != nil { return l, err } err = os.Chmod(addr, 0o660) - return l, err + if err != nil { + _ = l.Close() + return nil, err + } + return l, nil } return net.Listen("tcp", addr) } diff --git a/pkg/server/utils_test.go b/pkg/server/utils_test.go index 26963c3..24edcdd 100644 --- a/pkg/server/utils_test.go +++ b/pkg/server/utils_test.go @@ -116,3 +116,26 @@ func TestListenAndDialUnixSocket(t *testing.T) { require.NoError(t, err) assert.Equal(t, info.String(), netAddrToString(info)) } + +func TestUnixListenerDoesNotReplaceLiveSocket(t *testing.T) { + addr := filepath.Join(t.TempDir(), "live.sock") + first, err := listenTCPOrUnix(addr) + require.NoError(t, err) + defer first.Close() + _, err = listenTCPOrUnix(addr) + require.Error(t, err) + conn, err := net.Dial("unix", addr) + require.NoError(t, err) + defer conn.Close() +} + +func TestUnixListenerReclaimsStaleSocket(t *testing.T) { + addr := filepath.Join(t.TempDir(), "stale.sock") + first, err := net.Listen("unix", addr) + require.NoError(t, err) + first.(*net.UnixListener).SetUnlinkOnClose(false) + require.NoError(t, first.Close()) + second, err := listenTCPOrUnix(addr) + require.NoError(t, err) + defer second.Close() +} diff --git a/pkg/state/lock.go b/pkg/state/lock.go new file mode 100644 index 0000000..f826f71 --- /dev/null +++ b/pkg/state/lock.go @@ -0,0 +1,34 @@ +package state + +import ( + "errors" + "os" + + "golang.org/x/sys/unix" +) + +// Lock is an open-file-description lock. It must not be passed to children. +// Closing the file (including process death) releases it; no PID/heartbeat +// heuristic can mistakenly reclaim a still-live generation's tickets. +func Lock(path string) (*os.File, error) { + f, err := os.OpenFile(path, os.O_CREATE|os.O_RDWR, 0600) + if err != nil { + return nil, err + } + if err = unix.Flock(int(f.Fd()), unix.LOCK_EX|unix.LOCK_NB); err != nil { + _ = f.Close() + return nil, err + } + return f, nil +} + +func IsLocked(path string) (bool, error) { + f, err := Lock(path) + if errors.Is(err, unix.EWOULDBLOCK) { + return true, nil + } + if err != nil { + return false, err + } + return false, f.Close() +} diff --git a/pkg/state/store.go b/pkg/state/store.go new file mode 100644 index 0000000..a1f7f09 --- /dev/null +++ b/pkg/state/store.go @@ -0,0 +1,351 @@ +// Package state coordinates admission across generations of one local proxy. +// The database contains runtime state, not rsync payloads or durable history. +package state + +import ( + "context" + "database/sql" + "encoding/json" + "errors" + "fmt" + "net/url" + "os" + "path/filepath" + "time" + + _ "modernc.org/sqlite" // Register the CGo-free database/sql driver. +) + +const Protocol = 1 + +const ( + Starting = "starting" + Serving = "serving" + Draining = "draining" + Exited = "exited" +) + +var ( + ErrPerIP = errors.New("per-IP connection limit reached") + ErrFull = errors.New("upstream queue is full") + ErrNotCurrent = errors.New("generation is no longer current") +) + +type Limit struct { + Name string + Active, Queued, PerIP int +} + +type Generation struct { + ID string `json:"id"` + PID int `json:"pid"` + Version string `json:"version"` + Control string `json:"control"` + State string `json:"state"` + Deadline int64 `json:"deadline,omitempty"` + Snapshot json.RawMessage `json:"-"` +} + +type Ticket struct { + ID int64 + Active bool + Index, Queued, Max int +} + +type QueueInfo struct { + Limit + ActiveCount, QueuedCount int +} + +type Store struct { + db *sql.DB +} + +func Open(dir string) (*Store, error) { + if err := os.MkdirAll(dir, 0700); err != nil { + return nil, err + } + u := url.URL{Scheme: "file", Path: filepath.Join(dir, "state.db")} + q := u.Query() + q.Add("_pragma", "busy_timeout(2000)") + q.Add("_pragma", "foreign_keys(1)") + u.RawQuery = q.Encode() + db, err := sql.Open("sqlite", u.String()) + if err != nil { + return nil, err + } + db.SetMaxOpenConns(1) + s := &Store{db: db} + err = s.write(context.Background(), func(c *sql.Conn) error { + var version int + if err := c.QueryRowContext(context.Background(), "PRAGMA user_version").Scan(&version); err != nil { + return err + } + if version != 0 && version != Protocol { + return fmt.Errorf("incompatible state schema %d (supported %d)", version, Protocol) + } + if version == Protocol { + return nil + } + _, err := c.ExecContext(context.Background(), ` +CREATE TABLE generations(id TEXT PRIMARY KEY, pid INTEGER NOT NULL, version TEXT NOT NULL, control TEXT NOT NULL, state TEXT NOT NULL, deadline INTEGER NOT NULL DEFAULT 0, snapshot BLOB); +CREATE TABLE settings(key TEXT PRIMARY KEY, value TEXT NOT NULL); +INSERT INTO settings VALUES('current',''); +CREATE TABLE limits(name TEXT PRIMARY KEY, active INTEGER NOT NULL, queued INTEGER NOT NULL, per_ip INTEGER NOT NULL); +CREATE TABLE tickets(id INTEGER PRIMARY KEY AUTOINCREMENT, generation TEXT NOT NULL REFERENCES generations(id), upstream TEXT NOT NULL REFERENCES limits(name), ip TEXT NOT NULL, active INTEGER NOT NULL); +CREATE INDEX tickets_upstream ON tickets(upstream, active, id); +CREATE INDEX tickets_ip ON tickets(upstream, ip); +CREATE INDEX tickets_generation ON tickets(generation); +PRAGMA user_version=1;`) + return err + }) + if err == nil { + err = os.Chmod(filepath.Join(dir, "state.db"), 0600) + } + if err == nil { + _, err = db.Exec("PRAGMA journal_mode=WAL") + } + if err != nil { + _ = db.Close() + return nil, err + } + return s, nil +} + +func (s *Store) Close() error { return s.db.Close() } + +// Every read/modify/write operation starts as a writer. In particular, no +// deferred read transaction may race another process between COUNT and INSERT. +func (s *Store) write(ctx context.Context, f func(*sql.Conn) error) error { + c, err := s.db.Conn(ctx) + if err != nil { + return err + } + defer c.Close() + if _, err = c.ExecContext(ctx, "BEGIN IMMEDIATE"); err != nil { + return err + } + defer func() { _, _ = c.ExecContext(context.Background(), "ROLLBACK") }() + if err = f(c); err != nil { + return err + } + // Once commit starts, do not let request cancellation make its outcome + // ambiguous to the caller (which must release any successfully issued ticket). + if err = ctx.Err(); err != nil { + return err + } + _, err = c.ExecContext(context.Background(), "COMMIT") + return err +} + +func (s *Store) Register(g Generation) error { + _, err := s.db.Exec("INSERT INTO generations(id,pid,version,control,state) VALUES(?,?,?,?,?)", g.ID, g.PID, g.Version, g.Control, "starting") + return err +} + +func (s *Store) Current() (string, error) { + var id string + err := s.db.QueryRow("SELECT value FROM settings WHERE key='current'").Scan(&id) + return id, err +} + +func setLimits(ctx context.Context, c *sql.Conn, limits []Limit) error { + for _, l := range limits { + if l.Active < 0 || l.Queued < 0 || l.PerIP < 0 { + return errors.New("negative connection limit") + } + _, err := c.ExecContext(ctx, "INSERT INTO limits VALUES(?,?,?,?) ON CONFLICT(name) DO UPDATE SET active=excluded.active, queued=excluded.queued, per_ip=excluded.per_ip", l.Name, l.Active, l.Queued, l.PerIP) + if err != nil { + return err + } + } + return promote(ctx, c) +} + +// Activate is the commit point of an upgrade. Old generations retain their +// selected targets but consult these shared limits for every admission. +func (s *Store) Activate(id, parent string, limits []Limit) error { + ctx := context.Background() + return s.write(ctx, func(c *sql.Conn) error { + var current string + if err := c.QueryRowContext(ctx, "SELECT value FROM settings WHERE key='current'").Scan(¤t); err != nil { + return err + } + if current != parent { + return ErrNotCurrent + } + if err := setLimits(ctx, c, limits); err != nil { + return err + } + if _, err := c.ExecContext(ctx, "UPDATE settings SET value=? WHERE key='current'", id); err != nil { + return err + } + _, err := c.ExecContext(ctx, "UPDATE generations SET state='serving' WHERE id=?", id) + return err + }) +} + +func (s *Store) SetLimits(id string, limits []Limit) error { + ctx := context.Background() + return s.write(ctx, func(c *sql.Conn) error { + var current string + if err := c.QueryRowContext(ctx, "SELECT value FROM settings WHERE key='current'").Scan(¤t); err != nil { + return err + } + if current != id { + return ErrNotCurrent + } + return setLimits(ctx, c, limits) + }) +} + +func promote(ctx context.Context, c *sql.Conn) error { + // Rank only queued tickets. Already granted tickets (including grants whose + // owner hasn't polled yet) always consume capacity. + _, err := c.ExecContext(ctx, `WITH waiting AS ( + SELECT t.id, l.active AS max, ROW_NUMBER() OVER(PARTITION BY t.upstream ORDER BY t.id) AS position, + (SELECT COUNT(*) FROM tickets a WHERE a.upstream=t.upstream AND a.active=1) AS occupied + FROM tickets t JOIN limits l ON l.name=t.upstream WHERE t.active=0 +) UPDATE tickets SET active=1 WHERE id IN (SELECT id FROM waiting WHERE max=0 OR position<=max-occupied)`) + return err +} + +func (s *Store) Acquire(ctx context.Context, generation, upstream, ip string) (Ticket, error) { + var t Ticket + err := s.write(ctx, func(c *sql.Conn) error { + var max, queuedMax, perIP, active, queued, ips int + if err := c.QueryRowContext(ctx, "SELECT active,queued,per_ip FROM limits WHERE name=?", upstream).Scan(&max, &queuedMax, &perIP); err != nil { + return err + } + if err := c.QueryRowContext(ctx, "SELECT COUNT(*) FROM tickets WHERE upstream=? AND ip=?", upstream, ip).Scan(&ips); err != nil { + return err + } + if perIP > 0 && ips >= perIP { + return fmt.Errorf("per-IP cap of %d reached: %w", perIP, ErrPerIP) + } + if err := promote(ctx, c); err != nil { + return err + } + if err := c.QueryRowContext(ctx, "SELECT COUNT(*) FILTER(WHERE active=1),COUNT(*) FILTER(WHERE active=0) FROM tickets WHERE upstream=?", upstream).Scan(&active, &queued); err != nil { + return err + } + t.Active = max == 0 || (active < max && queued == 0) + if !t.Active && queuedMax > 0 && queued >= queuedMax { + return ErrFull + } + r, err := c.ExecContext(ctx, "INSERT INTO tickets(generation,upstream,ip,active) VALUES(?,?,?,?)", generation, upstream, ip, t.Active) + if err != nil { + return err + } + t.ID, err = r.LastInsertId() + t.Max = max + t.Index = queued + t.Queued = queued + 1 + return err + }) + return t, err +} + +func (s *Store) Release(ctx context.Context, generation string, id int64) error { + return s.write(ctx, func(c *sql.Conn) error { + if _, err := c.ExecContext(ctx, "DELETE FROM tickets WHERE generation=? AND id=?", generation, id); err != nil { + return err + } + return promote(ctx, c) + }) +} + +// Poll reads all of a generation's queued/granted tickets in one operation. +func (s *Store) Poll(generation string) (map[int64]Ticket, error) { + rows, err := s.db.Query(`WITH ranked AS ( + SELECT t.id,t.generation,t.active,l.active AS max, + COUNT(*) FILTER(WHERE t.active=0) OVER(PARTITION BY t.upstream ORDER BY t.id ROWS BETWEEN UNBOUNDED PRECEDING AND 1 PRECEDING) AS position, + COUNT(*) FILTER(WHERE t.active=0) OVER(PARTITION BY t.upstream) AS queued + FROM tickets t JOIN limits l ON l.name=t.upstream +) SELECT id,active,max,position,queued FROM ranked WHERE generation=?`, generation) + if err != nil { + return nil, err + } + defer rows.Close() + result := make(map[int64]Ticket) + for rows.Next() { + var t Ticket + if err := rows.Scan(&t.ID, &t.Active, &t.Max, &t.Index, &t.Queued); err != nil { + return nil, err + } + result[t.ID] = t + } + return result, rows.Err() +} + +func (s *Store) Queues() ([]QueueInfo, error) { + rows, err := s.db.Query(`SELECT l.name,l.active,l.queued,l.per_ip,COUNT(t.id) FILTER(WHERE t.active=1),COUNT(t.id) FILTER(WHERE t.active=0) FROM limits l LEFT JOIN tickets t ON t.upstream=l.name GROUP BY l.name ORDER BY l.name`) + if err != nil { + return nil, err + } + defer rows.Close() + var result []QueueInfo + for rows.Next() { + var q QueueInfo + if err := rows.Scan(&q.Name, &q.Active, &q.Queued, &q.PerIP, &q.ActiveCount, &q.QueuedCount); err != nil { + return nil, err + } + result = append(result, q) + } + return result, rows.Err() +} + +func (s *Store) Generations() ([]Generation, error) { + rows, err := s.db.Query("SELECT id,pid,version,control,state,deadline,snapshot FROM generations ORDER BY rowid") + if err != nil { + return nil, err + } + defer rows.Close() + var result []Generation + for rows.Next() { + var g Generation + var snapshot []byte + if err := rows.Scan(&g.ID, &g.PID, &g.Version, &g.Control, &g.State, &g.Deadline, &snapshot); err != nil { + return nil, err + } + g.Snapshot = json.RawMessage(snapshot) + result = append(result, g) + } + return result, rows.Err() +} + +func (s *Store) Publish(id string, snapshot []byte) error { + _, err := s.db.Exec("UPDATE generations SET snapshot=? WHERE id=?", snapshot, id) + return err +} + +func (s *Store) Drain(id string, deadline time.Time) error { + var seconds int64 + if !deadline.IsZero() { + seconds = deadline.Unix() + } + _, err := s.db.Exec("UPDATE generations SET state='draining',deadline=? WHERE id=? AND state='serving'", seconds, id) + return err +} + +func (s *Store) Finish(id string) error { + ctx := context.Background() + return s.write(ctx, func(c *sql.Conn) error { + if _, err := c.ExecContext(ctx, "DELETE FROM tickets WHERE generation=?", id); err != nil { + return err + } + if _, err := c.ExecContext(ctx, "UPDATE generations SET state='exited' WHERE id=?", id); err != nil { + return err + } + return promote(ctx, c) + }) +} + +// Reset is only called under the startup lock after confirming all generation +// locks are unheld. It deliberately preserves the schema and database inode. +func (s *Store) Reset() error { + return s.write(context.Background(), func(c *sql.Conn) error { + _, err := c.ExecContext(context.Background(), "DELETE FROM tickets; DELETE FROM generations; DELETE FROM limits; UPDATE settings SET value='' WHERE key='current';") + return err + }) +} diff --git a/pkg/state/store_test.go b/pkg/state/store_test.go new file mode 100644 index 0000000..ef3a86a --- /dev/null +++ b/pkg/state/store_test.go @@ -0,0 +1,246 @@ +package state + +import ( + "context" + "errors" + "fmt" + "os" + "os/exec" + "path/filepath" + "sort" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +func openTestStore(t *testing.T, dir, id string) *Store { + t.Helper() + s, err := Open(dir) + require.NoError(t, err) + t.Cleanup(func() { _ = s.Close() }) + require.NoError(t, s.Register(Generation{ID: id, PID: os.Getpid()})) + return s +} + +func TestGlobalFIFOAndPolicy(t *testing.T) { + dir := t.TempDir() + a := openTestStore(t, dir, "a") + b := openTestStore(t, dir, "b") + ctx := context.Background() + limits := []Limit{{Name: "u", Active: 1, Queued: 2, PerIP: 2}} + require.NoError(t, a.Activate("a", "", limits)) + one, err := a.Acquire(ctx, "a", "u", "ip") + require.NoError(t, err) + require.True(t, one.Active) + two, err := b.Acquire(ctx, "b", "u", "ip") + require.NoError(t, err) + require.False(t, two.Active) + _, err = b.Acquire(ctx, "b", "u", "ip") + require.ErrorIs(t, err, ErrPerIP) + three, err := a.Acquire(ctx, "a", "u", "other") + require.NoError(t, err) + require.False(t, three.Active) + _, err = b.Acquire(ctx, "b", "u", "third") + require.ErrorIs(t, err, ErrFull) + require.NoError(t, b.Activate("b", "a", limits)) + require.ErrorIs(t, a.SetLimits("a", limits), ErrNotCurrent) + require.NoError(t, a.Release(ctx, "a", one.ID)) + bt, err := b.Poll("b") + require.NoError(t, err) + require.True(t, bt[two.ID].Active) + at, err := a.Poll("a") + require.NoError(t, err) + require.False(t, at[three.ID].Active) + require.NoError(t, b.Finish("b")) + at, err = a.Poll("a") + require.NoError(t, err) + require.True(t, at[three.ID].Active) + require.NoError(t, a.Release(ctx, "a", three.ID)) + require.NoError(t, a.Release(ctx, "a", three.ID)) +} + +func TestLowerLimitAndUnlimitedQueue(t *testing.T) { + s := openTestStore(t, t.TempDir(), "a") + ctx := context.Background() + require.NoError(t, s.Activate("a", "", []Limit{{Name: "u", Active: 2}})) + one, err := s.Acquire(ctx, "a", "u", "ip") + require.NoError(t, err) + two, err := s.Acquire(ctx, "a", "u", "ip") + require.NoError(t, err) + require.NoError(t, s.SetLimits("a", []Limit{{Name: "u", Active: 1}})) + for range 50 { + ticket, err := s.Acquire(ctx, "a", "u", "ip") + require.NoError(t, err) + require.False(t, ticket.Active) + } + require.NoError(t, s.Release(ctx, "a", one.ID)) + q, err := s.Queues() + require.NoError(t, err) + require.Equal(t, 1, q[0].ActiveCount) + require.Equal(t, 50, q[0].QueuedCount) + require.NoError(t, s.Release(ctx, "a", two.ID)) + q, err = s.Queues() + require.NoError(t, err) + require.Equal(t, 1, q[0].ActiveCount) + require.Equal(t, 49, q[0].QueuedCount) +} + +func TestSchemaRejectionAndBusy(t *testing.T) { + dir := t.TempDir() + s := openTestStore(t, dir, "a") + _, err := s.db.Exec("PRAGMA user_version=999") + require.NoError(t, err) + _, err = Open(dir) + require.ErrorContains(t, err, "incompatible state schema") + _, err = s.db.Exec("PRAGMA user_version=1") + require.NoError(t, err) + b := openTestStore(t, dir, "b") + require.NoError(t, s.Activate("a", "", []Limit{{Name: "u", Active: 1}})) + c, err := s.db.Conn(context.Background()) + require.NoError(t, err) + _, err = c.ExecContext(context.Background(), "BEGIN IMMEDIATE") + require.NoError(t, err) + _, err = b.Acquire(context.Background(), "b", "u", "ip") + require.Error(t, err) + _, err = c.ExecContext(context.Background(), "ROLLBACK") + require.NoError(t, err) + require.NoError(t, c.Close()) + q, err := s.Queues() + require.NoError(t, err) + require.Zero(t, q[0].ActiveCount) +} + +func TestStateWorker(t *testing.T) { + dir := os.Getenv("RSYNC_STATE_TEST_DIR") + if dir == "" { + return + } + id := fmt.Sprint(os.Getpid()) + s, err := Open(dir) + if err != nil { + panic(err) + } + defer s.Close() + lock, err := Lock(filepath.Join(dir, id+".lock")) + if err != nil { + panic(err) + } + defer lock.Close() + if err = s.Register(Generation{ID: id, PID: os.Getpid()}); err != nil { + panic(err) + } + ctx := context.Background() + for range 40 { + ticket, err := s.Acquire(ctx, id, "u", id) + if errors.Is(err, ErrFull) || errors.Is(err, ErrPerIP) { + time.Sleep(time.Millisecond) + continue + } + if err != nil { + panic(err) + } + for !ticket.Active { + time.Sleep(time.Millisecond) + ts, err := s.Poll(id) + if err != nil { + panic(err) + } + ticket = ts[ticket.ID] + } + time.Sleep(2 * time.Millisecond) + if err = s.Release(ctx, id, ticket.ID); err != nil { + panic(err) + } + } +} + +func TestConcurrentProcessAdmission(t *testing.T) { + dir := t.TempDir() + s := openTestStore(t, dir, "parent") + require.NoError(t, s.Activate("parent", "", []Limit{{Name: "u", Active: 2, Queued: 20, PerIP: 2}})) + executable, err := os.Executable() + require.NoError(t, err) + commands := make([]*exec.Cmd, 0, 4) + for range 4 { + c := exec.Command(executable, "-test.run=^TestStateWorker$") + c.Env = append(os.Environ(), "RSYNC_STATE_TEST_DIR="+dir) + require.NoError(t, c.Start()) + commands = append(commands, c) + } + var wg sync.WaitGroup + errs := make(chan error, len(commands)) + for _, c := range commands { + wg.Go(func() { errs <- c.Wait() }) + } + done := make(chan struct{}) + go func() { wg.Wait(); close(done) }() + ticker := time.NewTicker(time.Millisecond) + defer ticker.Stop() + for { + select { + case <-done: + for range commands { + require.NoError(t, <-errs) + } + return + case <-ticker.C: + qs, err := s.Queues() + require.NoError(t, err) + require.LessOrEqual(t, qs[0].ActiveCount, 2) + } + } +} + +func TestAdmissionAt100ConnectionsPerSecond(t *testing.T) { + s := openTestStore(t, t.TempDir(), "a") + require.NoError(t, s.Activate("a", "", []Limit{{Name: "u", Active: 10, Queued: 100}})) + latencies := make([]time.Duration, 0, 200) + ticker := time.NewTicker(10 * time.Millisecond) + defer ticker.Stop() + for range 200 { + <-ticker.C + start := time.Now() + ticket, err := s.Acquire(context.Background(), "a", "u", "ip") + require.NoError(t, err) + latencies = append(latencies, time.Since(start)) + require.NoError(t, s.Release(context.Background(), "a", ticket.ID)) + } + sort.Slice(latencies, func(i, j int) bool { return latencies[i] < latencies[j] }) + t.Logf("100 connections/s, admission p50=%s p99=%s max=%s", latencies[100], latencies[198], latencies[199]) +} + +func TestQueueBurst(t *testing.T) { + dir := t.TempDir() + s := openTestStore(t, dir, "a") + _ = openTestStore(t, dir, "b") + ctx := context.Background() + require.NoError(t, s.Activate("a", "", []Limit{{Name: "u", Active: 1, Queued: 1000}})) + active, err := s.Acquire(ctx, "a", "u", "ip") + require.NoError(t, err) + var first Ticket + for i := range 1000 { + id := "a" + if i%2 == 0 { + id = "b" + } + ticket, err := s.Acquire(ctx, id, "u", "ip") + require.NoError(t, err) + require.False(t, ticket.Active) + if i == 0 { + first = ticket + } + } + start := time.Now() + tickets, err := s.Poll("b") + require.NoError(t, err) + t.Logf("1000 queued connections: batched poll took %s", time.Since(start)) + require.Len(t, tickets, 500) + require.Zero(t, tickets[first.ID].Index) + require.Equal(t, 1000, tickets[first.ID].Queued) + require.NoError(t, s.Release(ctx, "a", active.ID)) + tickets, err = s.Poll("b") + require.NoError(t, err) + require.True(t, tickets[first.ID].Active) +} diff --git a/test/systemd/upgrade.py b/test/systemd/upgrade.py new file mode 100644 index 0000000..ff6d099 --- /dev/null +++ b/test/systemd/upgrade.py @@ -0,0 +1,157 @@ +#!/usr/bin/env python3 +"""Exercise MAINPID handoff using an isolated user-level systemd service.""" + +import argparse +import contextlib +import http.client +import json +import os +from pathlib import Path +import shutil +import socket +import subprocess +import tempfile +import threading +import time + + +class UnixHTTP(http.client.HTTPConnection): + def __init__(self, path): + super().__init__("localhost", timeout=40) + self.path = str(path) + + def connect(self): + self.sock = socket.socket(socket.AF_UNIX) + self.sock.settimeout(self.timeout) + self.sock.connect(self.path) + + +def control(path, method, endpoint): + with contextlib.closing(UnixHTTP(path)) as client: + client.request(method, endpoint) + response = client.getresponse() + body = response.read() + if response.status != 200: + raise RuntimeError(f"{endpoint}: {response.status}: {body!r}") + return body + + +def echo(connection): + with connection, connection.makefile("rb") as reader: + if not reader.readline(): + return + connection.sendall(b"@RSYNCD: 32.0 sha512 sha256 sha1 md5 md4\n") + if not reader.readline(): + return + connection.sendall(b"READY\n") + while line := reader.readline(): + connection.sendall(line) + + +def accept(listener): + while True: + try: + connection, _ = listener.accept() + except OSError: + return + threading.Thread(target=echo, args=(connection,), daemon=True).start() + + +def property_value(unit, name): + return subprocess.check_output( + ["systemctl", "--user", "show", unit, f"--property={name}", "--value"], + text=True, + ).strip() + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--binary", required=True, type=Path) + args = parser.parse_args() + source = args.binary.resolve(strict=True) + unit = f"rsync-proxy-upgrade-test-{os.getpid()}.service" + with tempfile.TemporaryDirectory(prefix="rp-systemd-") as temporary: + directory = Path(temporary) + binary = directory / "proxy" + shutil.copy2(source, binary) + control_socket = directory / "http.sock" + # A Unix rsync listener also verifies that no retired process unlinks it. + rsync_socket = directory / "rsync.sock" + config = directory / "config.toml" + listener = socket.socket() + listener.bind(("127.0.0.1", 0)) + listener.listen() + port = listener.getsockname()[1] + threading.Thread(target=accept, args=(listener,), daemon=True).start() + config.write_text( + f'[proxy]\nlisten="{rsync_socket}"\nlisten_http="{control_socket}"\n' + f'state_dir="{directory}"\n[upstreams.u]\n' + f'address="127.0.0.1:{port}"\nmodules=["foo"]\n' + 'max_active_connections=3\n' + ) + connections = [] + readers = [] + try: + subprocess.run( + ["systemd-run", "--user", f"--unit={unit}", "--collect", + "--property=Type=notify", "--property=NotifyAccess=all", + "--property=KillMode=control-group", "--property=TimeoutStopSec=5s", + str(binary), "--config", str(config)], + check=True, + ) + pids = [] + for generation in range(3): + if generation: + shutil.copy2(source, directory / "proxy.next") + os.replace(directory / "proxy.next", binary) + control(control_socket, "POST", "/upgrade?timeout=30s") + assert property_value(unit, "ActiveState") == "active" + pid = property_value(unit, "MainPID") + assert pid != "0" and pid not in pids, pids + pids.append(pid) + connection = socket.socket(socket.AF_UNIX) + connection.settimeout(5) + connection.connect(str(rsync_socket)) + connections.append(connection) + reader = connection.makefile("rb") + readers.append(reader) + connection.sendall(b"@RSYNCD: 32.0\n") + assert reader.readline().startswith(b"@RSYNCD:") + connection.sendall(b"foo\n") + assert reader.readline() == b"READY\n" + for client, stream in zip(connections, readers): + client.sendall(b"ping\n") + assert stream.readline() == b"ping\n" + deadline = time.monotonic() + 5 + while time.monotonic() < deadline: + status = json.loads(control(control_socket, "GET", "/status")) + if status["count"] == 3: + break + time.sleep(0.1) + assert status["count"] == 3, status + for reader in readers: + reader.close() + for connection in connections: + connection.close() + deadline = time.monotonic() + 5 + while time.monotonic() < deadline: + status = json.loads(control(control_socket, "GET", "/status")) + if all(g["state"] == "exited" for g in status["generations"][:-1]): + break + time.sleep(0.1) + assert all(g["state"] == "exited" for g in status["generations"][:-1]), status + assert property_value(unit, "MainPID") == pids[-1] + assert property_value(unit, "ActiveState") == "active" + assert rsync_socket.exists() and control_socket.exists() + print(f"PASS: MAINPID {' -> '.join(pids)}; all streams survived; old generations exited") + finally: + for reader in readers: + reader.close() + for connection in connections: + connection.close() + subprocess.run(["systemctl", "--user", "stop", unit], check=False) + listener.close() + + +if __name__ == "__main__": + main() From 92afd5f4e7eaf98d8769ec63d9cf619d87fff64c Mon Sep 17 00:00:00 2001 From: taoky Date: Thu, 10 Sep 2026 03:35:03 +0800 Subject: [PATCH 2/3] test: report state worker failures with captured output Include worker PIDs, operation context, and subprocess output when concurrent admission tests fail. Assisted-by: OpenAI Codex (GPT-6) --- pkg/state/store_test.go | 45 +++++++++++++++++++++++++++++------------ 1 file changed, 32 insertions(+), 13 deletions(-) diff --git a/pkg/state/store_test.go b/pkg/state/store_test.go index ef3a86a..5061929 100644 --- a/pkg/state/store_test.go +++ b/pkg/state/store_test.go @@ -1,6 +1,7 @@ package state import ( + "bytes" "context" "errors" "fmt" @@ -120,16 +121,16 @@ func TestStateWorker(t *testing.T) { id := fmt.Sprint(os.Getpid()) s, err := Open(dir) if err != nil { - panic(err) + panic(fmt.Errorf("open state store: %w", err)) } defer s.Close() lock, err := Lock(filepath.Join(dir, id+".lock")) if err != nil { - panic(err) + panic(fmt.Errorf("lock generation: %w", err)) } defer lock.Close() if err = s.Register(Generation{ID: id, PID: os.Getpid()}); err != nil { - panic(err) + panic(fmt.Errorf("register generation: %w", err)) } ctx := context.Background() for range 40 { @@ -139,19 +140,19 @@ func TestStateWorker(t *testing.T) { continue } if err != nil { - panic(err) + panic(fmt.Errorf("acquire ticket: %w", err)) } for !ticket.Active { time.Sleep(time.Millisecond) ts, err := s.Poll(id) if err != nil { - panic(err) + panic(fmt.Errorf("poll tickets: %w", err)) } ticket = ts[ticket.ID] } time.Sleep(2 * time.Millisecond) if err = s.Release(ctx, id, ticket.ID); err != nil { - panic(err) + panic(fmt.Errorf("release ticket %d: %w", ticket.ID, err)) } } } @@ -162,17 +163,32 @@ func TestConcurrentProcessAdmission(t *testing.T) { require.NoError(t, s.Activate("parent", "", []Limit{{Name: "u", Active: 2, Queued: 20, PerIP: 2}})) executable, err := os.Executable() require.NoError(t, err) - commands := make([]*exec.Cmd, 0, 4) + type worker struct { + cmd *exec.Cmd + output *bytes.Buffer + } + workers := make([]worker, 0, 4) for range 4 { c := exec.Command(executable, "-test.run=^TestStateWorker$") c.Env = append(os.Environ(), "RSYNC_STATE_TEST_DIR="+dir) + output := new(bytes.Buffer) + c.Stdout = output + c.Stderr = output require.NoError(t, c.Start()) - commands = append(commands, c) + workers = append(workers, worker{cmd: c, output: output}) } var wg sync.WaitGroup - errs := make(chan error, len(commands)) - for _, c := range commands { - wg.Go(func() { errs <- c.Wait() }) + type workerResult struct { + pid int + err error + output string + } + results := make(chan workerResult, len(workers)) + for _, w := range workers { + wg.Go(func() { + err := w.cmd.Wait() + results <- workerResult{pid: w.cmd.Process.Pid, err: err, output: w.output.String()} + }) } done := make(chan struct{}) go func() { wg.Wait(); close(done) }() @@ -181,8 +197,11 @@ func TestConcurrentProcessAdmission(t *testing.T) { for { select { case <-done: - for range commands { - require.NoError(t, <-errs) + for range workers { + result := <-results + if result.err != nil { + t.Errorf("worker PID %d failed: %v\n%s", result.pid, result.err, result.output) + } } return case <-ticker.C: From 37ca8af209cd9f3eff94f1f3e503bdcfb261ebdf Mon Sep 17 00:00:00 2001 From: taoky Date: Thu, 10 Sep 2026 05:27:56 +0800 Subject: [PATCH 3/3] refactor: coordinate worker upgrades through a supervisor Replace SQLite with in-memory admission state and Unix socket IPC. Keep the supervisor PID and listeners stable while workers upgrade and drain, and reclaim capacity after worker exit. Cover supervisor failure, worker restart, and graceful shutdown; update systemd integration and migration documentation. Assisted-by: OpenAI Codex (GPT-6) --- README.md | 14 +- assets/config.example.toml | 2 +- assets/rsync-proxy.service | 4 +- cmd/cmd.go | 4 - cmd/daemon_linux.go | 269 +++--------------- cmd/daemon_linux_test.go | 165 ++++++++++- cmd/daemon_other.go | 8 +- cmd/supervisor_linux.go | 358 +++++++++++++++++++++++ go.mod | 9 - go.sum | 48 ---- pkg/server/admission.go | 65 +++-- pkg/server/snapshot_test.go | 3 +- pkg/state/ipc.go | 272 ++++++++++++++++++ pkg/state/store.go | 552 +++++++++++++++++++++--------------- pkg/state/store_test.go | 82 ++++-- test/systemd/upgrade.py | 31 +- 16 files changed, 1297 insertions(+), 589 deletions(-) create mode 100644 cmd/supervisor_linux.go create mode 100644 pkg/state/ipc.go diff --git a/README.md b/README.md index ee52a0e..14eca9d 100644 --- a/README.md +++ b/README.md @@ -67,17 +67,19 @@ sudo /usr/bin/rsync-proxy --config=/etc/rsync-proxy/config.toml upgrade --timeou `upgrade` 等待新进程就绪并完成监听交接后返回，不等待旧连接结束。已有明文/TLS 传输、握手和排队连接留在原进程；新连接由新进程接收。可以在旧连接未结束时继续升级，形成多代共存。`reload` 仍然只重读配置；`systemctl restart` 是普通重启，不是零停机更新入口。 -- `[proxy].state_dir` 默认为 `/run/rsync-proxy`，保存各代共享的 SQLite 运行状态。手动运行时须使用可写目录，各实例须使用独立目录；不要放在 NFS 等网络文件系统上，也不要在进程运行期间删除数据库、WAL 或锁文件。 +- `[proxy].state_dir` 默认为 `/run/rsync-proxy`，保存 supervisor 的 Unix socket 和进程锁文件。全局限额、FIFO 队列和各代快照保存在 supervisor 内存中，不使用数据库。手动运行时须使用可写目录，各实例须使用独立目录；不要在进程运行期间删除 socket 或锁文件。 - `[proxy].upgrade_drain_timeout` 单位为秒，默认 `0`。非零值从该代停止接收新连接时起计算，到期关闭剩余连接，包括握手和排队。采用本次成功升级的新配置值；后续更新不会改变或重置更早各代的期限。默认不增加期限，原有 idle/max-duration/throughput 规则继续生效。 - 所有代合计执行上游并发、排队和单 IP 限额，排队按全局 FIFO 晋升。调低限额不会主动断开已有连接；等待占用降到新限额以下再放行。旧连接继续使用已经选定的上游目标。 -- 热更新要求 `listen`、`listen_tls`、`listen_http` 和状态目录保持一致。新程序、配置、证书或协议不兼容，或就绪超过 `--timeout`，会报告失败并保留旧进程。数据库不可用时不绕过配额放行，已建立的传输继续运行。 -- `connections`、`/status` 和 `/metrics` 汇总各代，逐连接数据新增 `generation` 标识；Go/runtime 指标属于当前主进程。跨代快照每秒发布，正常退出前发布最终值，业务累计指标保留已退出代的贡献；异常退出可能损失最后一次未发布的统计。一次完全停止后的启动会重置运行状态。 -- 使用配套 systemd unit，使主 PID 在升级后指向新进程；需要支持 `sd_notify` barrier 的 systemd（246 或更新版本）。普通停止仍受 systemd 的停止超时约束。旧进程仍有连接时会继续占用内存和 FD，可通过退出期限限制。 +- 热更新要求 `listen`、`listen_tls`、`listen_http` 和状态目录保持一致。新程序、配置、证书或协议不兼容，或就绪超过 `--timeout`，会报告失败并保留旧进程。准入操作不做磁盘写入，队列变化通过 IPC 通知 worker。 +- `connections`、`/status` 和 `/metrics` 汇总各代，逐连接数据新增 `generation` 标识；Go/runtime 指标属于当前接收连接的 worker。跨代快照每秒发布，正常退出前发布最终值，业务累计指标保留已退出代的贡献；异常退出可能损失最后一次未发布的统计。一次完全停止后的启动会重置运行状态。 +- 使用配套 systemd unit，主 PID 始终是 supervisor。supervisor 保留监听 FD，通过 Unix socket 协调 worker；`upgrade` 启动安装路径上的新 worker，并让旧 worker 排空连接。普通停止仍受 systemd 的停止超时约束。旧 worker 仍有连接时会继续占用内存和 FD，可通过退出期限限制。 +- supervisor 确认 worker 退出后回收其额度；当前 worker 意外退出时，会从安装路径启动替代 worker。supervisor 或控制连接异常失效时，worker 停止接收并关闭连接，不会另建空状态继续放行。 +- `upgrade` 更新 worker，不替换常驻 supervisor；修改 supervisor 实现或 IPC 协议后，需要普通重启才能应用这部分更新。不兼容的 worker 会被拒绝，旧 worker 继续服务。 - 日志轮转使用 `rsync-proxy --config=... reopen-logs`，让所有存活的代重新打开日志。升级与日志重开命令读取配置中的管理地址，显式 `--host` 可覆盖。 -**首次从不支持该协议的旧版迁入需要一次普通重启。** deb 升级会对正在运行的服务调用 `upgrade`，失败不会自动强制重启；首次迁入或不兼容升级时，检查错误后手动 `systemctl restart rsync-proxy`，再完成包配置。进程崩溃、主机重启、主动停止和配置的连接超时不属于零停机保证。 +**首次从不支持 supervisor 协议的旧版（包括 SQLite 版本）迁入需要一次普通重启。** deb 升级会对正在运行的服务调用 `upgrade`，失败不会自动强制重启；首次迁入或不兼容升级时，检查错误后手动 `systemctl restart rsync-proxy`，再完成包配置。进程崩溃、主机重启、主动停止和配置的连接超时不属于零停机保证。 -开发验证：`go test -race ./...` 包含真实子进程的多代交接、SQLite 并发调度及故障测试。真实用户级 systemd 验证可运行 `python3 test/systemd/upgrade.py --binary /absolute/path/to/rsync-proxy`；该脚本只创建临时测试 unit、回环监听和临时目录。 +开发验证：`go test -race ./...` 包含真实子进程的多代交接、supervisor IPC 并发调度及故障测试。真实用户级 systemd 验证可运行 `python3 test/systemd/upgrade.py --binary /absolute/path/to/rsync-proxy`；该脚本只创建临时测试 unit、回环监听和临时目录。 ### 安装 fail2ban filter diff --git a/assets/config.example.toml b/assets/config.example.toml index b8a9eb1..6b356f8 100644 --- a/assets/config.example.toml +++ b/assets/config.example.toml @@ -11,7 +11,7 @@ error_log = "/var/log/rsync-proxy/error.log" tls_cert_file = "/etc/rsync-proxy/tls/server.crt" tls_key_file = "/etc/rsync-proxy/tls/server.key" -# Shared runtime state for all process generations (Linux). +# Supervisor Unix socket and process lock directory (Linux). # Use a local directory writable by the service user, never NFS. #state_dir = "/run/rsync-proxy" diff --git a/assets/rsync-proxy.service b/assets/rsync-proxy.service index 8a106c8..e1c47e5 100644 --- a/assets/rsync-proxy.service +++ b/assets/rsync-proxy.service @@ -6,8 +6,8 @@ StartLimitIntervalSec=10s [Service] Type=notify -NotifyAccess=all -KillMode=control-group +NotifyAccess=main +KillMode=mixed Restart=on-failure ExecStart=/usr/bin/rsync-proxy --config=/etc/rsync-proxy/config.toml diff --git a/cmd/cmd.go b/cmd/cmd.go index 46825da..1fa5146 100644 --- a/cmd/cmd.go +++ b/cmd/cmd.go @@ -356,10 +356,6 @@ func New() *cobra.Command { } log.SetOutput(cmd.ErrOrStderr()) - if err := s.ReadConfigFromFile(true); err != nil { - return fmt.Errorf("load config: %w", err) - } - return runDaemon(s) }, SilenceUsage: true, diff --git a/cmd/daemon_linux.go b/cmd/daemon_linux.go index 6426320..6a1a486 100644 --- a/cmd/daemon_linux.go +++ b/cmd/daemon_linux.go @@ -14,18 +14,14 @@ import ( "net" "net/http" "os" - "os/exec" "os/signal" "path/filepath" "strconv" - "strings" "sync" "sync/atomic" "syscall" "time" - "golang.org/x/sys/unix" - "github.com/ustclug/rsync-proxy/pkg/server" "github.com/ustclug/rsync-proxy/pkg/state" ) @@ -46,7 +42,7 @@ type daemon struct { ready atomic.Bool s *server.Server store *state.Store - id, dir, binary string + id, dir string mu sync.Mutex // serializes reload and upgrade through the commit point private *http.Server privateListener net.Listener @@ -55,11 +51,12 @@ type daemon struct { done chan struct{} backgroundDone chan struct{} drainOnce sync.Once + drainMu sync.Mutex drainTimer *time.Timer stopAdmission func() } -func runDaemon(s *server.Server) error { +func runWorker(s *server.Server) error { dir := s.StateDir if dir == "" { dir = "/run/rsync-proxy" @@ -72,25 +69,16 @@ func runDaemon(s *server.Server) error { if err = os.MkdirAll(dir, 0700); err != nil { return err } - startup, err := state.Lock(filepath.Join(dir, "startup.lock")) - if err != nil { - return fmt.Errorf("another startup is in progress: %w", err) - } - defer startup.Close() store, err := state.Open(dir) if err != nil { return err } defer store.Close() - binary, err := os.Executable() - if err != nil { - return err - } idBytes := make([]byte, 12) if _, err = rand.Read(idBytes); err != nil { return err } - d := &daemon{s: s, store: store, dir: dir, id: hex.EncodeToString(idBytes), binary: binary, runResult: make(chan error, 1), done: make(chan struct{}), backgroundDone: make(chan struct{})} + d := &daemon{s: s, store: store, dir: dir, id: hex.EncodeToString(idBytes), runResult: make(chan error, 1), done: make(chan struct{}), backgroundDone: make(chan struct{})} var parent net.Conn var decoder *json.Decoder var hello handoffMessage @@ -113,27 +101,9 @@ func runDaemon(s *server.Server) error { return errors.New("incompatible handoff protocol or state_dir") } } else { - gens, err := store.Generations() - if err != nil { - return err - } - for _, g := range gens { - live, err := state.IsLocked(filepath.Join(dir, g.ID+".lock")) - if err != nil { - return err - } - if live { - return fmt.Errorf("generation %s is still running; use upgrade", g.ID) - } - } - if err = store.Reset(); err != nil { - return err - } - for _, g := range gens { - _ = os.Remove(filepath.Join(dir, g.ID+".lock")) - _ = os.Remove(g.Control) - } + return errors.New("worker requires supervisor handoff") } + d.lock, err = state.Lock(filepath.Join(dir, d.id+".lock")) if err != nil { return err @@ -151,11 +121,7 @@ func runDaemon(s *server.Server) error { if err = store.Register(state.Generation{ID: d.id, PID: os.Getpid(), Version: Version, Control: control}); err != nil { return err } - defer func() { - if err := store.Finish(d.id); err != nil { - log.Printf("[ERROR] retire generation: %v", err) - } - }() + s.Generation = d.id // Configuration was loaded before attaching the store: a prospective child // must not publish new limits until the parent authorizes the commit. @@ -190,7 +156,11 @@ func runDaemon(s *server.Server) error { return } if r.URL.Path == "/internal/stop" { - d.beginDrain(0) + timeout, err := time.ParseDuration(r.URL.Query().Get("timeout")) + if err != nil { + timeout = 0 + } + d.beginDrain(timeout) return } if !d.control(w, r) { @@ -237,25 +207,21 @@ func runDaemon(s *server.Server) error { if err = s.PublishSnapshot(); err != nil { return err } - if parent == nil { - if err = notifySystemd("READY=1\nMAINPID=" + strconv.Itoa(os.Getpid())); err != nil { - return err - } - } else { - if err = json.NewEncoder(parent).Encode(handoffMessage{Stage: "active", ID: d.id}); err != nil { - return err - } - // Keep the parent alive until it has transferred systemd ownership. - var ack handoffMessage - if err = decoder.Decode(&ack); err != nil { - return err - } - if ack.Stage != "owned" { - return errors.New("missing ownership acknowledgement") - } + if err = json.NewEncoder(parent).Encode(handoffMessage{Stage: "active", ID: d.id}); err != nil { + return err + } + var ack handoffMessage + if err = decoder.Decode(&ack); err != nil { + return err + } + if ack.Stage != "owned" { + return errors.New("missing supervisor acknowledgement") } - _ = startup.Close() + d.ready.Store(true) + if err = json.NewEncoder(parent).Encode(handoffMessage{Stage: "serving", ID: d.id}); err != nil { + return err + } go d.background() defer func() { close(d.done); <-d.backgroundDone }() signals := make(chan os.Signal, 2) @@ -264,20 +230,30 @@ func runDaemon(s *server.Server) error { var runErr error select { case runErr = <-d.runResult: + case <-store.Done(): + s.Close() + <-d.runResult + runErr = errors.New("supervisor connection lost") case <-signals: - d.stopAll() + d.beginDrain(0) select { case runErr = <-d.runResult: + case <-store.Done(): + s.Close() + <-d.runResult + runErr = errors.New("supervisor connection lost while draining") case <-signals: s.AbortConnections() runErr = <-d.runResult } } runFinished = true + d.drainMu.Lock() if d.drainTimer != nil { d.drainTimer.Stop() } - // Finish HTTP handlers before closing their database/log dependencies. + d.drainMu.Unlock() + // Finish HTTP handlers before closing their supervisor/log dependencies. ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() if s.HTTPServer != nil { @@ -286,19 +262,14 @@ func runDaemon(s *server.Server) error { if err = s.PublishSnapshot(); err != nil { log.Printf("[ERROR] final snapshot: %v", err) } - current, _ := store.Current() - if current == d.id { - for _, addr := range s.ConfiguredListeners() { - if strings.HasPrefix(addr, "/") { - _ = os.Remove(addr) - } - } - } + return runErr } func (d *daemon) beginDrain(timeout time.Duration) { d.drainOnce.Do(func() { + d.drainMu.Lock() + defer d.drainMu.Unlock() var deadline time.Time if timeout > 0 { deadline = time.Now().Add(timeout) @@ -324,40 +295,8 @@ func (d *daemon) background() { if err := d.s.PublishSnapshot(); err != nil { log.Printf("[ERROR] publish snapshot: %v", err) } - gens, err := d.store.Generations() - if err != nil { - continue - } - for _, g := range gens { - if g.ID == d.id || g.State == state.Exited { - continue - } - live, err := state.IsLocked(filepath.Join(d.dir, g.ID+".lock")) - if err == nil && !live { - if err := d.store.Finish(g.ID); err != nil { - log.Printf("[ERROR] reclaim generation: %v", err) - } - } - } - } - } -} - -func (d *daemon) stopAll() { - current, err := d.store.Current() - if err == nil && current == d.id { - gens, err := d.store.Generations() - if err == nil { - for _, g := range gens { - if g.ID != d.id && g.State != state.Exited { - if err := privatePost(g.Control, "/internal/stop"); err != nil { - log.Printf("[ERROR] stop generation %s: %v", g.ID, err) - } - } - } } } - d.beginDrain(0) } func privatePost(addr, path string) error { @@ -456,7 +395,7 @@ func (d *daemon) control(w http.ResponseWriter, r *http.Request) bool { return true } } - if err := d.upgrade(timeout); err != nil { + if err := d.store.Upgrade(d.id, timeout); err != nil { http.Error(w, err.Error(), http.StatusInternalServerError) } else { _, _ = fmt.Fprintln(w, `{"message":"Upgrade completed"}`) @@ -464,129 +403,3 @@ func (d *daemon) control(w http.ResponseWriter, r *http.Request) bool { } return true } - -func (d *daemon) upgrade(timeout time.Duration) (resultErr error) { - files, err := d.s.ListenerFiles() - if err != nil { - return err - } - defer func() { - for _, f := range files { - _ = f.Close() - } - }() - pair, err := unix.Socketpair(unix.AF_UNIX, unix.SOCK_STREAM|unix.SOCK_CLOEXEC, 0) - if err != nil { - return err - } - pf := os.NewFile(uintptr(pair[0]), "parent") - childFile := os.NewFile(uintptr(pair[1]), "child") - defer childFile.Close() - conn, err := net.FileConn(pf) - _ = pf.Close() - if err != nil { - return err - } - defer conn.Close() - _ = conn.SetDeadline(time.Now().Add(timeout)) - env := make([]string, 0, len(os.Environ())+1) - for _, v := range os.Environ() { - if !strings.HasPrefix(v, handoffEnv+"=") { - env = append(env, v) - } - } - env = append(env, handoffEnv+"="+strconv.Itoa(state.Protocol)) - child := &exec.Cmd{Path: d.binary, Args: append([]string{d.binary}, os.Args[1:]...), Env: env, ExtraFiles: append([]*os.File{childFile}, files...), Stdin: os.Stdin, Stdout: os.Stdout, Stderr: os.Stderr} - if err = child.Start(); err != nil { - return err - } - _ = childFile.Close() - wait := make(chan error, 1) - go func() { wait <- child.Wait() }() - success := false - defer func() { - if success { - return - } - _ = child.Process.Kill() - <-wait - // The old listeners have not been closed. Restore admission policy if - // the child failed after committing but before acknowledging ownership. - current, e := d.store.Current() - if e == nil && current != d.id { - if e = d.store.Activate(d.id, current, d.s.Limits()); e != nil { - log.Printf("[ERROR] restore serving generation: %v", e) - } - } - _ = notifySystemd("MAINPID=" + strconv.Itoa(os.Getpid())) - }() - enc := json.NewEncoder(conn) - dec := json.NewDecoder(conn) - if err = enc.Encode(handoffMessage{Protocol: state.Protocol, Parent: d.id, StateDir: d.dir, Listeners: d.s.ConfiguredListeners()}); err != nil { - return err - } - var ready handoffMessage - if err = dec.Decode(&ready); err != nil { - return fmt.Errorf("new generation did not become ready: %w", err) - } - if ready.Stage != "ready" { - return errors.New("invalid readiness acknowledgement") - } - if err = enc.Encode(handoffMessage{Stage: "commit"}); err != nil { - return err - } - var active handoffMessage - if err = dec.Decode(&active); err != nil { - return fmt.Errorf("new generation failed to activate: %w", err) - } - if active.Stage != "active" || active.ID != ready.ID { - return errors.New("invalid activation acknowledgement") - } - if err = notifySystemd("MAINPID=" + strconv.Itoa(child.Process.Pid)); err != nil { - return fmt.Errorf("transfer systemd ownership: %w", err) - } - if err = enc.Encode(handoffMessage{Stage: "owned"}); err != nil { - return err - } - success = true - d.beginDrain(ready.DrainTimeout) - log.Printf("[INFO] upgraded %s to %s (pid %d)", d.id, ready.ID, child.Process.Pid) - return nil -} - -// A barrier makes ownership changes visible to PID 1 before the old main -// process is allowed to exit. Only the current main process sends MAINPID. -func notifySystemd(message string) error { - path := os.Getenv("NOTIFY_SOCKET") - if path == "" { - return nil - } - fd, err := unix.Socket(unix.AF_UNIX, unix.SOCK_DGRAM|unix.SOCK_CLOEXEC, 0) - if err != nil { - return err - } - defer unix.Close(fd) - addr := &unix.SockaddrUnix{Name: path} - if err = unix.Sendto(fd, []byte(message), 0, addr); err != nil { - return err - } - r, w, err := os.Pipe() - if err != nil { - return err - } - defer r.Close() - defer w.Close() - if err = unix.Sendmsg(fd, []byte("BARRIER=1"), unix.UnixRights(int(w.Fd())), addr, 0); err != nil { - return err - } - _ = w.Close() - if err = r.SetReadDeadline(time.Now().Add(3 * time.Second)); err != nil { - return err - } - var b [1]byte - _, err = r.Read(b[:]) - if errors.Is(err, io.EOF) { - return nil - } - return err -} diff --git a/cmd/daemon_linux_test.go b/cmd/daemon_linux_test.go index c03cf89..b729b72 100644 --- a/cmd/daemon_linux_test.go +++ b/cmd/daemon_linux_test.go @@ -5,6 +5,7 @@ package cmd import ( "bufio" "bytes" + "context" "crypto/ecdsa" "crypto/elliptic" "crypto/rand" @@ -13,6 +14,7 @@ import ( "crypto/x509/pkix" "encoding/json" "encoding/pem" + "errors" "fmt" "io" "math/big" @@ -21,6 +23,7 @@ import ( "os" "os/exec" "path/filepath" + "strconv" "strings" "sync" "syscall" @@ -432,7 +435,17 @@ func TestDeadGenerationReleasesCapacity(t *testing.T) { require.Contains(t, line, "queued") _, err = br.ReadString('\n') require.NoError(t, err) - require.NoError(t, syscall.Kill(f.pid, syscall.SIGKILL)) + st, err := f.status() + require.NoError(t, err) + var oldPID int + for _, g := range st.Generations { + if g.Version == "A" { + oldPID = g.PID + } + } + require.NotZero(t, oldPID) + require.NotEqual(t, f.pid, oldPID) + require.NoError(t, syscall.Kill(oldPID, syscall.SIGKILL)) checkEcho(t, b, br) } @@ -578,3 +591,153 @@ func TestRealRsyncTransferAcrossUpgrades(t *testing.T) { require.NoError(t, err) require.Equal(t, data, actual) } + +func TestSupervisorRestartsCurrentWorker(t *testing.T) { + f := newDaemonFixture(t, 1, 0) + a, ar := f.connect(false) + checkEcho(t, a, ar) + before, err := f.status() + require.NoError(t, err) + require.Len(t, before.Generations, 1) + old := before.Generations[0] + require.NotEqual(t, f.pid, old.PID) + require.NoError(t, syscall.Kill(old.PID, syscall.SIGKILL)) + require.Eventually(t, func() bool { + st, e := f.status() + if e != nil { + return false + } + for _, g := range st.Generations { + if g.ID != old.ID && g.State == state.Serving { + return true + } + } + return false + }, 5*time.Second, 20*time.Millisecond) + b, br := f.connect(false) + checkEcho(t, b, br) + require.NoError(t, syscall.Kill(f.pid, 0)) +} + +func TestSupervisorDeathStopsWorkers(t *testing.T) { + f := newDaemonFixture(t, 1, 0) + a, ar := f.connect(false) + checkEcho(t, a, ar) + f.install("B") + f.upgrade() + b, br := f.connect(false) + line, err := br.ReadString('\n') + require.NoError(t, err) + require.Contains(t, line, "queued") + _, err = br.ReadString('\n') + require.NoError(t, err) + st, err := f.status() + require.NoError(t, err) + require.NoError(t, syscall.Kill(f.pid, syscall.SIGKILL)) + for _, c := range []net.Conn{a, b} { + require.NoError(t, c.SetReadDeadline(time.Now().Add(5*time.Second))) + _, err = c.Read(make([]byte, 1)) + require.Error(t, err) + var ne net.Error + if errors.As(err, &ne) { + require.False(t, ne.Timeout(), "worker kept connection alive after supervisor died") + } + } + require.Eventually(t, func() bool { + for _, g := range st.Generations { + locked, e := state.IsLocked(filepath.Join(f.dir, g.ID+".lock")) + if e != nil || locked { + return false + } + } + return true + }, 5*time.Second, 20*time.Millisecond) +} + +func TestSupervisorShutdownDrainsAllWorkers(t *testing.T) { + f := newDaemonFixture(t, 0, 0) + a, ar := f.connect(false) + checkEcho(t, a, ar) + f.install("B") + f.upgrade() + b, br := f.connect(false) + checkEcho(t, b, br) + st, err := f.status() + require.NoError(t, err) + require.NoError(t, syscall.Kill(f.pid, syscall.SIGTERM)) + // A graceful stop must not interrupt either generation's accepted sessions. + checkEcho(t, a, ar) + checkEcho(t, b, br) + require.NoError(t, a.Close()) + require.NoError(t, b.Close()) + require.Eventually(t, func() bool { + for _, g := range st.Generations { + live, e := state.IsLocked(filepath.Join(f.dir, g.ID+".lock")) + if e != nil || live { + return false + } + } + live, e := state.IsLocked(filepath.Join(f.dir, "startup.lock")) + return e == nil && !live + }, 5*time.Second, 20*time.Millisecond) +} + +func TestSupervisorDeathDuringWorkerStartup(t *testing.T) { + f := newDaemonFixture(t, 0, 0) + pidFile := filepath.Join(f.dir, "starting.pid") + // This executable never connects to the supervisor or completes handoff. + script := fmt.Sprintf("#!/bin/sh\necho $$ > %q\nexec sleep 30\n", pidFile) + require.NoError(t, os.WriteFile(f.binary+".new", []byte(script), 0700)) + require.NoError(t, os.Rename(f.binary+".new", f.binary)) + finished := make(chan struct{}) + go func() { + defer close(finished) + _ = sendControl(f.socket, "/upgrade?timeout=10s", 10*time.Second, io.Discard) + }() + var pid int + require.Eventually(t, func() bool { + b, e := os.ReadFile(pidFile) + if e != nil { + return false + } + pid, e = strconv.Atoi(strings.TrimSpace(string(b))) + return e == nil && pid > 0 + }, 3*time.Second, 10*time.Millisecond) + t.Cleanup(func() { _ = syscall.Kill(pid, syscall.SIGKILL) }) + require.NoError(t, syscall.Kill(f.pid, syscall.SIGKILL)) + require.Eventually(t, func() bool { + b, e := os.ReadFile(fmt.Sprintf("/proc/%d/stat", pid)) + if os.IsNotExist(e) { + return true + } + if e != nil { + return false + } + // A killed orphan can briefly remain a zombie until init reaps it. + fields := strings.Fields(string(b)) + return len(fields) > 2 && fields[2] == "Z" + }, 3*time.Second, 10*time.Millisecond) + select { + case <-finished: + case <-time.After(3 * time.Second): + t.Fatal("upgrade request did not terminate") + } +} + +func TestRetireDoesNotReclaimReusedPID(t *testing.T) { + store := state.New() + require.NoError(t, store.Register(state.Generation{ID: "old", PID: 42})) + require.NoError(t, store.Finish("old")) + require.NoError(t, store.Register(state.Generation{ID: "new", PID: 42})) + require.NoError(t, store.Activate("new", "", []state.Limit{{Name: "u", Active: 1}})) + _, err := store.Acquire(context.Background(), "new", "u", "ip") + require.NoError(t, err) + p := &supervisor{store: store, dir: t.TempDir()} + p.retire(&workerProcess{cmd: &exec.Cmd{Process: &os.Process{Pid: 42}}, id: "old"}) + q, err := store.Queues() + require.NoError(t, err) + require.Equal(t, 1, q[0].ActiveCount) + gens, err := store.Generations() + require.NoError(t, err) + require.Equal(t, state.Serving, gens[1].State) +} diff --git a/cmd/daemon_other.go b/cmd/daemon_other.go index c449dd2..76db529 100644 --- a/cmd/daemon_other.go +++ b/cmd/daemon_other.go @@ -2,9 +2,15 @@ package cmd -import "github.com/ustclug/rsync-proxy/pkg/server" +import ( + "fmt" + "github.com/ustclug/rsync-proxy/pkg/server" +) func runDaemon(s *server.Server) error { + if err := s.ReadConfigFromFile(true); err != nil { + return fmt.Errorf("load config: %w", err) + } if err := s.Listen(); err != nil { return err } diff --git a/cmd/supervisor_linux.go b/cmd/supervisor_linux.go new file mode 100644 index 0000000..1f3eec3 --- /dev/null +++ b/cmd/supervisor_linux.go @@ -0,0 +1,358 @@ +//go:build linux + +package cmd + +import ( + "encoding/json" + "errors" + "fmt" + "log" + "net" + "os" + "os/exec" + "os/signal" + "path/filepath" + "runtime" + "strconv" + "strings" + "sync" + "syscall" + "time" + + "golang.org/x/sys/unix" + + "github.com/ustclug/rsync-proxy/pkg/server" + "github.com/ustclug/rsync-proxy/pkg/state" +) + +type workerProcess struct { + cmd *exec.Cmd + done chan struct{} + err error // written before done closes + id string // protected by supervisor.mu +} + +type supervisor struct { + mu sync.Mutex + s *server.Server + store *state.Store + dir, binary string + workers map[int]*workerProcess + exited chan *workerProcess + stopping bool +} + +func runDaemon(s *server.Server) error { + worker := os.Getenv(handoffEnv) != "" + // Only workers discover upstream modules and open transfer logs. The + // supervisor needs configuration solely to own the listening sockets. + if err := s.ReadConfigFromFile(worker); err != nil { + return fmt.Errorf("load config: %w", err) + } + if worker { + return runWorker(s) + } + dir := s.StateDir + if dir == "" { + dir = "/run/rsync-proxy" + } + dir, err := filepath.Abs(dir) + if err != nil { + return err + } + s.StateDir = dir + if err = os.MkdirAll(dir, 0700); err != nil { + return err + } + lock, err := state.Lock(filepath.Join(dir, "startup.lock")) + if err != nil { + return fmt.Errorf("instance already running: %w", err) + } + defer lock.Close() + // A previous supervisor may have died before its workers finished exiting. + // Refuse to start with an empty authority while those workers are still live. + locks, err := filepath.Glob(filepath.Join(dir, "*.lock")) + if err != nil { + return err + } + for _, path := range locks { + if path == filepath.Join(dir, "startup.lock") { + continue + } + live, e := state.IsLocked(path) + if e != nil { + return e + } + if live { + return fmt.Errorf("worker still running (%s); wait for it to exit before restarting", path) + } + } + path := filepath.Join(dir, "supervisor.sock") + if err = os.Remove(path); err != nil && !os.IsNotExist(err) { + return err + } + binary, err := os.Executable() + if err != nil { + return err + } + p := &supervisor{s: s, store: state.New(), dir: dir, binary: binary, workers: make(map[int]*workerProcess), exited: make(chan *workerProcess, 64)} + closeIPC, err := state.Serve(dir, p.store, p.upgrade) + if err != nil { + return err + } + defer closeIPC() + defer s.CloseLogs() + if err = s.Listen(); err != nil { + return err + } + defer s.Close() + defer p.killAndWait() + signals := make(chan os.Signal, 2) + signal.Notify(signals, syscall.SIGTERM, syscall.SIGINT) + defer signal.Stop(signals) + p.mu.Lock() + err = p.startWorker("", 30*time.Second) + p.mu.Unlock() + if err != nil { + return err + } + if err = notifySystemd("READY=1\nMAINPID=" + strconv.Itoa(os.Getpid())); err != nil { + return err + } + for { + select { + case <-signals: + p.mu.Lock() + if p.stopping { + for _, w := range p.workers { + _ = w.cmd.Process.Kill() + } + } else { + p.stopping = true + for _, w := range p.workers { + _ = w.cmd.Process.Signal(syscall.SIGTERM) + } + } + empty := len(p.workers) == 0 + p.mu.Unlock() + if empty { + return nil + } + case w := <-p.exited: + p.mu.Lock() + delete(p.workers, w.cmd.Process.Pid) + p.retire(w) + current, _ := p.store.Current() + if w.err != nil { + log.Printf("[ERROR] worker %s exited: %v", w.id, w.err) + } + if !p.stopping && current == w.id { + log.Printf("[INFO] restarting current worker %s", w.id) + err = p.startWorker(current, 30*time.Second) + } + empty := len(p.workers) == 0 + p.mu.Unlock() + if err != nil { + return err + } + if empty { + return nil + } + } + } +} + +func (p *supervisor) retire(w *workerProcess) { + // Registration may have succeeded even if the readiness response was lost. + gens, _ := p.store.Generations() + for _, g := range gens { + if (w.id != "" && g.ID == w.id) || (w.id == "" && g.PID == w.cmd.Process.Pid && g.State != state.Exited) { + w.id = g.ID + _ = p.store.Finish(g.ID) + _ = os.Remove(g.Control) + _ = os.Remove(filepath.Join(p.dir, g.ID+".lock")) + } + } +} + +func (p *supervisor) killAndWait() { + p.mu.Lock() + defer p.mu.Unlock() + p.stopping = true + for _, w := range p.workers { + _ = w.cmd.Process.Kill() + } + for _, w := range p.workers { + <-w.done + p.retire(w) + } +} + +func (p *supervisor) upgrade(id string, timeout time.Duration) error { + if timeout <= 0 { + return errors.New("invalid upgrade timeout") + } + if !p.mu.TryLock() { + return errors.New("supervisor is busy") + } + defer p.mu.Unlock() + if p.stopping { + return errors.New("supervisor is stopping") + } + current, err := p.store.Current() + if err != nil { + return err + } + if id != current { + return state.ErrNotCurrent + } + return p.startWorker(current, timeout) +} + +// startWorker runs with mu held. Listener ownership remains in the supervisor, +// so failure before readiness never removes the old worker's listening sockets. +func (p *supervisor) startWorker(parent string, timeout time.Duration) (result error) { + files, err := p.s.ListenerFiles() + if err != nil { + return err + } + defer func() { + for _, f := range files { + _ = f.Close() + } + }() + pair, err := unix.Socketpair(unix.AF_UNIX, unix.SOCK_STREAM|unix.SOCK_CLOEXEC, 0) + if err != nil { + return err + } + pf := os.NewFile(uintptr(pair[0]), "supervisor-handoff") + childFile := os.NewFile(uintptr(pair[1]), "worker-handoff") + defer childFile.Close() + conn, err := net.FileConn(pf) + _ = pf.Close() + if err != nil { + return err + } + defer conn.Close() + _ = conn.SetDeadline(time.Now().Add(timeout)) + env := make([]string, 0, len(os.Environ())+1) + for _, v := range os.Environ() { + if !strings.HasPrefix(v, handoffEnv+"=") && !strings.HasPrefix(v, "NOTIFY_SOCKET=") { + env = append(env, v) + } + } + env = append(env, handoffEnv+"="+strconv.Itoa(state.Protocol)) + child := &exec.Cmd{Path: p.binary, Args: append([]string{p.binary}, os.Args[1:]...), Env: env, ExtraFiles: append([]*os.File{childFile}, files...), Stdin: os.Stdin, Stdout: os.Stdout, Stderr: os.Stderr, SysProcAttr: &syscall.SysProcAttr{Pdeathsig: syscall.SIGKILL}} + w := &workerProcess{cmd: child, done: make(chan struct{})} + started := make(chan error, 1) + go func() { + // Linux delivers Pdeathsig when the creating thread dies. Keep that + // thread alive until Wait completes, even if Go retires other threads. + // This also covers workers stuck in config loading, before IPC starts. + runtime.LockOSThread() + defer runtime.UnlockOSThread() + e := child.Start() + started <- e + if e != nil { + return + } + w.err = child.Wait() + close(w.done) + p.exited <- w + }() + if err = <-started; err != nil { + return err + } + _ = childFile.Close() + p.workers[child.Process.Pid] = w + + previous, _ := p.store.Queues() + oldLimits := make([]state.Limit, 0, len(previous)) + for _, q := range previous { + oldLimits = append(oldLimits, q.Limit) + } + success := false + defer func() { + if success { + return + } + _ = child.Process.Kill() + <-w.done + p.retire(w) + delete(p.workers, child.Process.Pid) + current, _ := p.store.Current() + if parent != "" && current != parent { + if e := p.store.Activate(parent, current, oldLimits); e != nil { + log.Printf("[ERROR] restore previous worker: %v", e) + } + } + }() + enc, dec := json.NewEncoder(conn), json.NewDecoder(conn) + if err = enc.Encode(handoffMessage{Protocol: state.Protocol, Parent: parent, StateDir: p.dir, Listeners: p.s.ConfiguredListeners()}); err != nil { + return err + } + var ready handoffMessage + if err = dec.Decode(&ready); err != nil { + return fmt.Errorf("worker did not become ready: %w", err) + } + if ready.Stage != "ready" || ready.ID == "" { + return errors.New("invalid worker readiness acknowledgement") + } + w.id = ready.ID + if err = enc.Encode(handoffMessage{Stage: "commit"}); err != nil { + return err + } + var active handoffMessage + if err = dec.Decode(&active); err != nil { + return fmt.Errorf("worker failed to activate: %w", err) + } + if active.Stage != "active" || active.ID != ready.ID { + return errors.New("invalid activation acknowledgement") + } + if err = enc.Encode(handoffMessage{Stage: "owned"}); err != nil { + return err + } + var serving handoffMessage + if err = dec.Decode(&serving); err != nil { + return fmt.Errorf("worker did not confirm ownership: %w", err) + } + if serving.Stage != "serving" || serving.ID != ready.ID { + return errors.New("invalid serving acknowledgement") + } + success = true + if parent != "" { + gens, _ := p.store.Generations() + for _, g := range gens { + if g.ID == parent && g.State != state.Exited { + if e := privatePost(g.Control, "/internal/stop?timeout="+ready.DrainTimeout.String()); e != nil { + // StopAccepting must be confirmed before acknowledging an upgrade. If the + // old worker is unresponsive, kill and reap it before releasing its grants. + for _, old := range p.workers { + if old.id == parent { + _ = old.cmd.Process.Kill() + <-old.done + p.retire(old) + } + } + log.Printf("[ERROR] retiring unresponsive worker: %v", e) + } + } + } + } + log.Printf("[INFO] supervisor %d activated worker %s (pid %d)", os.Getpid(), w.id, child.Process.Pid) + return nil +} + +// Only the fixed supervisor sends readiness; no MAINPID handoff is required. +func notifySystemd(message string) error { + path := os.Getenv("NOTIFY_SOCKET") + if path == "" { + return nil + } + fd, err := unix.Socket(unix.AF_UNIX, unix.SOCK_DGRAM|unix.SOCK_CLOEXEC, 0) + if err != nil { + return err + } + defer unix.Close(fd) + return unix.Sendto(fd, []byte(message), 0, &unix.SockaddrUnix{Name: path}) +} diff --git a/go.mod b/go.mod index 31bf013..fcc9e61 100644 --- a/go.mod +++ b/go.mod @@ -12,7 +12,6 @@ require ( github.com/spf13/cobra v1.10.2 github.com/stretchr/testify v1.11.1 golang.org/x/sys v0.37.0 - modernc.org/sqlite v1.46.1 ) require ( @@ -21,27 +20,19 @@ require ( github.com/clipperhouse/displaywidth v0.10.0 // indirect github.com/clipperhouse/uax29/v2 v2.6.0 // indirect github.com/davecgh/go-spew v1.1.1 // indirect - github.com/dustin/go-humanize v1.0.1 // indirect - github.com/google/uuid v1.6.0 // indirect github.com/inconshreveable/mousetrap v1.1.0 // indirect github.com/kr/text v0.2.0 // indirect github.com/mattn/go-colorable v0.1.14 // indirect github.com/mattn/go-isatty v0.0.20 // indirect github.com/mattn/go-runewidth v0.0.19 // indirect github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect - github.com/ncruces/go-strftime v1.0.0 // indirect github.com/olekukonko/cat v0.0.0-20250911104152-50322a0618f6 // indirect github.com/olekukonko/errors v1.2.0 // indirect github.com/olekukonko/ll v0.1.6 // indirect github.com/pmezard/go-difflib v1.0.0 // indirect github.com/prometheus/procfs v0.16.1 // indirect - github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect github.com/spf13/pflag v1.0.10 // indirect go.yaml.in/yaml/v2 v2.4.2 // indirect - golang.org/x/exp v0.0.0-20251023183803-a4bb9ffd2546 // indirect google.golang.org/protobuf v1.36.8 // indirect gopkg.in/yaml.v3 v3.0.1 // indirect - modernc.org/libc v1.67.6 // indirect - modernc.org/mathutil v1.7.1 // indirect - modernc.org/memory v1.11.0 // indirect ) diff --git a/go.sum b/go.sum index c2ca8c7..551af35 100644 --- a/go.sum +++ b/go.sum @@ -10,18 +10,10 @@ github.com/cpuguy83/go-md2man/v2 v2.0.6/go.mod h1:oOW0eioCTA6cOiMLiUPZOpcVxMig6N github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E= 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/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY= -github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto= github.com/fatih/color v1.18.0 h1:S8gINlzdQ840/4pfAwic/ZE0djQEH3wM94VfqLTZcOM= github.com/fatih/color v1.18.0/go.mod h1:4FelSpRwEGDpQ12mAdzqdOukCy4u8WUtOY6lkT/6HfU= 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/pprof v0.0.0-20250317173921-a4b03ec1a45e h1:ijClszYn+mADRFY17kjQEVQ1XRhq2/JR1M3sGqeJoxs= -github.com/google/pprof v0.0.0-20250317173921-a4b03ec1a45e/go.mod h1:boTsfXsheKC2y+lKOCMpSfarhxDeIzfZG1jqGcPl3cA= -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/hashicorp/golang-lru/v2 v2.0.7 h1:a+bsQ5rvGLjzHuww6tVxozPZFVghXaHOwFs4luLUK2k= -github.com/hashicorp/golang-lru/v2 v2.0.7/go.mod h1:QeFd9opnmA6QUJc5vARoKUSoFhyfM2/ZepoAG6RGpeM= github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8= github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw= github.com/klauspost/compress v1.18.0 h1:c/Cqfb0r+Yi+JtIEq73FWXVkRonBlf0CRNYc8Zttxdo= @@ -40,8 +32,6 @@ github.com/mattn/go-runewidth v0.0.19 h1:v++JhqYnZuu5jSKrk9RbgF5v4CGUjqRfBm05byF github.com/mattn/go-runewidth v0.0.19/go.mod h1:XBkDxAl56ILZc9knddidhrOlY5R/pDhgLpndooCuJAs= 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/ncruces/go-strftime v1.0.0 h1:HMFp8mLCTPp341M/ZnA4qaf7ZlsbTc+miZjCLOFAw7w= -github.com/ncruces/go-strftime v1.0.0/go.mod h1:Fwc5htZGVVkseilnfgOVb9mKy6w1naJmn9CehxcKcls= github.com/olekukonko/cat v0.0.0-20250911104152-50322a0618f6 h1:zrbMGy9YXpIeTnGj4EljqMiZsIcE09mmF8XsD5AYOJc= github.com/olekukonko/cat v0.0.0-20250911104152-50322a0618f6/go.mod h1:rEKTHC9roVVicUIfZK7DYrdIoM0EOr8mK1Hj5s3JjH0= github.com/olekukonko/errors v1.2.0 h1:10Zcn4GeV59t/EGqJc8fUjtFT/FuUh5bTMzZ1XwmCRo= @@ -62,8 +52,6 @@ github.com/prometheus/common v0.66.1 h1:h5E0h5/Y8niHc5DlaLlWLArTQI7tMrsfQjHV+d9Z github.com/prometheus/common v0.66.1/go.mod h1:gcaUsgf3KfRSwHY4dIMXLPV0K/Wg1oZ8+SbZk/HH/dA= github.com/prometheus/procfs v0.16.1 h1:hZ15bTNuirocR6u0JZ6BAHHmwS1p8B4P6MRqxtzMyRg= github.com/prometheus/procfs v0.16.1/go.mod h1:teAbpZRB1iIAJYREa1LsoWUXykVXA1KlTmWl8x/U+Is= -github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec h1:W09IVJc94icq4NjY3clb7Lk8O1qJ8BdBEF8z0ibU0rE= -github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec/go.mod h1:qqbHyh8v60DhA7CoWK5oRCqLrMHRGoxYCSS9EjAz6Eo= 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/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM= @@ -79,17 +67,9 @@ go.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE= go.yaml.in/yaml/v2 v2.4.2 h1:DzmwEr2rDGHl7lsFgAHxmNz/1NlQ7xLIrlN2h5d1eGI= go.yaml.in/yaml/v2 v2.4.2/go.mod h1:081UH+NErpNdqlCXm3TtEran0rJZGxAYx9hb/ELlsPU= go.yaml.in/yaml/v3 v3.0.4/go.mod h1:DhzuOOF2ATzADvBadXxruRBLzYTpT36CKvDb3+aBEFg= -golang.org/x/exp v0.0.0-20251023183803-a4bb9ffd2546 h1:mgKeJMpvi0yx/sU5GsxQ7p6s2wtOnGAHZWCHUM4KGzY= -golang.org/x/exp v0.0.0-20251023183803-a4bb9ffd2546/go.mod h1:j/pmGrbnkbPtQfxEe5D0VQhZC6qKbfKifgD0oM7sR70= -golang.org/x/mod v0.29.0 h1:HV8lRxZC4l2cr3Zq1LvtOsi/ThTgWnUk/y64QSs8GwA= -golang.org/x/mod v0.29.0/go.mod h1:NyhrlYXJ2H4eJiRy/WDBO6HMqZQ6q9nk4JzS3NuCK+w= -golang.org/x/sync v0.17.0 h1:l60nONMj9l5drqw6jlhIELNv9I0A4OFgRsG9k2oT9Ug= -golang.org/x/sync v0.17.0/go.mod h1:9KTHXmSnoGruLpwFjVSX0lNNA75CykiMECbovNTZqGI= golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.37.0 h1:fdNQudmxPjkdUTPnLn5mdQv7Zwvbvpaxqs831goi9kQ= golang.org/x/sys v0.37.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks= -golang.org/x/tools v0.38.0 h1:Hx2Xv8hISq8Lm16jvBZ2VQf+RLmbd7wVUsALibYI/IQ= -golang.org/x/tools v0.38.0/go.mod h1:yEsQ/d/YK8cjh0L6rZlY8tgtlKiBNTL14pGDJPJpYQs= google.golang.org/protobuf v1.36.8 h1:xHScyCOEuuwZEc6UtSOvPbAT4zRh0xcNRYekJwfqyMc= google.golang.org/protobuf v1.36.8/go.mod h1:fuxRtAxBytpl4zzqUh6/eyUujkJdNiuEkXntxiD/uRU= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= @@ -97,31 +77,3 @@ gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntN gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q= gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= -modernc.org/cc/v4 v4.27.1 h1:9W30zRlYrefrDV2JE2O8VDtJ1yPGownxciz5rrbQZis= -modernc.org/cc/v4 v4.27.1/go.mod h1:uVtb5OGqUKpoLWhqwNQo/8LwvoiEBLvZXIQ/SmO6mL0= -modernc.org/ccgo/v4 v4.30.1 h1:4r4U1J6Fhj98NKfSjnPUN7Ze2c6MnAdL0hWw6+LrJpc= -modernc.org/ccgo/v4 v4.30.1/go.mod h1:bIOeI1JL54Utlxn+LwrFyjCx2n2RDiYEaJVSrgdrRfM= -modernc.org/fileutil v1.3.40 h1:ZGMswMNc9JOCrcrakF1HrvmergNLAmxOPjizirpfqBA= -modernc.org/fileutil v1.3.40/go.mod h1:HxmghZSZVAz/LXcMNwZPA/DRrQZEVP9VX0V4LQGQFOc= -modernc.org/gc/v2 v2.6.5 h1:nyqdV8q46KvTpZlsw66kWqwXRHdjIlJOhG6kxiV/9xI= -modernc.org/gc/v2 v2.6.5/go.mod h1:YgIahr1ypgfe7chRuJi2gD7DBQiKSLMPgBQe9oIiito= -modernc.org/gc/v3 v3.1.1 h1:k8T3gkXWY9sEiytKhcgyiZ2L0DTyCQ/nvX+LoCljoRE= -modernc.org/gc/v3 v3.1.1/go.mod h1:HFK/6AGESC7Ex+EZJhJ2Gni6cTaYpSMmU/cT9RmlfYY= -modernc.org/goabi0 v0.2.0 h1:HvEowk7LxcPd0eq6mVOAEMai46V+i7Jrj13t4AzuNks= -modernc.org/goabi0 v0.2.0/go.mod h1:CEFRnnJhKvWT1c1JTI3Avm+tgOWbkOu5oPA8eH8LnMI= -modernc.org/libc v1.67.6 h1:eVOQvpModVLKOdT+LvBPjdQqfrZq+pC39BygcT+E7OI= -modernc.org/libc v1.67.6/go.mod h1:JAhxUVlolfYDErnwiqaLvUqc8nfb2r6S6slAgZOnaiE= -modernc.org/mathutil v1.7.1 h1:GCZVGXdaN8gTqB1Mf/usp1Y/hSqgI2vAGGP4jZMCxOU= -modernc.org/mathutil v1.7.1/go.mod h1:4p5IwJITfppl0G4sUEDtCr4DthTaT47/N3aT6MhfgJg= -modernc.org/memory v1.11.0 h1:o4QC8aMQzmcwCK3t3Ux/ZHmwFPzE6hf2Y5LbkRs+hbI= -modernc.org/memory v1.11.0/go.mod h1:/JP4VbVC+K5sU2wZi9bHoq2MAkCnrt2r98UGeSK7Mjw= -modernc.org/opt v0.1.4 h1:2kNGMRiUjrp4LcaPuLY2PzUfqM/w9N23quVwhKt5Qm8= -modernc.org/opt v0.1.4/go.mod h1:03fq9lsNfvkYSfxrfUhZCWPk1lm4cq4N+Bh//bEtgns= -modernc.org/sortutil v1.2.1 h1:+xyoGf15mM3NMlPDnFqrteY07klSFxLElE2PVuWIJ7w= -modernc.org/sortutil v1.2.1/go.mod h1:7ZI3a3REbai7gzCLcotuw9AC4VZVpYMjDzETGsSMqJE= -modernc.org/sqlite v1.46.1 h1:eFJ2ShBLIEnUWlLy12raN0Z1plqmFX9Qe3rjQTKt6sU= -modernc.org/sqlite v1.46.1/go.mod h1:CzbrU2lSB1DKUusvwGz7rqEKIq+NUd8GWuBBZDs9/nA= -modernc.org/strutil v1.2.1 h1:UneZBkQA+DX2Rp35KcM69cSsNES9ly8mQWD71HKlOA0= -modernc.org/strutil v1.2.1/go.mod h1:EHkiggD70koQxjVdSBM3JKM7k6L0FbGE5eymy9i3B9A= -modernc.org/token v1.1.0 h1:Xl7Ap9dKaEs5kLoOQeQmPWevfnk/DM5qcLcYlA8ys6Y= -modernc.org/token v1.1.0/go.mod h1:UGzOrNV1mAFSEB63lOFHIpNRUVMvYTc6yu1SMY/XTDM= diff --git a/pkg/server/admission.go b/pkg/server/admission.go index d613860..af48226 100644 --- a/pkg/server/admission.go +++ b/pkg/server/admission.go @@ -6,7 +6,6 @@ import ( "fmt" "log" "sync" - "time" "github.com/ustclug/rsync-proxy/pkg/queue" "github.com/ustclug/rsync-proxy/pkg/state" @@ -36,31 +35,48 @@ type admissionWait struct { } type admission struct { - s *Server - mu sync.Mutex - waiting map[int64]*admissionWait - releases map[int64]bool - done chan struct{} + s *Server + mu sync.Mutex + waiting map[int64]*admissionWait + wake chan struct{} + done chan struct{} } func (s *Server) StartAdmission() func() { - a := &admission{s: s, waiting: make(map[int64]*admissionWait), releases: make(map[int64]bool), done: make(chan struct{})} + a := &admission{s: s, waiting: make(map[int64]*admissionWait), done: make(chan struct{}), wake: make(chan struct{}, 1)} s.admission = a ctx, cancel := context.WithCancel(context.Background()) + events := make(chan struct{}, 1) + watchDone := make(chan struct{}) + go func() { + defer close(watchDone) + var revision uint64 + for { + next, err := s.Shared.Watch(ctx, revision) + if err != nil { + return + } + revision = next + select { + case events <- struct{}{}: + default: + } + } + }() go func() { defer close(a.done) - ticker := time.NewTicker(100 * time.Millisecond) - defer ticker.Stop() for { select { case <-ctx.Done(): return - case <-ticker.C: + case <-events: + a.poll() + case <-a.wake: a.poll() } } }() - return func() { cancel(); <-a.done } + return func() { cancel(); <-watchDone; <-a.done } } func offerStatus(c chan queue.Status, status queue.Status) { @@ -73,15 +89,8 @@ func offerStatus(c chan queue.Status, status queue.Status) { func (a *admission) poll() { a.mu.Lock() - ids := make([]int64, 0, len(a.releases)) - for id := range a.releases { - ids = append(ids, id) - } hasWaiting := len(a.waiting) > 0 a.mu.Unlock() - for _, id := range ids { - a.release(id) - } if !hasWaiting { return } @@ -112,16 +121,10 @@ func (a *admission) release(id int64) { a.mu.Lock() delete(a.waiting, id) a.mu.Unlock() - ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) - defer cancel() - err := a.s.Shared.Release(ctx, a.s.Generation, id) - a.mu.Lock() - defer a.mu.Unlock() - if err != nil { - a.releases[id] = true - log.Printf("[ERROR] release admission (will retry): %v", err) - } else { - delete(a.releases, id) + if err := a.s.Shared.Release(context.Background(), a.s.Generation, id); err != nil { + // Transport failure closes Shared.Done. The worker then exits, and the + // supervisor reclaims its remaining tickets only after reaping it. + log.Printf("[ERROR] release admission: %v", err) } } @@ -152,6 +155,12 @@ func (s *Server) acquire(ctx context.Context, upstream, ip string, q *queue.Queu } c <- queue.Status{Ok: t.Active, Index: t.Index, Max: t.Queued} a.mu.Unlock() + if !t.Active { + select { + case a.wake <- struct{}{}: + default: + } + } var once sync.Once return admissionHandle{c: c, release: func() { once.Do(func() { a.release(t.ID) }) }}, nil } diff --git a/pkg/server/snapshot_test.go b/pkg/server/snapshot_test.go index 007c27a..d152aad 100644 --- a/pkg/server/snapshot_test.go +++ b/pkg/server/snapshot_test.go @@ -13,8 +13,7 @@ import ( ) func TestAggregateRetainsCountersWithoutRetiredGauges(t *testing.T) { - store, err := state.Open(t.TempDir()) - require.NoError(t, err) + store := state.New() defer store.Close() makeGeneration := func(id string) *Server { s := New() diff --git a/pkg/state/ipc.go b/pkg/state/ipc.go new file mode 100644 index 0000000..79f6f53 --- /dev/null +++ b/pkg/state/ipc.go @@ -0,0 +1,272 @@ +package state + +import ( + "context" + "errors" + "fmt" + "net" + "net/rpc" + "os" + "path/filepath" + "sync" + "time" +) + +// Request and Response form the versioned, private supervisor protocol. +// gob framing and request IDs are supplied by net/rpc on a persistent socket. +const operationUpgrade = "upgrade" + +type Request struct { + Protocol int + Op, ID, Parent, Upstream, IP string + Generation Generation + Limits []Limit + TicketID int64 + Snapshot []byte + Deadline time.Time + Revision uint64 + Timeout time.Duration +} +type Response struct { + Error string + Current string + Ticket Ticket + Tickets map[int64]Ticket + Queues []QueueInfo + Generations []Generation + Revision uint64 +} + +// A stopped peer must not leave an RPC sender blocked indefinitely before its +// response timeout even starts. net/rpc serializes writes on each connection. +type deadlineConn struct{ net.Conn } + +func (c deadlineConn) Write(b []byte) (int, error) { + if err := c.SetWriteDeadline(time.Now().Add(5 * time.Second)); err != nil { + return 0, err + } + return c.Conn.Write(b) +} + +type remoteStore struct { + client *rpc.Client + done chan struct{} + once sync.Once +} + +func (r *remoteStore) close() error { + var err error + r.once.Do(func() { close(r.done); err = r.client.Close() }) + return err +} + +func Open(dir string) (*Store, error) { + conn, err := net.DialTimeout("unix", filepath.Join(dir, "supervisor.sock"), 3*time.Second) + if err != nil { + return nil, fmt.Errorf("connect to supervisor: %w", err) + } + s := &Store{remote: &remoteStore{client: rpc.NewClient(deadlineConn{conn}), done: make(chan struct{})}} + if _, err = s.call(Request{Op: "hello"}); err != nil { + _ = s.Close() + return nil, err + } + return s, nil +} + +// Done closes if the control connection fails. Workers must then stop serving; +// they cannot establish a fresh authority and silently forget existing grants. +func (s *Store) Done() <-chan struct{} { + if s.remote == nil { + return nil + } + return s.remote.done +} + +func decodeError(msg string) error { + switch msg { + case "": + return nil + case ErrFull.Error(): + return ErrFull + case ErrPerIP.Error(): + return ErrPerIP + case ErrNotCurrent.Error(): + return ErrNotCurrent + default: + return errors.New(msg) + } +} + +func (s *Store) call(req Request) (Response, error) { + req.Protocol = Protocol + var result Response + call := s.remote.client.Go("Coordinator.Call", req, &result, make(chan *rpc.Call, 1)) + timeout := 5 * time.Second + if req.Op == operationUpgrade { + timeout = req.Timeout + 5*time.Second + } + timer := time.NewTimer(timeout) + defer timer.Stop() + select { + case completed := <-call.Done: + if completed.Error != nil { + _ = s.Close() + return Response{}, completed.Error + } + return result, decodeError(result.Error) + case <-timer.C: + _ = s.Close() + return Response{}, errors.New("supervisor request timed out") + } +} + +func (s *Store) watchRemote(ctx context.Context, after uint64) (uint64, error) { + for { + var r Response + call := s.remote.client.Go("Coordinator.Call", Request{Protocol: Protocol, Op: "watch", Revision: after}, &r, make(chan *rpc.Call, 1)) + timer := time.NewTimer(5 * time.Second) + select { + case <-ctx.Done(): + timer.Stop() + return after, ctx.Err() + case <-timer.C: + _ = s.Close() + return after, errors.New("supervisor watch timed out") + case c := <-call.Done: + timer.Stop() + if c.Error != nil { + _ = s.Close() + return after, c.Error + } + if err := decodeError(r.Error); err != nil { + return after, err + } + if r.Revision != after { + return r.Revision, nil + } + } + } +} + +func (s *Store) Upgrade(id string, timeout time.Duration) error { + if s.remote == nil { + return errors.New("upgrade requires supervisor client") + } + _, e := s.call(Request{Op: operationUpgrade, ID: id, Timeout: timeout}) + return e +} + +type coordinator struct { + store *Store + upgrade func(string, time.Duration) error +} + +func (c *coordinator) Call(q Request, r *Response) error { + if q.Protocol != Protocol { + r.Error = "incompatible supervisor protocol" + return nil + } + var err error + switch q.Op { + case "hello": + case "current": + r.Current, err = c.store.Current() + case "register": + err = c.store.Register(q.Generation) + case "activate": + err = c.store.Activate(q.ID, q.Parent, q.Limits) + case "limits": + err = c.store.SetLimits(q.ID, q.Limits) + case "acquire": + r.Ticket, err = c.store.Acquire(context.Background(), q.ID, q.Upstream, q.IP) + case "release": + err = c.store.Release(context.Background(), q.ID, q.TicketID) + case "poll": + r.Tickets, err = c.store.Poll(q.ID) + case "queues": + r.Queues, err = c.store.Queues() + case "generations": + r.Generations, err = c.store.Generations() + case "publish": + err = c.store.Publish(q.ID, q.Snapshot) + case "drain": + err = c.store.Drain(q.ID, q.Deadline) + case "watch": + // Bound abandoned watches without changing their event-driven semantics. + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + r.Revision, err = c.store.Watch(ctx, q.Revision) + if errors.Is(err, context.DeadlineExceeded) { + err = nil + } + case operationUpgrade: + if c.upgrade == nil { + err = errors.New("upgrade unavailable") + } else { + err = c.upgrade(q.ID, q.Timeout) + } + default: + err = fmt.Errorf("unsupported supervisor operation %q", q.Op) + } + if err != nil { + r.Error = err.Error() + } + return nil +} + +// Serve exposes a store owned by the supervisor. The caller must hold the +// instance lifetime lock before removing a stale socket or starting this server. +func Serve(dir string, s *Store, upgrade func(string, time.Duration) error) (func(), error) { + if err := os.MkdirAll(dir, 0700); err != nil { + return nil, err + } + path := filepath.Join(dir, "supervisor.sock") + l, err := net.Listen("unix", path) + if err != nil { + return nil, err + } + if err = os.Chmod(path, 0600); err != nil { + _ = l.Close() + return nil, err + } + srv := rpc.NewServer() + if err = srv.RegisterName("Coordinator", &coordinator{store: s, upgrade: upgrade}); err != nil { + _ = l.Close() + return nil, err + } + var mu sync.Mutex + clients := make(map[net.Conn]bool) + var wg sync.WaitGroup + accepted := make(chan struct{}) + go func() { + defer close(accepted) + for { + conn, e := l.Accept() + if e != nil { + return + } + mu.Lock() + clients[conn] = true + mu.Unlock() + wg.Go(func() { + srv.ServeConn(deadlineConn{conn}) + mu.Lock() + delete(clients, conn) + mu.Unlock() + }) + } + }() + var once sync.Once + return func() { + once.Do(func() { + _ = l.Close() + <-accepted + mu.Lock() + for c := range clients { + _ = c.Close() + } + mu.Unlock() + wg.Wait() + }) + }, nil +} diff --git a/pkg/state/store.go b/pkg/state/store.go index a1f7f09..1b5869b 100644 --- a/pkg/state/store.go +++ b/pkg/state/store.go @@ -1,22 +1,17 @@ -// Package state coordinates admission across generations of one local proxy. -// The database contains runtime state, not rsync payloads or durable history. +// Package state coordinates worker generations through the supervisor's memory. package state import ( "context" - "database/sql" "encoding/json" "errors" "fmt" - "net/url" - "os" - "path/filepath" + "sort" + "sync" "time" - - _ "modernc.org/sqlite" // Register the CGo-free database/sql driver. ) -const Protocol = 1 +const Protocol = 2 const ( Starting = "starting" @@ -57,295 +52,380 @@ type QueueInfo struct { ActiveCount, QueuedCount int } +// Store is either the supervisor's in-memory store or a Unix socket client. +// The mutex covers only memory operations; no I/O is performed while holding it. type Store struct { - db *sql.DB + mu sync.Mutex + current string + generations map[string]Generation + order []string + limits map[string]Limit + tickets map[int64]entry + next int64 + revision uint64 + changed chan struct{} + remote *remoteStore } -func Open(dir string) (*Store, error) { - if err := os.MkdirAll(dir, 0700); err != nil { - return nil, err - } - u := url.URL{Scheme: "file", Path: filepath.Join(dir, "state.db")} - q := u.Query() - q.Add("_pragma", "busy_timeout(2000)") - q.Add("_pragma", "foreign_keys(1)") - u.RawQuery = q.Encode() - db, err := sql.Open("sqlite", u.String()) - if err != nil { - return nil, err - } - db.SetMaxOpenConns(1) - s := &Store{db: db} - err = s.write(context.Background(), func(c *sql.Conn) error { - var version int - if err := c.QueryRowContext(context.Background(), "PRAGMA user_version").Scan(&version); err != nil { - return err - } - if version != 0 && version != Protocol { - return fmt.Errorf("incompatible state schema %d (supported %d)", version, Protocol) - } - if version == Protocol { - return nil - } - _, err := c.ExecContext(context.Background(), ` -CREATE TABLE generations(id TEXT PRIMARY KEY, pid INTEGER NOT NULL, version TEXT NOT NULL, control TEXT NOT NULL, state TEXT NOT NULL, deadline INTEGER NOT NULL DEFAULT 0, snapshot BLOB); -CREATE TABLE settings(key TEXT PRIMARY KEY, value TEXT NOT NULL); -INSERT INTO settings VALUES('current',''); -CREATE TABLE limits(name TEXT PRIMARY KEY, active INTEGER NOT NULL, queued INTEGER NOT NULL, per_ip INTEGER NOT NULL); -CREATE TABLE tickets(id INTEGER PRIMARY KEY AUTOINCREMENT, generation TEXT NOT NULL REFERENCES generations(id), upstream TEXT NOT NULL REFERENCES limits(name), ip TEXT NOT NULL, active INTEGER NOT NULL); -CREATE INDEX tickets_upstream ON tickets(upstream, active, id); -CREATE INDEX tickets_ip ON tickets(upstream, ip); -CREATE INDEX tickets_generation ON tickets(generation); -PRAGMA user_version=1;`) - return err - }) - if err == nil { - err = os.Chmod(filepath.Join(dir, "state.db"), 0600) - } - if err == nil { - _, err = db.Exec("PRAGMA journal_mode=WAL") - } - if err != nil { - _ = db.Close() - return nil, err - } - return s, nil +type entry struct { + Ticket + Generation, Upstream, IP string } -func (s *Store) Close() error { return s.db.Close() } +func New() *Store { + return &Store{generations: make(map[string]Generation), limits: make(map[string]Limit), tickets: make(map[int64]entry), changed: make(chan struct{})} +} -// Every read/modify/write operation starts as a writer. In particular, no -// deferred read transaction may race another process between COUNT and INSERT. -func (s *Store) write(ctx context.Context, f func(*sql.Conn) error) error { - c, err := s.db.Conn(ctx) - if err != nil { - return err - } - defer c.Close() - if _, err = c.ExecContext(ctx, "BEGIN IMMEDIATE"); err != nil { - return err +func (s *Store) Close() error { + if s.remote != nil { + return s.remote.close() } - defer func() { _, _ = c.ExecContext(context.Background(), "ROLLBACK") }() - if err = f(c); err != nil { - return err - } - // Once commit starts, do not let request cancellation make its outcome - // ambiguous to the caller (which must release any successfully issued ticket). - if err = ctx.Err(); err != nil { - return err - } - _, err = c.ExecContext(context.Background(), "COMMIT") - return err + return nil } -func (s *Store) Register(g Generation) error { - _, err := s.db.Exec("INSERT INTO generations(id,pid,version,control,state) VALUES(?,?,?,?,?)", g.ID, g.PID, g.Version, g.Control, "starting") - return err +func (s *Store) changedLocked() { + s.revision++ + close(s.changed) + s.changed = make(chan struct{}) } func (s *Store) Current() (string, error) { - var id string - err := s.db.QueryRow("SELECT value FROM settings WHERE key='current'").Scan(&id) - return id, err + if s.remote != nil { + r, e := s.call(Request{Op: "current"}) + return r.Current, e + } + s.mu.Lock() + defer s.mu.Unlock() + return s.current, nil +} + +func (s *Store) Register(g Generation) error { + if s.remote != nil { + _, e := s.call(Request{Op: "register", Generation: g}) + return e + } + s.mu.Lock() + defer s.mu.Unlock() + if g.ID == "" { + return errors.New("empty generation") + } + if _, ok := s.generations[g.ID]; ok { + return errors.New("generation already registered") + } + g.State = Starting + s.generations[g.ID] = g + s.order = append(s.order, g.ID) + return nil } -func setLimits(ctx context.Context, c *sql.Conn, limits []Limit) error { +func validateLimits(limits []Limit) error { for _, l := range limits { - if l.Active < 0 || l.Queued < 0 || l.PerIP < 0 { - return errors.New("negative connection limit") - } - _, err := c.ExecContext(ctx, "INSERT INTO limits VALUES(?,?,?,?) ON CONFLICT(name) DO UPDATE SET active=excluded.active, queued=excluded.queued, per_ip=excluded.per_ip", l.Name, l.Active, l.Queued, l.PerIP) - if err != nil { - return err + if l.Name == "" || l.Active < 0 || l.Queued < 0 || l.PerIP < 0 { + return errors.New("invalid admission limit") } } - return promote(ctx, c) + return nil } -// Activate is the commit point of an upgrade. Old generations retain their -// selected targets but consult these shared limits for every admission. func (s *Store) Activate(id, parent string, limits []Limit) error { - ctx := context.Background() - return s.write(ctx, func(c *sql.Conn) error { - var current string - if err := c.QueryRowContext(ctx, "SELECT value FROM settings WHERE key='current'").Scan(¤t); err != nil { - return err - } - if current != parent { - return ErrNotCurrent - } - if err := setLimits(ctx, c, limits); err != nil { - return err - } - if _, err := c.ExecContext(ctx, "UPDATE settings SET value=? WHERE key='current'", id); err != nil { - return err - } - _, err := c.ExecContext(ctx, "UPDATE generations SET state='serving' WHERE id=?", id) + if s.remote != nil { + _, e := s.call(Request{Op: "activate", ID: id, Parent: parent, Limits: limits}) + return e + } + if err := validateLimits(limits); err != nil { return err - }) + } + s.mu.Lock() + defer s.mu.Unlock() + if s.current != parent { + return ErrNotCurrent + } + g, ok := s.generations[id] + if !ok || g.State == Exited { + return errors.New("unknown or exited generation") + } + for _, l := range limits { + s.limits[l.Name] = l + } + s.current = id + g.State = Serving + g.Deadline = 0 + s.generations[id] = g + s.promoteLocked() + s.changedLocked() + return nil } func (s *Store) SetLimits(id string, limits []Limit) error { - ctx := context.Background() - return s.write(ctx, func(c *sql.Conn) error { - var current string - if err := c.QueryRowContext(ctx, "SELECT value FROM settings WHERE key='current'").Scan(¤t); err != nil { - return err - } - if current != id { - return ErrNotCurrent - } - return setLimits(ctx, c, limits) - }) + if s.remote != nil { + _, e := s.call(Request{Op: "limits", ID: id, Limits: limits}) + return e + } + if err := validateLimits(limits); err != nil { + return err + } + s.mu.Lock() + defer s.mu.Unlock() + if s.current != id { + return ErrNotCurrent + } + for _, l := range limits { + s.limits[l.Name] = l + } + s.promoteLocked() + s.changedLocked() + return nil } -func promote(ctx context.Context, c *sql.Conn) error { - // Rank only queued tickets. Already granted tickets (including grants whose - // owner hasn't polled yet) always consume capacity. - _, err := c.ExecContext(ctx, `WITH waiting AS ( - SELECT t.id, l.active AS max, ROW_NUMBER() OVER(PARTITION BY t.upstream ORDER BY t.id) AS position, - (SELECT COUNT(*) FROM tickets a WHERE a.upstream=t.upstream AND a.active=1) AS occupied - FROM tickets t JOIN limits l ON l.name=t.upstream WHERE t.active=0 -) UPDATE tickets SET active=1 WHERE id IN (SELECT id FROM waiting WHERE max=0 OR position<=max-occupied)`) - return err +func (s *Store) sortedLocked() []int64 { + ids := make([]int64, 0, len(s.tickets)) + for id := range s.tickets { + ids = append(ids, id) + } + sort.Slice(ids, func(i, j int) bool { return ids[i] < ids[j] }) + return ids } -func (s *Store) Acquire(ctx context.Context, generation, upstream, ip string) (Ticket, error) { - var t Ticket - err := s.write(ctx, func(c *sql.Conn) error { - var max, queuedMax, perIP, active, queued, ips int - if err := c.QueryRowContext(ctx, "SELECT active,queued,per_ip FROM limits WHERE name=?", upstream).Scan(&max, &queuedMax, &perIP); err != nil { - return err - } - if err := c.QueryRowContext(ctx, "SELECT COUNT(*) FROM tickets WHERE upstream=? AND ip=?", upstream, ip).Scan(&ips); err != nil { - return err - } - if perIP > 0 && ips >= perIP { - return fmt.Errorf("per-IP cap of %d reached: %w", perIP, ErrPerIP) +func (s *Store) promoteLocked() { + active := make(map[string]int) + for _, t := range s.tickets { + if t.Active { + active[t.Upstream]++ } - if err := promote(ctx, c); err != nil { - return err + } + for _, id := range s.sortedLocked() { + t := s.tickets[id] + l := s.limits[t.Upstream] + if !t.Active && (l.Active == 0 || active[t.Upstream] < l.Active) { + t.Active = true + s.tickets[id] = t + active[t.Upstream]++ } - if err := c.QueryRowContext(ctx, "SELECT COUNT(*) FILTER(WHERE active=1),COUNT(*) FILTER(WHERE active=0) FROM tickets WHERE upstream=?", upstream).Scan(&active, &queued); err != nil { - return err + } +} + +func (s *Store) Acquire(ctx context.Context, generation, upstream, ip string) (Ticket, error) { + if err := ctx.Err(); err != nil { + return Ticket{}, err + } + if s.remote != nil { + // Once sent, await the definite result. Cancellation must not orphan a grant. + r, e := s.call(Request{Op: "acquire", ID: generation, Upstream: upstream, IP: ip}) + return r.Ticket, e + } + s.mu.Lock() + defer s.mu.Unlock() + g, ok := s.generations[generation] + if !ok || g.State == Exited { + return Ticket{}, errors.New("unknown or exited generation") + } + l, ok := s.limits[upstream] + if !ok { + return Ticket{}, fmt.Errorf("unknown upstream %q", upstream) + } + active, queued, ips := 0, 0, 0 + for _, t := range s.tickets { + if t.Upstream != upstream { + continue } - t.Active = max == 0 || (active < max && queued == 0) - if !t.Active && queuedMax > 0 && queued >= queuedMax { - return ErrFull + if t.IP == ip { + ips++ } - r, err := c.ExecContext(ctx, "INSERT INTO tickets(generation,upstream,ip,active) VALUES(?,?,?,?)", generation, upstream, ip, t.Active) - if err != nil { - return err + if t.Active { + active++ + } else { + queued++ } - t.ID, err = r.LastInsertId() - t.Max = max - t.Index = queued - t.Queued = queued + 1 - return err - }) - return t, err + } + if l.PerIP > 0 && ips >= l.PerIP { + return Ticket{}, ErrPerIP + } + granted := queued == 0 && (l.Active == 0 || active < l.Active) + if !granted && l.Queued > 0 && queued >= l.Queued { + return Ticket{}, ErrFull + } + s.next++ + t := Ticket{ID: s.next, Active: granted, Index: queued, Queued: queued + 1, Max: l.Active} + s.tickets[t.ID] = entry{Ticket: t, Generation: generation, Upstream: upstream, IP: ip} + s.changedLocked() + return t, nil } func (s *Store) Release(ctx context.Context, generation string, id int64) error { - return s.write(ctx, func(c *sql.Conn) error { - if _, err := c.ExecContext(ctx, "DELETE FROM tickets WHERE generation=? AND id=?", generation, id); err != nil { - return err - } - return promote(ctx, c) - }) + if err := ctx.Err(); err != nil { + return err + } + if s.remote != nil { + _, e := s.call(Request{Op: "release", ID: generation, TicketID: id}) + return e + } + s.mu.Lock() + defer s.mu.Unlock() + if t, ok := s.tickets[id]; ok && t.Generation == generation { + delete(s.tickets, id) + s.promoteLocked() + s.changedLocked() + } + return nil } -// Poll reads all of a generation's queued/granted tickets in one operation. func (s *Store) Poll(generation string) (map[int64]Ticket, error) { - rows, err := s.db.Query(`WITH ranked AS ( - SELECT t.id,t.generation,t.active,l.active AS max, - COUNT(*) FILTER(WHERE t.active=0) OVER(PARTITION BY t.upstream ORDER BY t.id ROWS BETWEEN UNBOUNDED PRECEDING AND 1 PRECEDING) AS position, - COUNT(*) FILTER(WHERE t.active=0) OVER(PARTITION BY t.upstream) AS queued - FROM tickets t JOIN limits l ON l.name=t.upstream -) SELECT id,active,max,position,queued FROM ranked WHERE generation=?`, generation) - if err != nil { - return nil, err - } - defer rows.Close() + if s.remote != nil { + r, e := s.call(Request{Op: "poll", ID: generation}) + return r.Tickets, e + } + s.mu.Lock() + defer s.mu.Unlock() + return s.pollLocked(generation), nil +} + +func (s *Store) pollLocked(generation string) map[int64]Ticket { + queued := make(map[string]int) + position := make(map[string]int) + for _, t := range s.tickets { + if !t.Active { + queued[t.Upstream]++ + } + } result := make(map[int64]Ticket) - for rows.Next() { - var t Ticket - if err := rows.Scan(&t.ID, &t.Active, &t.Max, &t.Index, &t.Queued); err != nil { - return nil, err + for _, id := range s.sortedLocked() { + t := s.tickets[id] + if t.Generation == generation { + v := t.Ticket + v.Index = position[t.Upstream] + v.Queued = queued[t.Upstream] + v.Max = s.limits[t.Upstream].Active + result[id] = v + } + if !t.Active { + position[t.Upstream]++ + } + } + return result +} + +// Watch waits for an admission change. The revision closes the race between +// taking a snapshot and subscribing. No periodic queue scan is necessary. +func (s *Store) Watch(ctx context.Context, after uint64) (uint64, error) { + if s.remote != nil { + return s.watchRemote(ctx, after) + } + for { + s.mu.Lock() + rev, ch := s.revision, s.changed + s.mu.Unlock() + if rev != after { + return rev, nil + } + select { + case <-ctx.Done(): + return after, ctx.Err() + case <-ch: } - result[t.ID] = t } - return result, rows.Err() } func (s *Store) Queues() ([]QueueInfo, error) { - rows, err := s.db.Query(`SELECT l.name,l.active,l.queued,l.per_ip,COUNT(t.id) FILTER(WHERE t.active=1),COUNT(t.id) FILTER(WHERE t.active=0) FROM limits l LEFT JOIN tickets t ON t.upstream=l.name GROUP BY l.name ORDER BY l.name`) - if err != nil { - return nil, err - } - defer rows.Close() - var result []QueueInfo - for rows.Next() { - var q QueueInfo - if err := rows.Scan(&q.Name, &q.Active, &q.Queued, &q.PerIP, &q.ActiveCount, &q.QueuedCount); err != nil { - return nil, err + if s.remote != nil { + r, e := s.call(Request{Op: "queues"}) + return r.Queues, e + } + s.mu.Lock() + defer s.mu.Unlock() + counts := make(map[string]QueueInfo) + for name, l := range s.limits { + counts[name] = QueueInfo{Limit: l} + } + for _, t := range s.tickets { + q := counts[t.Upstream] + if t.Active { + q.ActiveCount++ + } else { + q.QueuedCount++ } + counts[t.Upstream] = q + } + result := make([]QueueInfo, 0, len(counts)) + for _, q := range counts { result = append(result, q) } - return result, rows.Err() + sort.Slice(result, func(i, j int) bool { return result[i].Name < result[j].Name }) + return result, nil } func (s *Store) Generations() ([]Generation, error) { - rows, err := s.db.Query("SELECT id,pid,version,control,state,deadline,snapshot FROM generations ORDER BY rowid") - if err != nil { - return nil, err - } - defer rows.Close() - var result []Generation - for rows.Next() { - var g Generation - var snapshot []byte - if err := rows.Scan(&g.ID, &g.PID, &g.Version, &g.Control, &g.State, &g.Deadline, &snapshot); err != nil { - return nil, err - } - g.Snapshot = json.RawMessage(snapshot) + if s.remote != nil { + r, e := s.call(Request{Op: "generations"}) + return r.Generations, e + } + s.mu.Lock() + defer s.mu.Unlock() + result := make([]Generation, 0, len(s.order)) + for _, id := range s.order { + g := s.generations[id] + g.Snapshot = append(json.RawMessage(nil), g.Snapshot...) result = append(result, g) } - return result, rows.Err() + return result, nil } func (s *Store) Publish(id string, snapshot []byte) error { - _, err := s.db.Exec("UPDATE generations SET snapshot=? WHERE id=?", snapshot, id) - return err + if s.remote != nil { + _, e := s.call(Request{Op: "publish", ID: id, Snapshot: snapshot}) + return e + } + s.mu.Lock() + defer s.mu.Unlock() + g, ok := s.generations[id] + if !ok { + return errors.New("unknown generation") + } + g.Snapshot = append(json.RawMessage(nil), snapshot...) + s.generations[id] = g + return nil } func (s *Store) Drain(id string, deadline time.Time) error { - var seconds int64 - if !deadline.IsZero() { - seconds = deadline.Unix() + if s.remote != nil { + _, e := s.call(Request{Op: "drain", ID: id, Deadline: deadline}) + return e + } + s.mu.Lock() + defer s.mu.Unlock() + g, ok := s.generations[id] + if !ok { + return errors.New("unknown generation") + } + if g.State == Serving { + g.State = Draining + if !deadline.IsZero() { + g.Deadline = deadline.Unix() + } + s.generations[id] = g } - _, err := s.db.Exec("UPDATE generations SET state='draining',deadline=? WHERE id=? AND state='serving'", seconds, id) - return err + return nil } +// Finish is called by the supervisor only after Wait confirms process death. +// Losing an IPC connection alone is not proof that its transfers have ended. func (s *Store) Finish(id string) error { - ctx := context.Background() - return s.write(ctx, func(c *sql.Conn) error { - if _, err := c.ExecContext(ctx, "DELETE FROM tickets WHERE generation=?", id); err != nil { - return err - } - if _, err := c.ExecContext(ctx, "UPDATE generations SET state='exited' WHERE id=?", id); err != nil { - return err + if s.remote != nil { + return errors.New("only the supervisor can retire workers") + } + s.mu.Lock() + defer s.mu.Unlock() + g, ok := s.generations[id] + if !ok { + return nil + } + for key, t := range s.tickets { + if t.Generation == id { + delete(s.tickets, key) } - return promote(ctx, c) - }) -} - -// Reset is only called under the startup lock after confirming all generation -// locks are unheld. It deliberately preserves the schema and database inode. -func (s *Store) Reset() error { - return s.write(context.Background(), func(c *sql.Conn) error { - _, err := c.ExecContext(context.Background(), "DELETE FROM tickets; DELETE FROM generations; DELETE FROM limits; UPDATE settings SET value='' WHERE key='current';") - return err - }) + } + g.State = Exited + s.generations[id] = g + s.promoteLocked() + s.changedLocked() + return nil } diff --git a/pkg/state/store_test.go b/pkg/state/store_test.go index 5061929..cb27f82 100644 --- a/pkg/state/store_test.go +++ b/pkg/state/store_test.go @@ -5,6 +5,8 @@ import ( "context" "errors" "fmt" + "net" + "net/rpc" "os" "os/exec" "path/filepath" @@ -16,6 +18,22 @@ import ( "github.com/stretchr/testify/require" ) +func startTestCoordinator(t *testing.T, dir string) *Store { + t.Helper() + owner := New() + stop, err := Serve(dir, owner, nil) + require.NoError(t, err) + t.Cleanup(stop) + return owner +} + +func newTestStore(t *testing.T, id string) *Store { + t.Helper() + dir := t.TempDir() + startTestCoordinator(t, dir) + return openTestStore(t, dir, id) +} + func openTestStore(t *testing.T, dir, id string) *Store { t.Helper() s, err := Open(dir) @@ -27,6 +45,7 @@ func openTestStore(t *testing.T, dir, id string) *Store { func TestGlobalFIFOAndPolicy(t *testing.T) { dir := t.TempDir() + owner := startTestCoordinator(t, dir) a := openTestStore(t, dir, "a") b := openTestStore(t, dir, "b") ctx := context.Background() @@ -54,7 +73,7 @@ func TestGlobalFIFOAndPolicy(t *testing.T) { at, err := a.Poll("a") require.NoError(t, err) require.False(t, at[three.ID].Active) - require.NoError(t, b.Finish("b")) + require.NoError(t, owner.Finish("b")) at, err = a.Poll("a") require.NoError(t, err) require.True(t, at[three.ID].Active) @@ -63,7 +82,7 @@ func TestGlobalFIFOAndPolicy(t *testing.T) { } func TestLowerLimitAndUnlimitedQueue(t *testing.T) { - s := openTestStore(t, t.TempDir(), "a") + s := newTestStore(t, "a") ctx := context.Background() require.NoError(t, s.Activate("a", "", []Limit{{Name: "u", Active: 2}})) one, err := s.Acquire(ctx, "a", "u", "ip") @@ -88,29 +107,54 @@ func TestLowerLimitAndUnlimitedQueue(t *testing.T) { require.Equal(t, 49, q[0].QueuedCount) } -func TestSchemaRejectionAndBusy(t *testing.T) { +func TestProtocolRejectionAndCancellation(t *testing.T) { dir := t.TempDir() - s := openTestStore(t, dir, "a") - _, err := s.db.Exec("PRAGMA user_version=999") + startTestCoordinator(t, dir) + conn, err := net.Dial("unix", filepath.Join(dir, "supervisor.sock")) require.NoError(t, err) - _, err = Open(dir) - require.ErrorContains(t, err, "incompatible state schema") - _, err = s.db.Exec("PRAGMA user_version=1") + client := rpc.NewClient(conn) + defer client.Close() + var response Response + require.NoError(t, client.Call("Coordinator.Call", Request{Protocol: 999, Op: "hello"}, &response)) + require.Contains(t, response.Error, "incompatible supervisor protocol") + s := openTestStore(t, dir, "a") + require.NoError(t, s.Activate("a", "", []Limit{{Name: "u", Active: 1}})) + ctx, cancel := context.WithCancel(context.Background()) + cancel() + _, err = s.Acquire(ctx, "a", "u", "ip") + require.ErrorIs(t, err, context.Canceled) + q, err := s.Queues() require.NoError(t, err) + require.Zero(t, q[0].ActiveCount) +} + +func TestDisconnectDoesNotReleaseLiveWorkers(t *testing.T) { + dir := t.TempDir() + owner := startTestCoordinator(t, dir) + a := openTestStore(t, dir, "a") b := openTestStore(t, dir, "b") - require.NoError(t, s.Activate("a", "", []Limit{{Name: "u", Active: 1}})) - c, err := s.db.Conn(context.Background()) + require.NoError(t, a.Activate("a", "", []Limit{{Name: "u", Active: 1}})) + _, err := a.Acquire(context.Background(), "a", "u", "ip") require.NoError(t, err) - _, err = c.ExecContext(context.Background(), "BEGIN IMMEDIATE") + queued, err := b.Acquire(context.Background(), "b", "u", "other") require.NoError(t, err) - _, err = b.Acquire(context.Background(), "b", "u", "ip") - require.Error(t, err) - _, err = c.ExecContext(context.Background(), "ROLLBACK") + require.False(t, queued.Active) + require.NoError(t, a.Close()) + q, err := b.Queues() require.NoError(t, err) - require.NoError(t, c.Close()) - q, err := s.Queues() + require.Equal(t, 1, q[0].ActiveCount) + require.Equal(t, 1, q[0].QueuedCount) + rev, err := owner.Watch(context.Background(), 0) require.NoError(t, err) - require.Zero(t, q[0].ActiveCount) + changed := make(chan uint64, 1) + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + go func() { next, _ := b.Watch(ctx, rev); changed <- next }() + require.NoError(t, owner.Finish("a")) + require.Greater(t, <-changed, rev) + tickets, err := b.Poll("b") + require.NoError(t, err) + require.True(t, tickets[queued.ID].Active) } func TestStateWorker(t *testing.T) { @@ -159,6 +203,7 @@ func TestStateWorker(t *testing.T) { func TestConcurrentProcessAdmission(t *testing.T) { dir := t.TempDir() + startTestCoordinator(t, dir) s := openTestStore(t, dir, "parent") require.NoError(t, s.Activate("parent", "", []Limit{{Name: "u", Active: 2, Queued: 20, PerIP: 2}})) executable, err := os.Executable() @@ -213,7 +258,7 @@ func TestConcurrentProcessAdmission(t *testing.T) { } func TestAdmissionAt100ConnectionsPerSecond(t *testing.T) { - s := openTestStore(t, t.TempDir(), "a") + s := newTestStore(t, "a") require.NoError(t, s.Activate("a", "", []Limit{{Name: "u", Active: 10, Queued: 100}})) latencies := make([]time.Duration, 0, 200) ticker := time.NewTicker(10 * time.Millisecond) @@ -232,6 +277,7 @@ func TestAdmissionAt100ConnectionsPerSecond(t *testing.T) { func TestQueueBurst(t *testing.T) { dir := t.TempDir() + startTestCoordinator(t, dir) s := openTestStore(t, dir, "a") _ = openTestStore(t, dir, "b") ctx := context.Background() diff --git a/test/systemd/upgrade.py b/test/systemd/upgrade.py index ff6d099..a2ba675 100644 --- a/test/systemd/upgrade.py +++ b/test/systemd/upgrade.py @@ -1,5 +1,5 @@ #!/usr/bin/env python3 -"""Exercise MAINPID handoff using an isolated user-level systemd service.""" +"""Exercise worker upgrades with a stable supervisor MAINPID using an isolated user-level systemd service.""" import argparse import contextlib @@ -94,8 +94,8 @@ def main(): try: subprocess.run( ["systemd-run", "--user", f"--unit={unit}", "--collect", - "--property=Type=notify", "--property=NotifyAccess=all", - "--property=KillMode=control-group", "--property=TimeoutStopSec=5s", + "--property=Type=notify", "--property=NotifyAccess=main", + "--property=KillMode=mixed", "--property=TimeoutStopSec=5s", str(binary), "--config", str(config)], check=True, ) @@ -107,7 +107,7 @@ def main(): control(control_socket, "POST", "/upgrade?timeout=30s") assert property_value(unit, "ActiveState") == "active" pid = property_value(unit, "MainPID") - assert pid != "0" and pid not in pids, pids + assert pid != "0" and (not pids or pid == pids[0]), pids pids.append(pid) connection = socket.socket(socket.AF_UNIX) connection.settimeout(5) @@ -143,7 +143,28 @@ def main(): assert property_value(unit, "MainPID") == pids[-1] assert property_value(unit, "ActiveState") == "active" assert rsync_socket.exists() and control_socket.exists() - print(f"PASS: MAINPID {' -> '.join(pids)}; all streams survived; old generations exited") + # TERM goes only to the supervisor, which drains its workers. + # Sending TERM to the whole cgroup as well would deliver it twice. + connection = socket.socket(socket.AF_UNIX) + connection.settimeout(5) + connection.connect(str(rsync_socket)) + connections.append(connection) + reader = connection.makefile("rb") + readers.append(reader) + connection.sendall(b"@RSYNCD: 32.0\n") + assert reader.readline().startswith(b"@RSYNCD:") + connection.sendall(b"foo\n") + assert reader.readline() == b"READY\n" + subprocess.run(["systemctl", "--user", "stop", "--no-block", unit], check=True) + deadline = time.monotonic() + 3 + while property_value(unit, "ActiveState") != "deactivating": + assert time.monotonic() < deadline + time.sleep(0.02) + connection.sendall(b"still draining\n") + assert reader.readline() == b"still draining\n" + reader.close() + connection.close() + print(f"PASS: supervisor MAINPID {pids[0]} stayed fixed; all streams survived; old workers exited; stop drained connections") finally: for reader in readers: reader.close()