diff --git a/README.md b/README.md index 8b1726f..14eca9d 100644 --- a/README.md +++ b/README.md @@ -55,6 +55,32 @@ 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`，保存 supervisor 的 Unix socket 和进程锁文件。全局限额、FIFO 队列和各代快照保存在 supervisor 内存中，不使用数据库。手动运行时须使用可写目录，各实例须使用独立目录；不要在进程运行期间删除 socket 或锁文件。 +- `[proxy].upgrade_drain_timeout` 单位为秒，默认 `0`。非零值从该代停止接收新连接时起计算，到期关闭剩余连接，包括握手和排队。采用本次成功升级的新配置值；后续更新不会改变或重置更早各代的期限。默认不增加期限，原有 idle/max-duration/throughput 规则继续生效。 +- 所有代合计执行上游并发、排队和单 IP 限额，排队按全局 FIFO 晋升。调低限额不会主动断开已有连接；等待占用降到新限额以下再放行。旧连接继续使用已经选定的上游目标。 +- 热更新要求 `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` 可覆盖。 + +**首次从不支持 supervisor 协议的旧版（包括 SQLite 版本）迁入需要一次普通重启。** deb 升级会对正在运行的服务调用 `upgrade`，失败不会自动强制重启；首次迁入或不兼容升级时，检查错误后手动 `systemctl restart rsync-proxy`，再完成包配置。进程崩溃、主机重启、主动停止和配置的连接超时不属于零停机保证。 + +开发验证：`go test -race ./...` 包含真实子进程的多代交接、supervisor IPC 并发调度及故障测试。真实用户级 systemd 验证可运行 `python3 test/systemd/upgrade.py --binary /absolute/path/to/rsync-proxy`；该脚本只创建临时测试 unit、回环监听和临时目录。 + ### 安装 fail2ban filter ```shell @@ -90,7 +116,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 +135,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..6b356f8 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" +# Supervisor Unix socket and process lock directory (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..e1c47e5 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=main +KillMode=mixed 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..1fa5146 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, @@ -353,14 +356,7 @@ func New() *cobra.Command { } log.SetOutput(cmd.ErrOrStderr()) - if err := s.ReadConfigFromFile(true); err != nil { - 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 +368,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..6a1a486 --- /dev/null +++ b/cmd/daemon_linux.go @@ -0,0 +1,405 @@ +//go:build linux + +package cmd + +import ( + "context" + "crypto/rand" + "encoding/hex" + "encoding/json" + "errors" + "fmt" + "io" + "log" + "net" + "net/http" + "os" + "os/signal" + "path/filepath" + "strconv" + "sync" + "sync/atomic" + "syscall" + "time" + + "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 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 + drainMu sync.Mutex + drainTimer *time.Timer + stopAdmission func() +} + +func runWorker(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 + } + store, err := state.Open(dir) + if err != nil { + return err + } + defer store.Close() + 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), 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 { + return errors.New("worker requires supervisor handoff") + } + + 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 + } + + 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" { + timeout, err := time.ParseDuration(r.URL.Query().Get("timeout")) + if err != nil { + timeout = 0 + } + d.beginDrain(timeout) + 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 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") + } + + 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) + signal.Notify(signals, syscall.SIGTERM, syscall.SIGINT) + defer signal.Stop(signals) + var runErr error + select { + case runErr = <-d.runResult: + case <-store.Done(): + s.Close() + <-d.runResult + runErr = errors.New("supervisor connection lost") + case <-signals: + 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() + } + 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 { + _ = s.HTTPServer.Shutdown(ctx) + } + if err = s.PublishSnapshot(); err != nil { + log.Printf("[ERROR] final snapshot: %v", err) + } + + 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) + 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) + } + } + } +} + +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.store.Upgrade(d.id, timeout); err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + } else { + _, _ = fmt.Fprintln(w, `{"message":"Upgrade completed"}`) + } + } + return true +} diff --git a/cmd/daemon_linux_test.go b/cmd/daemon_linux_test.go new file mode 100644 index 0000000..b729b72 --- /dev/null +++ b/cmd/daemon_linux_test.go @@ -0,0 +1,743 @@ +//go:build linux + +package cmd + +import ( + "bufio" + "bytes" + "context" + "crypto/ecdsa" + "crypto/elliptic" + "crypto/rand" + "crypto/tls" + "crypto/x509" + "crypto/x509/pkix" + "encoding/json" + "encoding/pem" + "errors" + "fmt" + "io" + "math/big" + "net" + "net/http" + "os" + "os/exec" + "path/filepath" + "strconv" + "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) + 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) +} + +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) +} + +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 new file mode 100644 index 0000000..76db529 --- /dev/null +++ b/cmd/daemon_other.go @@ -0,0 +1,18 @@ +//go:build !linux + +package cmd + +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 + } + return s.Run() +} 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/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..fcc9e61 100644 --- a/go.mod +++ b/go.mod @@ -7,8 +7,11 @@ 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 ) require ( @@ -27,12 +30,9 @@ require ( 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/spf13/pflag v1.0.10 // indirect go.yaml.in/yaml/v2 v2.4.2 // indirect - golang.org/x/sys v0.35.0 // indirect google.golang.org/protobuf v1.36.8 // indirect gopkg.in/yaml.v3 v3.0.1 // indirect ) diff --git a/go.sum b/go.sum index d09b71d..551af35 100644 --- a/go.sum +++ b/go.sum @@ -68,8 +68,8 @@ 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/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= 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= diff --git a/pkg/server/admission.go b/pkg/server/admission.go new file mode 100644 index 0000000..af48226 --- /dev/null +++ b/pkg/server/admission.go @@ -0,0 +1,166 @@ +package server + +import ( + "context" + "errors" + "fmt" + "log" + "sync" + + "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 + wake chan struct{} + done chan struct{} +} + +func (s *Server) StartAdmission() func() { + 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) + for { + select { + case <-ctx.Done(): + return + case <-events: + a.poll() + case <-a.wake: + a.poll() + } + } + }() + return func() { cancel(); <-watchDone; <-a.done } +} + +func offerStatus(c chan queue.Status, status queue.Status) { + select { + case <-c: + default: + } + c <- status +} + +func (a *admission) poll() { + a.mu.Lock() + hasWaiting := len(a.waiting) > 0 + a.mu.Unlock() + 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() + 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) + } +} + +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() + 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/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..d152aad --- /dev/null +++ b/pkg/server/snapshot_test.go @@ -0,0 +1,58 @@ +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 := state.New() + 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/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/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..1b5869b --- /dev/null +++ b/pkg/state/store.go @@ -0,0 +1,431 @@ +// Package state coordinates worker generations through the supervisor's memory. +package state + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "sort" + "sync" + "time" +) + +const Protocol = 2 + +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 +} + +// 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 { + 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 +} + +type entry struct { + Ticket + Generation, Upstream, IP string +} + +func New() *Store { + return &Store{generations: make(map[string]Generation), limits: make(map[string]Limit), tickets: make(map[int64]entry), changed: make(chan struct{})} +} + +func (s *Store) Close() error { + if s.remote != nil { + return s.remote.close() + } + return nil +} + +func (s *Store) changedLocked() { + s.revision++ + close(s.changed) + s.changed = make(chan struct{}) +} + +func (s *Store) Current() (string, error) { + 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 validateLimits(limits []Limit) error { + for _, l := range limits { + if l.Name == "" || l.Active < 0 || l.Queued < 0 || l.PerIP < 0 { + return errors.New("invalid admission limit") + } + } + return nil +} + +func (s *Store) Activate(id, parent string, limits []Limit) error { + 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 { + 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 (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) promoteLocked() { + active := make(map[string]int) + for _, t := range s.tickets { + if t.Active { + active[t.Upstream]++ + } + } + 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]++ + } + } +} + +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 + } + if t.IP == ip { + ips++ + } + if t.Active { + active++ + } else { + queued++ + } + } + 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 { + 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 +} + +func (s *Store) Poll(generation string) (map[int64]Ticket, error) { + 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 _, 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: + } + } +} + +func (s *Store) Queues() ([]QueueInfo, error) { + 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) + } + sort.Slice(result, func(i, j int) bool { return result[i].Name < result[j].Name }) + return result, nil +} + +func (s *Store) Generations() ([]Generation, error) { + 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, nil +} + +func (s *Store) Publish(id string, snapshot []byte) error { + 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 { + 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 + } + 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 { + 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) + } + } + 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 new file mode 100644 index 0000000..cb27f82 --- /dev/null +++ b/pkg/state/store_test.go @@ -0,0 +1,311 @@ +package state + +import ( + "bytes" + "context" + "errors" + "fmt" + "net" + "net/rpc" + "os" + "os/exec" + "path/filepath" + "sort" + "sync" + "testing" + "time" + + "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) + 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() + owner := startTestCoordinator(t, dir) + 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, owner.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 := 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") + 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 TestProtocolRejectionAndCancellation(t *testing.T) { + dir := t.TempDir() + startTestCoordinator(t, dir) + conn, err := net.Dial("unix", filepath.Join(dir, "supervisor.sock")) + require.NoError(t, err) + 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, a.Activate("a", "", []Limit{{Name: "u", Active: 1}})) + _, err := a.Acquire(context.Background(), "a", "u", "ip") + require.NoError(t, err) + queued, err := b.Acquire(context.Background(), "b", "u", "other") + require.NoError(t, err) + require.False(t, queued.Active) + require.NoError(t, a.Close()) + q, err := b.Queues() + require.NoError(t, err) + 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) + 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) { + dir := os.Getenv("RSYNC_STATE_TEST_DIR") + if dir == "" { + return + } + id := fmt.Sprint(os.Getpid()) + s, err := Open(dir) + if err != nil { + panic(fmt.Errorf("open state store: %w", err)) + } + defer s.Close() + lock, err := Lock(filepath.Join(dir, id+".lock")) + if err != nil { + panic(fmt.Errorf("lock generation: %w", err)) + } + defer lock.Close() + if err = s.Register(Generation{ID: id, PID: os.Getpid()}); err != nil { + panic(fmt.Errorf("register generation: %w", 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(fmt.Errorf("acquire ticket: %w", err)) + } + for !ticket.Active { + time.Sleep(time.Millisecond) + ts, err := s.Poll(id) + if err != nil { + 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(fmt.Errorf("release ticket %d: %w", ticket.ID, err)) + } + } +} + +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() + require.NoError(t, err) + 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()) + workers = append(workers, worker{cmd: c, output: output}) + } + var wg sync.WaitGroup + 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) }() + ticker := time.NewTicker(time.Millisecond) + defer ticker.Stop() + for { + select { + case <-done: + 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: + qs, err := s.Queues() + require.NoError(t, err) + require.LessOrEqual(t, qs[0].ActiveCount, 2) + } + } +} + +func TestAdmissionAt100ConnectionsPerSecond(t *testing.T) { + 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) + 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() + startTestCoordinator(t, dir) + 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..a2ba675 --- /dev/null +++ b/test/systemd/upgrade.py @@ -0,0 +1,178 @@ +#!/usr/bin/env python3 +"""Exercise worker upgrades with a stable supervisor MAINPID 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=main", + "--property=KillMode=mixed", "--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 (not pids or pid == pids[0]), 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() + # 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() + for connection in connections: + connection.close() + subprocess.run(["systemctl", "--user", "stop", unit], check=False) + listener.close() + + +if __name__ == "__main__": + main()