diff --git a/cmd/gocached/gocached.go b/cmd/gocached/gocached.go new file mode 100644 index 0000000..abbb29c --- /dev/null +++ b/cmd/gocached/gocached.go @@ -0,0 +1,147 @@ +// Copyright (c) Tailscale Inc & AUTHORS +// SPDX-License-Identifier: BSD-3-Clause + +// The gocached command is a small example binary showing how to use +// the gocached library package. +package main + +import ( + "context" + "flag" + "fmt" + "log" + "net" + "net/http" + "os" + "strings" + "time" + + "github.com/bradfitz/parentdeath" + "github.com/tailscale/tb/gocache/gocached" + _ "modernc.org/sqlite" +) + +var ( + dir = flag.String("cache-dir", "", "cache directory, if empty defaults to /gocached") + sqliteDir = flag.String("sqlite-dir", "", "directory for the SQLite metadata database; if empty, defaults to the cache directory") + hotDir = flag.String("hot-dir", "", "if non-empty, enable storage tiering with this directory as a fast tier (e.g. local NVMe) holding a bounded copy of recently used blobs; the cache directory remains the source of truth") + hotCapacity = flag.Int("hot-capacity-gb", 600, "maximum size of the hot tier directory in GiB; only used with --hot-dir") + putSpoolCap = flag.Int("put-spool-gb", 0, "capacity in GiB of the put queue's spooled lane: accepted PUT bytes waiting on local disk (the hot tier's filesystem, with --hot-dir) for a copy into the cache directory; 0 means one eighth of that filesystem's size") + verbose = flag.Bool("verbose", false, "be verbose") + listen = flag.String("listen", ":31364", "listen address for the build-facing HTTP server") + debugListen = flag.String("debug-listen", "", "if non-empty, listen address for the debug HTTP server (pprof, metrics, etc)") + + maxSize = flag.Int("max-size-gb", 50, "maximum size of the cache in GiB; 0 means no limit") + maxAge = flag.Int("max-age-days", 60, "maximum age of objects in the cache in days; 0 means no limit") + + globalGeneration = flag.Int("global-generation", 1, "generation number of the global namespace; incrementing it effectively wipes the global namespace, with the old generation's objects aging out via normal LRU eviction") + + shardPrefixLen = flag.Int("shard-prefix-len", 2, "number of SHA256 hex characters per usage-stats shard key (valid 1..4); total shards = 16^n. Defaults to 2 (256 shards). Increase if a single shard's stats scan becomes too slow on a large DB") + + jwtIssuer = flag.String("jwt-issuer", "", "the issuer to trust JWTs from; if set, all requests will require auth, and must set at least one -jwt-claim") + // See example GitHub token claims for what can be available: + // https://docs.github.com/en/actions/concepts/security/openid-connect + jwtClaims = make(jwtClaimValue) + globalJWTClaims = make(jwtClaimValue) +) + +type jwtClaimValue map[string]string + +func (v jwtClaimValue) String() string { + return fmt.Sprintf("%v", map[string]string(v)) +} + +func (v jwtClaimValue) Set(s string) error { + claim, value, ok := strings.Cut(s, "=") + if !ok || claim == "" || value == "" { + return fmt.Errorf("bad claim %q, want x=y", s) + } + v[claim] = value + return nil +} + +func main() { + flag.Var(&jwtClaims, "jwt-claim", "a claim in the form x=y that any JWT presented must have to start a session; may be specified more than once") + flag.Var(&globalJWTClaims, "global-jwt-claim", "an additional claim in the form x=y that a JWT must have to allow writing to the cache's global namespace; may be specified more than once") + flag.Parse() + + parentdeath.Monitor(func() { + log.Printf("gocached: parent process died, exiting") + os.Exit(0) + }) + + opts := []gocached.ServerOption{ + gocached.WithDir(*dir), + gocached.WithSQLiteDir(*sqliteDir), + gocached.WithVerbose(*verbose), + gocached.WithMaxSize(int64(*maxSize) << 30), + gocached.WithMaxAge(time.Duration(*maxAge) * 24 * time.Hour), + gocached.WithShardPrefixLen(*shardPrefixLen), + gocached.WithGlobalGeneration(*globalGeneration), + } + + if *hotDir != "" { + if *hotCapacity <= 0 { + log.Fatal("--hot-capacity-gb must be positive when --hot-dir is set") + } + opts = append(opts, + gocached.WithHotDir(*hotDir), + gocached.WithHotCapacity(int64(*hotCapacity)<<30), + ) + } + if *putSpoolCap > 0 { + opts = append(opts, gocached.WithPutSpoolCapacity(int64(*putSpoolCap)<<30)) + } + + if *jwtIssuer != "" { + if len(jwtClaims) == 0 { + log.Fatal("must specify --jwt-claim at least once when --jwt-issuer is set") + } + + opts = append(opts, + gocached.WithJWTAuth(*jwtIssuer), + gocached.WithNamespaceMapping(func(ctx context.Context, claims map[string]any) (*gocached.NamespaceGrant, error) { + var ns gocached.Namespace + for k, want := range jwtClaims { + if got := claims[k]; got != want { + return nil, fmt.Errorf("claim %q = %v, want %v", k, got, want) + } + if ns != "" { + ns += "," + } + ns += gocached.Namespace(fmt.Sprintf("%s=%s", k, want)) + } + for k, want := range globalJWTClaims { + if got := claims[k]; got != want { + return &gocached.NamespaceGrant{ + WriteNamespace: &ns, + ExtraReadNamespace: &ns, + }, nil + } + } + return &gocached.NamespaceGrant{ + WriteNamespace: new(gocached.GlobalNamespace), + ExtraReadNamespace: &ns, + }, nil + }), + ) + } + + srv, err := gocached.NewServer(opts...) + if err != nil { + log.Fatalf("starting gocached: %v", err) + } + + if *debugListen != "" { + debugLn, err := net.Listen("tcp", *debugListen) + if err != nil { + log.Fatalf("debug listen: %v", err) + } + go func() { + log.Fatal(http.Serve(debugLn, http.HandlerFunc(srv.ServeHTTPDebug))) + }() + } + + log.Printf("gocached: listening on %s ...", *listen) + log.Fatal(http.ListenAndServe(*listen, srv)) +} diff --git a/cmd/gocached/gocached_test.go b/cmd/gocached/gocached_test.go new file mode 100644 index 0000000..30dbde8 --- /dev/null +++ b/cmd/gocached/gocached_test.go @@ -0,0 +1,59 @@ +// Copyright (c) Tailscale Inc & AUTHORS +// SPDX-License-Identifier: BSD-3-Clause + +package main + +import ( + "testing" + + "github.com/google/go-cmp/cmp" +) + +func TestJWTClaimFlag(t *testing.T) { + for name, tc := range map[string]struct { + input []string + want jwtClaimValue + }{ + "single_claim": { + input: []string{"role=builder"}, + want: jwtClaimValue{"role": "builder"}, + }, + "multiple_claims": { + input: []string{"env=prod", "team=devops"}, + want: jwtClaimValue{"env": "prod", "team": "devops"}, + }, + "duplicate_keys": { + input: []string{"key=value1", "key=value2"}, + want: jwtClaimValue{"key": "value2"}, + }, + "value_with_equals": { + input: []string{"data=a=b=c"}, + want: jwtClaimValue{"data": "a=b=c"}, + }, + } { + t.Run(name, func(t *testing.T) { + claim := make(jwtClaimValue) + for _, v := range tc.input { + if err := claim.Set(v); err != nil { + t.Fatalf("Set failed: %v", err) + } + } + if diff := cmp.Diff(claim, tc.want); diff != "" { + t.Errorf("jwtClaimValue mismatch (-got +want):\n%s", diff) + } + }) + } +} + +func TestInvalidJWTClaimFlags(t *testing.T) { + for _, s := range []string{ + "invalidclaim", + "=nokey", + "novalue=", + } { + claim := make(jwtClaimValue) + if err := claim.Set(s); err == nil { + t.Fatalf("Set with invalid claim %q did not return error", s) + } + } +} diff --git a/go.mod b/go.mod index 3d1a87f..a0e51da 100644 --- a/go.mod +++ b/go.mod @@ -3,23 +3,40 @@ module github.com/tailscale/tb go 1.27.1 require ( + github.com/bradfitz/parentdeath v0.0.0-20260315043412-764506aeb900 github.com/cespare/xxhash/v2 v2.3.0 + github.com/go-jose/go-jose/v4 v4.1.3 + github.com/golang-jwt/jwt/v5 v5.3.1 + github.com/google/go-cmp v0.7.0 + github.com/pierrec/lz4/v4 v4.1.26 + github.com/prometheus/client_golang v1.24.1 + github.com/prometheus/client_model v0.6.3 + golang.org/x/sync v0.23.0 + modernc.org/sqlite v1.51.0 tailscale.com v1.103.0-pre.0.20260920011920-7d96cf5a62ef ) require ( filippo.io/edwards25519 v1.2.0 // indirect github.com/akutz/memconn v0.1.0 // indirect + github.com/beorn7/perks v1.0.1 // indirect github.com/coder/websocket v1.8.15 // indirect github.com/dblohm7/wingoes v0.0.0-20260526185140-fb298caac7ca // indirect + github.com/dustin/go-humanize v1.0.1 // indirect github.com/fxamacker/cbor/v2 v2.9.3 // indirect github.com/go-json-experiment/json v0.0.0-20260820222146-c27c302e5fc3 // indirect - github.com/google/go-cmp v0.7.0 // indirect + github.com/google/uuid v1.6.0 // indirect github.com/hdevalence/ed25519consensus v0.2.0 // indirect github.com/jsimonetti/rtnetlink v1.4.2 // indirect + github.com/mattn/go-isatty v0.0.24 // indirect github.com/mdlayher/netlink v1.11.2 // indirect github.com/mdlayher/socket v0.7.0 // indirect github.com/mitchellh/go-ps v1.0.0 // indirect + github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect + github.com/ncruces/go-strftime v1.0.0 // indirect + github.com/prometheus/common v0.71.0 // indirect + github.com/prometheus/procfs v0.21.1 // indirect + github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect github.com/tailscale/go-winio v0.0.0-20231025203758-c4f33415bf55 // indirect github.com/x448/float16 v0.8.4 // indirect go4.org/mem v0.0.0-20240501181205-ae6ca9944745 // indirect @@ -27,7 +44,10 @@ require ( golang.org/x/crypto v0.57.0 // indirect golang.org/x/exp v0.0.0-20260908205506-85c1c2202aba // indirect golang.org/x/net v0.59.0 // indirect - golang.org/x/sync v0.23.0 // indirect golang.org/x/sys v0.48.0 // indirect golang.zx2c4.com/wireguard/windows v1.0.1 // indirect + google.golang.org/protobuf v1.36.12 // indirect + modernc.org/libc v1.72.3 // indirect + modernc.org/mathutil v1.7.1 // indirect + modernc.org/memory v1.11.0 // indirect ) diff --git a/go.sum b/go.sum index 41f69f0..b0cc92a 100644 --- a/go.sum +++ b/go.sum @@ -34,6 +34,10 @@ github.com/aws/aws-sdk-go-v2/service/sts v1.49.0 h1:N7Ey8obY3uSui+cxl0OUzFlFmkxS github.com/aws/aws-sdk-go-v2/service/sts v1.49.0/go.mod h1:zMBwjSf4Pt8a1OHYiZ5rPD0PJRK1kQrUaxhS/Dbld8E= github.com/aws/smithy-go v1.28.1 h1:R/nXH00c8qcfCzQVELtRw+eLQWtzv+VAIEFJ1/xxXlQ= github.com/aws/smithy-go v1.28.1/go.mod h1:YE2RhdIuDbA5E5bTdciG9KrW3+TiEONeUWCqxX9i1Fc= +github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM= +github.com/beorn7/perks v1.0.1/go.mod h1:G2ZrVWU2WbWT9wwq4/hrbKbnv/1ERSJQ0ibhJ6rlkpw= +github.com/bradfitz/parentdeath v0.0.0-20260315043412-764506aeb900 h1:YTPrKVaBJlca4xkN+YU2VG3WHIDFme0vItEKuj0wrBQ= +github.com/bradfitz/parentdeath v0.0.0-20260315043412-764506aeb900/go.mod h1:nmAuQ8iUcGbqQsa847JnlFq6PQMMxpEeK0QCPAPYf3w= github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs= github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= github.com/cilium/ebpf v0.22.0 h1:v2ktp0roffpMOj2MMf3idtCQZOsAoC4BJbAJN+ke2bY= @@ -50,18 +54,24 @@ github.com/dblohm7/wingoes v0.0.0-20260526185140-fb298caac7ca h1:h1Awca4lQOspNR/ github.com/dblohm7/wingoes v0.0.0-20260526185140-fb298caac7ca/go.mod h1:2TGl1jRJrRpbzykmg7asHm3h08TqutUgQqY5v9k/g3c= github.com/digitalocean/go-smbios v0.0.0-20180907143718-390a4f403a8e h1:vUmf0yezR0y7jJ5pceLHthLaYf4bA5T14B6q39S4q2Q= github.com/digitalocean/go-smbios v0.0.0-20180907143718-390a4f403a8e/go.mod h1:YTIHhz/QFSYnu/EhlF2SpU2Uk+32abacUYA5ZPljz1A= +github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY= +github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto= github.com/frankban/quicktest v1.14.6 h1:7Xjx+VpznH+oBnejlPUj8oUpdxnVs4f8XU8WnHkI4W8= github.com/frankban/quicktest v1.14.6/go.mod h1:4ptaffx2x8+WTWXmUCuVU6aPUX1/Mz7zb5vbUoiM6w0= github.com/fxamacker/cbor/v2 v2.9.3 h1:oQBnFATpNdY8gJHTndDDv5Xl4QqNaz51G5LLEPhng3Q= github.com/fxamacker/cbor/v2 v2.9.3/go.mod h1:vM4b+DJCtHn+zz7h3FFp/hDAI9WNWCsZj23V5ytsSxQ= github.com/gaissmai/bart v0.29.0 h1:wO6HGE8g9YE0Wm0bCpYxwRzfQ4+fbJKOhL64e5ACGCI= github.com/gaissmai/bart v0.29.0/go.mod h1:GREWQfTLRWz/c5FTOsIw+KkscuFkIV5t8Rp7Nd1Td5c= +github.com/go-jose/go-jose/v4 v4.1.3 h1:CVLmWDhDVRa6Mi/IgCgaopNosCaHz7zrMeF9MlZRkrs= +github.com/go-jose/go-jose/v4 v4.1.3/go.mod h1:x4oUasVrzR7071A4TnHLGSPpNOm2a21K9Kf04k1rs08= github.com/go-json-experiment/json v0.0.0-20260820222146-c27c302e5fc3 h1:UADEEmDKgfXbtnGJZ97beY5XLo9ZechG1nlU4KnRrkE= github.com/go-json-experiment/json v0.0.0-20260820222146-c27c302e5fc3/go.mod h1:tphK2c80bpPhMOI4v6bIc2xWywPfbqi1Z06+RcrMkDg= github.com/go-ole/go-ole v1.3.0 h1:Dt6ye7+vXGIKZ7Xtk4s6/xVdGDQynvom7xCFEdWr6uE= github.com/go-ole/go-ole v1.3.0/go.mod h1:5LS6F96DhAwUc7C+1HLexzMXY1xGRSryjyPPKW6zv78= github.com/godbus/dbus/v5 v5.2.2 h1:TUR3TgtSVDmjiXOgAAyaZbYmIeP3DPkld3jgKGV8mXQ= github.com/godbus/dbus/v5 v5.2.2/go.mod h1:3AAv2+hPq5rdnr5txxxRwiGjPXamgoIHgz9FPBfOp3c= +github.com/golang-jwt/jwt/v5 v5.3.1 h1:kYf81DTWFe7t+1VvL7eS+jKFVWaUnK9cB1qbwn63YCY= +github.com/golang-jwt/jwt/v5 v5.3.1/go.mod h1:fxCRLWMO43lRc8nhHWY6LGqRcf+1gQWArsqaEUEa5bE= github.com/golang/groupcache v0.0.0-20241129210726-2c02b8208cf8 h1:f+oWsMOmNPc8JmEHVZIycC7hBoQxHH9pNKQORJNozsQ= github.com/golang/groupcache v0.0.0-20241129210726-2c02b8208cf8/go.mod h1:wcDNUvekVysuuOpQKo3191zZyTpiI6se1N1ULghS0sw= github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8= @@ -70,8 +80,13 @@ github.com/google/go-tpm v0.9.8 h1:slArAR9Ft+1ybZu0lBwpSmpwhRXaa85hWtMinMyRAWo= github.com/google/go-tpm v0.9.8/go.mod h1:h9jEsEECg7gtLis0upRBQU+GhYVH6jMjrFxI8u6bVUY= github.com/google/nftables v0.3.0 h1:bkyZ0cbpVeMHXOrtlFc8ISmfVqq5gPJukoYieyVmITg= github.com/google/nftables v0.3.0/go.mod h1:BCp9FsrbF1Fn/Yu6CLUc9GGZFw/+hsxfluNXXmxBfRM= +github.com/google/pprof v0.0.0-20250317173921-a4b03ec1a45e h1:ijClszYn+mADRFY17kjQEVQ1XRhq2/JR1M3sGqeJoxs= +github.com/google/pprof v0.0.0-20250317173921-a4b03ec1a45e/go.mod h1:boTsfXsheKC2y+lKOCMpSfarhxDeIzfZG1jqGcPl3cA= github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/hashicorp/golang-lru v1.0.2 h1:dV3g9Z/unq5DpblPpw+Oqcv4dU/1omnb4Ok8iPY6p1c= +github.com/hashicorp/golang-lru/v2 v2.0.7 h1:a+bsQ5rvGLjzHuww6tVxozPZFVghXaHOwFs4luLUK2k= +github.com/hashicorp/golang-lru/v2 v2.0.7/go.mod h1:QeFd9opnmA6QUJc5vARoKUSoFhyfM2/ZepoAG6RGpeM= github.com/hdevalence/ed25519consensus v0.2.0 h1:37ICyZqdyj0lAZ8P4D1d1id3HqbbG1N3iBb1Tb4rdcU= github.com/hdevalence/ed25519consensus v0.2.0/go.mod h1:w3BHWjwJbFU29IRHL1Iqkw3sus+7FctEyM4RqDxYNzo= github.com/huin/goupnp v1.3.0 h1:UvLUlWDNpoUdYzb2TCn+MuTWtcjXKSza2n6CBdQ0xXc= @@ -90,6 +105,10 @@ github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE= github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk= github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY= github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE= +github.com/kylelemons/godebug v1.1.0 h1:RPNrshWIDI6G2gRW9EHilWtl7Z6Sb1BR0xunSBf0SNc= +github.com/kylelemons/godebug v1.1.0/go.mod h1:9/0rRGxNHcop5bhtWyNeEfOS8JIWk580+fNqagV/RAw= +github.com/mattn/go-isatty v0.0.24 h1:tGZZoVgT/KiqK1c8ocVLeDS8BSWMRd47J3Lbz7vsReI= +github.com/mattn/go-isatty v0.0.24/go.mod h1:nMCL3Zebbrt45jsMDgnfIwz6ydEQApk5oEI3HqDio6A= github.com/mdlayher/genetlink v1.4.0 h1:f/Xs7Y2T+GyX9b3dbiUhnLE9InGs5F9RxJ2JwBMl71o= github.com/mdlayher/genetlink v1.4.0/go.mod h1:d1hrKr8fwZU2JkcAtQUAzeTrI7nbgQSl+5k1cC0biSA= github.com/mdlayher/netlink v1.11.2 h1:HKh2jqe+omdSWcQ88nrT7INE61B0NXfiSPFdgL4YbNI= @@ -102,18 +121,28 @@ github.com/mitchellh/go-ps v1.0.0 h1:i6ampVEEF4wQFF+bkYfwYgY+F/uYJDktmvLPf7qIgjc github.com/mitchellh/go-ps v1.0.0/go.mod h1:J4lOc8z8yJs6vUwklHw2XEIiT4z4C40KtWVN3nvg8Pg= github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 h1:C3w9PqII01/Oq1c1nUAm88MOHcQC9l5mIlSMApZMrHA= github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822/go.mod h1:+n7T8mK8HuQTcFwEeznm/DIxMOiR9yIdICNftLE1DvQ= +github.com/ncruces/go-strftime v1.0.0 h1:HMFp8mLCTPp341M/ZnA4qaf7ZlsbTc+miZjCLOFAw7w= +github.com/ncruces/go-strftime v1.0.0/go.mod h1:Fwc5htZGVVkseilnfgOVb9mKy6w1naJmn9CehxcKcls= github.com/pierrec/lz4/v4 v4.1.26 h1:GrpZw1gZttORinvzBdXPUXATeqlJjqUG/D87TKMnhjY= github.com/pierrec/lz4/v4 v4.1.26/go.mod h1:EoQMVJgeeEOMsCqCzqFm2O0cJvljX2nGZjcRIPL34O4= github.com/pires/go-proxyproto v0.15.0 h1:dTshmNbFm/D+0+sbrxUuddPOZ5Y0B7c5NhtsBkm6LqI= github.com/pires/go-proxyproto v0.15.0/go.mod h1:OXsCrKwrK2tXS9YrI5tkHx5xaQlO8FH3lFW76orFh24= +github.com/prometheus/client_golang v1.24.1 h1:JnJkREXzWxUdCuPFpIWZiPispT9xVV59uiuyR2bPlnU= +github.com/prometheus/client_golang v1.24.1/go.mod h1:F+oSRECHg4sse5ucfYpYDeIv/hu68Zo0uoHKetWnzcE= github.com/prometheus/client_model v0.6.3 h1:O0jaTVAYNxTHYInEPFJt5I3+sN8zqBtVMPTB1qyxiEo= github.com/prometheus/client_model v0.6.3/go.mod h1:gpN5P9S7Rr6Yr92PiQ+Ixvhf6JZEkF1dnxsYL2aPBEM= github.com/prometheus/common v0.71.0 h1:9KDAKb7Mj3HEVKyFCK6Dc/HIwlBzZIN2l7/lrHl3KK8= github.com/prometheus/common v0.71.0/go.mod h1:CLJ5H8TEsGX8bl31BdMkfhIZ+QmZ9tBPPotUxUbfcmk= +github.com/prometheus/procfs v0.21.1 h1:GljZCt+zSTS+NZq88cyQ1LjZ+RCHp3uVuabBWA5+OJI= +github.com/prometheus/procfs v0.21.1/go.mod h1:aB55Cww9pdSJVHk0hUf0inxWyyjPogFIjmHKYgMKmtY= +github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec h1:W09IVJc94icq4NjY3clb7Lk8O1qJ8BdBEF8z0ibU0rE= +github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec/go.mod h1:qqbHyh8v60DhA7CoWK5oRCqLrMHRGoxYCSS9EjAz6Eo= github.com/rogpeppe/go-internal v1.16.0 h1:O9DK+vNMDVGLr2BeZqmpLeMjiMNkuXfcqntWbZV6S5g= github.com/rogpeppe/go-internal v1.16.0/go.mod h1:DrUVZyrJU+txYW5/1kwtXQSMFio52ZOxX7yM1VHvnxs= github.com/safchain/ethtool v0.7.0 h1:rlJzfDetsVvT61uz8x1YIcFn12akMfuPulHtZjtb7Is= github.com/safchain/ethtool v0.7.0/go.mod h1:MenQKEjXdfkjD3mp2QdCk8B/hwvkrlOTm/FD4gTpFxQ= +github.com/stretchr/testify v1.12.1 h1:EuwCh5fleGS7H32xRwO3wRGT7DxrDhLAT6FF8MpWDWE= +github.com/stretchr/testify v1.12.1/go.mod h1:MDEgiDPPsNp5cuIrHPPCyornHKgEVbtFUmoNlxoYthg= github.com/tailscale/certstore v0.1.1-0.20260409135935-3638fb84b77d h1:JcGKBZAL7ePLwOhUdN8qGQZlP5GueEiIZwY7R62pejE= github.com/tailscale/certstore v0.1.1-0.20260409135935-3638fb84b77d/go.mod h1:XrBNfAFN+pwoWuksbFS9Ccxnopa15zJGgXRFN90l3K4= github.com/tailscale/go-winio v0.0.0-20231025203758-c4f33415bf55 h1:Gzfnfk2TWrk8Jj4P4c1a3CtQyMaTVCznlkLZI++hok4= @@ -134,6 +163,12 @@ github.com/vishvananda/netns v0.0.5 h1:DfiHV+j8bA32MFM7bfEunvT8IAqQ/NzSJHtcmW5zd github.com/vishvananda/netns v0.0.5/go.mod h1:SpkAiCQRtJ6TvvxPnOSyH3BMl6unz3xZlaprSwhNNJM= github.com/x448/float16 v0.8.4 h1:qLwI1I70+NjRFUR3zs1JPUCgaCXSh3SW62uAKT1mSBM= github.com/x448/float16 v0.8.4/go.mod h1:14CWIYCyZA/cWjXOioeEpHeN/83MdbZDRQHoFcYsOfg= +go.uber.org/goleak v1.3.0 h1:2K3zAYmnTNqV73imy9J1T3WC+gmCePx2hEGkimedGto= +go.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE= +go.yaml.in/yaml/v2 v2.4.4 h1:tuyd0P+2Ont/d6e2rl3be67goVK4R6deVxCUX5vyPaQ= +go.yaml.in/yaml/v2 v2.4.4/go.mod h1:gMZqIpDtDqOfM0uNfy0SkpRhvUryYH0Z6wdMYcacYXQ= +go.yaml.in/yaml/v3 v3.0.5 h1:N6y/pJk8buWs9NY5ERU2HSMfm+IuD/OtfdAnq6kESPw= +go.yaml.in/yaml/v3 v3.0.5/go.mod h1:HVTZu1O7/Vkt2N+BFy8Zza+lnLsABggaTM2ZpNIGuKg= go4.org/mem v0.0.0-20240501181205-ae6ca9944745 h1:Tl++JLUCe4sxGu8cTpDzRLd3tN7US4hOxG5YpKCzkek= go4.org/mem v0.0.0-20240501181205-ae6ca9944745/go.mod h1:reUoABIJ9ikfM5sgtSF3Wushcza7+WeD01VB9Lirh3g= go4.org/netipx v0.0.0-20260823151212-3075585bcbeb h1:XBM4hvfwGAttkkiTIFfeigdfcL1xIfdKXqFdgiHGtDs= @@ -166,5 +201,33 @@ google.golang.org/protobuf v1.36.12 h1:pJOKDDOyeXErUroCihFAd5LQuwXBSpVnKGrj5o/fw google.golang.org/protobuf v1.36.12/go.mod h1:HTf+CrKn2C3g5S8VImy6tdcUvCska2kB7j23XfzDpco= gvisor.dev/gvisor v0.0.0-20260915211658-a6f909f08a72 h1:EytKYr5WrVs+Aah/0xDO8ZAZZzS/iTEVC8ln4ipDbL0= gvisor.dev/gvisor v0.0.0-20260915211658-a6f909f08a72/go.mod h1:8aLQqUBHDH8fY5y60lzmwDpMMbQCcT3EBfoSwhfaGCY= +modernc.org/cc/v4 v4.28.2 h1:3tQ0lf2ADtoby2EtSP+J7IE2SHwEJdP8ioR59wx7XpY= +modernc.org/cc/v4 v4.28.2/go.mod h1:OnovgIhbbMXMu1aISnJ0wvVD1KnW+cAUJkIrAWh+kVI= +modernc.org/ccgo/v4 v4.34.0 h1:yRLPFZieg532OT4rp4JFNIVcquwalMX26G95WQDqwCQ= +modernc.org/ccgo/v4 v4.34.0/go.mod h1:AS5WYMyBakQ+fhsHhtP8mWB82KTGPkNNJDGfGQCe0/A= +modernc.org/fileutil v1.4.0 h1:j6ZzNTftVS054gi281TyLjHPp6CPHr2KCxEXjEbD6SM= +modernc.org/fileutil v1.4.0/go.mod h1:EqdKFDxiByqxLk8ozOxObDSfcVOv/54xDs/DUHdvCUU= +modernc.org/gc/v2 v2.6.5 h1:nyqdV8q46KvTpZlsw66kWqwXRHdjIlJOhG6kxiV/9xI= +modernc.org/gc/v2 v2.6.5/go.mod h1:YgIahr1ypgfe7chRuJi2gD7DBQiKSLMPgBQe9oIiito= +modernc.org/gc/v3 v3.1.2 h1:ZtDCnhonXSZexk/AYsegNRV1lJGgaNZJuKjJSWKyEqo= +modernc.org/gc/v3 v3.1.2/go.mod h1:HFK/6AGESC7Ex+EZJhJ2Gni6cTaYpSMmU/cT9RmlfYY= +modernc.org/goabi0 v0.2.0 h1:HvEowk7LxcPd0eq6mVOAEMai46V+i7Jrj13t4AzuNks= +modernc.org/goabi0 v0.2.0/go.mod h1:CEFRnnJhKvWT1c1JTI3Avm+tgOWbkOu5oPA8eH8LnMI= +modernc.org/libc v1.72.3 h1:ZnDF4tXn4NBXFutMMQC4vtbTFSXhhKzR73fv0beZEAU= +modernc.org/libc v1.72.3/go.mod h1:dn0dZNnnn1clLyvRxLxYExxiKRZIRENOfqQ8XEeg4Qs= +modernc.org/mathutil v1.7.1 h1:GCZVGXdaN8gTqB1Mf/usp1Y/hSqgI2vAGGP4jZMCxOU= +modernc.org/mathutil v1.7.1/go.mod h1:4p5IwJITfppl0G4sUEDtCr4DthTaT47/N3aT6MhfgJg= +modernc.org/memory v1.11.0 h1:o4QC8aMQzmcwCK3t3Ux/ZHmwFPzE6hf2Y5LbkRs+hbI= +modernc.org/memory v1.11.0/go.mod h1:/JP4VbVC+K5sU2wZi9bHoq2MAkCnrt2r98UGeSK7Mjw= +modernc.org/opt v0.2.0 h1:tGyef5ApycA7FSEOMraay9SaTk5zmbx7Tu+cJs4QKZg= +modernc.org/opt v0.2.0/go.mod h1:03fq9lsNfvkYSfxrfUhZCWPk1lm4cq4N+Bh//bEtgns= +modernc.org/sortutil v1.2.1 h1:+xyoGf15mM3NMlPDnFqrteY07klSFxLElE2PVuWIJ7w= +modernc.org/sortutil v1.2.1/go.mod h1:7ZI3a3REbai7gzCLcotuw9AC4VZVpYMjDzETGsSMqJE= +modernc.org/sqlite v1.51.0 h1:aH/MMSoayAIhozZ7uJbVTT9QO/VhzBf0J9tymmmuC/U= +modernc.org/sqlite v1.51.0/go.mod h1:tcNzv5p84E0skkmJn038y+hWJbLQXQqEnQfeh5r2JLM= +modernc.org/strutil v1.2.1 h1:UneZBkQA+DX2Rp35KcM69cSsNES9ly8mQWD71HKlOA0= +modernc.org/strutil v1.2.1/go.mod h1:EHkiggD70koQxjVdSBM3JKM7k6L0FbGE5eymy9i3B9A= +modernc.org/token v1.1.0 h1:Xl7Ap9dKaEs5kLoOQeQmPWevfnk/DM5qcLcYlA8ys6Y= +modernc.org/token v1.1.0/go.mod h1:UGzOrNV1mAFSEB63lOFHIpNRUVMvYTc6yu1SMY/XTDM= tailscale.com v1.103.0-pre.0.20260920011920-7d96cf5a62ef h1:N5JizbSj4+lLBUSwIc89JAjDnxYbxEKyNhIhNwBvBHo= tailscale.com v1.103.0-pre.0.20260920011920-7d96cf5a62ef/go.mod h1:oUmJ/nA/pTiEi4J7pOaPdEL6Bxi+OhgIXGkz1KdPAjA= diff --git a/gocache/.gitignore b/gocache/.gitignore new file mode 100644 index 0000000..b25c15b --- /dev/null +++ b/gocache/.gitignore @@ -0,0 +1 @@ +*~ diff --git a/gocache/README.md b/gocache/README.md new file mode 100644 index 0000000..2034a80 --- /dev/null +++ b/gocache/README.md @@ -0,0 +1,45 @@ +# go-tool-cache + +Do you like Go's built-in build & test caching but wish it weren't purely stored on local disk in the `$GOCACHE` directory? + +Want to share your cache over the network between your various machines, coworkers, and CI runs without all that GitHub actions/caches tarring and untarring? + +Go's [GOCACHEPROG](https://pkg.go.dev/cmd/go/internal/cacheprog) lets you do that! + +This was a demonstration repro for when GOCACHEPROG was still a +[proposal](https://github.com/golang/go/issues/59719). Now it just contains some misc +examples. + +## Status + +GOCACHEPROG shipped as an experiment in Go 1.21. It became official in Go 1.24. + +## Using + +First, build your cache child process. For example, + +```sh +$ go install github.com/bradfitz/go-tool-cache/cmd/go-cacher@latest +``` + +Then tell Go to use it: + +```sh +$ GOCACHEPROG=$HOME/go/bin/go-cacher go install std +``` + +See some stats: + +```sh +$ GOCACHEPROG="$HOME/go/bin/go-cacher --verbose" go install std +Defaulting to cache dir /home/bradfitz/.cache/go-cacher ... +cacher: closing; 548 gets (0 hits, 548 misses, 0 errors); 1090 puts (0 errors) +``` + +Run it again and watch the hit rate go up: + +```sh +$ GOCACHEPROG="$HOME/go/bin/go-cacher --verbose" go install std +Defaulting to cache dir /home/bradfitz/.cache/go-cacher ... +cacher: closing; 808 gets (808 hits, 0 misses, 0 errors); 0 puts (0 errors) +``` diff --git a/gocache/cacheproc/cacheproc.go b/gocache/cacheproc/cacheproc.go new file mode 100644 index 0000000..ed53e70 --- /dev/null +++ b/gocache/cacheproc/cacheproc.go @@ -0,0 +1,222 @@ +// Copyright 2023 The Go Authors. All rights reserved. +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +// Package cacheproc implements the mechanics of talking to cmd/go's GOCACHE protocol +// so you can write a caching child process at a higher level. +package cacheproc + +import ( + "bufio" + "bytes" + "context" + "encoding/hex" + "encoding/json" + "errors" + "fmt" + "io" + "log" + "os" + "sync" + "sync/atomic" + + "github.com/tailscale/tb/gocache/wire" +) + +// Process implements the cmd/go JSON protocol over stdin & stdout via three +// funcs that callers can optionally implement. +type Process struct { + // Get optionally specifies a func to look up something from the cache. If + // nil, all gets are treated as cache misses.touch + // + // The actionID is a lowercase hex string of unspecified format or length. + // + // The returned outputID must be the same outputID provided to Put earlier; + // it will be a lowercase hex string of unspecified hash function or length. + // + // On cache miss, return all zero values (no error). On cache hit, diskPath + // must be the absolute path to a regular file; its size and modtime are + // returned to cmd/go. + // + // If the returned diskPath doesn't exist, it's treated as a cache miss. + Get func(ctx context.Context, actionID string) (outputID, diskPath string, _ error) + + // Put optionally specifies a func to add something to the cache. + // The actionID and outputID is a lowercase hex string of unspecified format or length. + // On success, diskPath must be the absolute path to a regular file. + // If nil, cmd/go may write to disk elsewhere as needed. + Put func(ctx context.Context, actionID, outputID string, size int64, r io.Reader) (diskPath string, _ error) + + // Close optionally specifies a func to run when the cmd/go tool is + // shutting down. + Close func() error + + Gets atomic.Int64 + GetHits atomic.Int64 + GetMisses atomic.Int64 + GetErrors atomic.Int64 + Puts atomic.Int64 + PutErrors atomic.Int64 +} + +func (p *Process) Run() error { + br := bufio.NewReader(os.Stdin) + jd := json.NewDecoder(br) + + bw := bufio.NewWriter(os.Stdout) + je := json.NewEncoder(bw) + + var caps []wire.Cmd + if p.Get != nil { + caps = append(caps, "get") + } + if p.Put != nil { + caps = append(caps, "put") + } + if p.Close != nil { + caps = append(caps, "close") + } + je.Encode(&wire.Response{KnownCommands: caps}) + if err := bw.Flush(); err != nil { + return err + } + + var wmu sync.Mutex // guards writing responses + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + for { + var req wire.Request + if err := jd.Decode(&req); err != nil { + if errors.Is(err, io.EOF) { + return nil + } + return err + } + // For Go1.23 backward compatibility, remove in Go1.25. + if len(req.OutputID) == 0 && len(req.ObjectID) != 0 { + req.OutputID = req.ObjectID + } + if req.Command == wire.CmdPut && req.BodySize > 0 { + // TODO(bradfitz): stream this and pass a checksum-validating + // io.Reader that validates on EOF. + var bodyb []byte + if err := jd.Decode(&bodyb); err != nil { + log.Fatal(err) + } + if int64(len(bodyb)) != req.BodySize { + log.Fatalf("only got %d bytes of declared %d", len(bodyb), req.BodySize) + } + req.Body = bytes.NewReader(bodyb) + } + go func() { + res := &wire.Response{ID: req.ID} + ctx := ctx // TODO: include req ID as a context.Value for tracing? + if err := p.handleRequest(ctx, &req, res); err != nil { + res.Err = err.Error() + } + wmu.Lock() + defer wmu.Unlock() + je.Encode(res) + bw.Flush() + }() + } +} + +func (p *Process) handleRequest(ctx context.Context, req *wire.Request, res *wire.Response) error { + switch req.Command { + default: + return errors.New("unknown command") + case "close": + if p.Close != nil { + return p.Close() + } + return nil + case "get": + return p.handleGet(ctx, req, res) + case "put": + return p.handlePut(ctx, req, res) + } +} + +func (p *Process) handleGet(ctx context.Context, req *wire.Request, res *wire.Response) (retErr error) { + p.Gets.Add(1) + defer func() { + if retErr != nil { + p.GetErrors.Add(1) + } else if res.Miss { + p.GetMisses.Add(1) + } else { + p.GetHits.Add(1) + } + }() + if p.Get == nil { + res.Miss = true + return nil + } + outputID, diskPath, err := p.Get(ctx, fmt.Sprintf("%x", req.ActionID)) + if err != nil { + return err + } + if outputID == "" && diskPath == "" { + res.Miss = true + return nil + } + if outputID == "" { + return errors.New("no outputID") + } + res.OutputID, err = hex.DecodeString(outputID) + if err != nil { + return fmt.Errorf("invalid OutputID: %v", err) + } + fi, err := os.Stat(diskPath) + if err != nil { + if os.IsNotExist(err) { + res.Miss = true + return nil + } + return err + } + if !fi.Mode().IsRegular() { + return fmt.Errorf("not a regular file") + } + res.Size = fi.Size() + res.TimeNanos = fi.ModTime().UnixNano() + res.DiskPath = diskPath + return nil +} + +func (p *Process) handlePut(ctx context.Context, req *wire.Request, res *wire.Response) (retErr error) { + actionID, outputID := fmt.Sprintf("%x", req.ActionID), fmt.Sprintf("%x", req.OutputID) + p.Puts.Add(1) + defer func() { + if retErr != nil { + p.PutErrors.Add(1) + log.Printf("put(action %s, obj %s, %v bytes): %v", actionID, outputID, req.BodySize, retErr) + } + }() + if p.Put == nil { + if req.Body != nil { + io.Copy(io.Discard, req.Body) + } + return nil + } + var body io.Reader = req.Body + if body == nil { + body = bytes.NewReader(nil) + } + diskPath, err := p.Put(ctx, actionID, outputID, req.BodySize, body) + if err != nil { + return err + } + fi, err := os.Stat(diskPath) + if err != nil { + return fmt.Errorf("stat after successful Put: %w", err) + } + if fi.Size() != req.BodySize { + return fmt.Errorf("failed to write file to disk with right size: disk=%v; wanted=%v", fi.Size(), req.BodySize) + } + res.DiskPath = diskPath + return nil +} diff --git a/gocache/cachers/disk.go b/gocache/cachers/disk.go new file mode 100644 index 0000000..e8f5048 --- /dev/null +++ b/gocache/cachers/disk.go @@ -0,0 +1,138 @@ +package cachers + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io" + "log" + "os" + "path/filepath" + "time" +) + +// indexEntry is the metadata that DiskCache stores on disk for an ActionID. +type indexEntry struct { + Version int `json:"v"` + OutputID string `json:"o"` + Size int64 `json:"n"` + TimeNanos int64 `json:"t"` +} + +type DiskCache struct { + Dir string + Verbose bool + Logf func(format string, args ...any) // optional alt logger +} + +func (dc *DiskCache) logf(format string, args ...any) { + if dc.Logf != nil { + dc.Logf(format, args...) + } else if dc.Verbose { + log.Printf(format, args...) + } +} + +func (dc *DiskCache) Get(ctx context.Context, actionID string) (outputID, diskPath string, err error) { + if !validHex(actionID) { + return "", "", fmt.Errorf("actionID must be valid hex strings") + } + + actionFile := dc.ActionFilename(actionID) + ij, err := os.ReadFile(actionFile) + if err != nil { + if os.IsNotExist(err) { + err = nil + if dc.Verbose { + dc.logf("disk miss: %v", actionID) + } + } + return "", "", err + } + var ie indexEntry + if err := json.Unmarshal(ij, &ie); err != nil { + dc.logf("Warning: JSON error for action %q: %v", actionID, err) + return "", "", nil + } + if !validHex(ie.OutputID) { + // Protect against malicious non-hex OutputID on disk + return "", "", nil + } + return ie.OutputID, dc.OutputFilename(ie.OutputID), nil +} + +func (dc *DiskCache) OutputFilename(outputID string) string { + if !validHex(outputID) { + return "" + } + return filepath.Join(dc.Dir, outputID[:2], fmt.Sprintf("o-%s", outputID)) +} + +func (dc *DiskCache) ActionFilename(actionID string) string { + if !validHex(actionID) { + return "" + } + return filepath.Join(dc.Dir, actionID[:2], fmt.Sprintf("a-%s", actionID)) +} + +func validHex(x string) bool { + if len(x) < 4 || len(x) > 100 { + return false + } + for _, b := range x { + if b >= '0' && b <= '9' || b >= 'a' && b <= 'f' { + continue + } + return false + } + return true +} + +func (dc *DiskCache) Put(ctx context.Context, actionID, outputID string, size int64, body io.Reader) (diskPath string, _ error) { + if len(actionID) < 4 || len(outputID) < 4 { + return "", fmt.Errorf("actionID and outputID must be at least 4 characters long") + } + if !validHex(actionID) { + log.Printf("diskcache: got invalid actionID %q", actionID) + return "", errors.New("actionID must be hex") + } + if !validHex(outputID) { + log.Printf("diskcache: got invalid outputID %q", outputID) + return "", errors.New("outputID must be hex") + } + + actionFile := dc.ActionFilename(actionID) + outputFile := dc.OutputFilename(outputID) + actionDir := filepath.Dir(actionFile) + outputDir := filepath.Dir(outputFile) + + if err := os.MkdirAll(actionDir, 0755); err != nil { + return "", fmt.Errorf("failed to create action directory: %w", err) + } + if err := os.MkdirAll(outputDir, 0755); err != nil { + return "", fmt.Errorf("failed to create output directory: %w", err) + } + + wrote, err := writeOutputFile(outputFile, body, size, outputID) + if err != nil { + return "", err + } + if wrote != size { + return "", fmt.Errorf("wrote %d bytes, expected %d", wrote, size) + } + + ij, err := json.Marshal(indexEntry{ + Version: 1, + OutputID: outputID, + Size: size, + TimeNanos: time.Now().UnixNano(), + }) + if err != nil { + return "", err + } + if err := writeActionFile(actionFile, ij); err != nil { + return "", err + } + return outputFile, nil +} diff --git a/gocache/cachers/disk_notwindows.go b/gocache/cachers/disk_notwindows.go new file mode 100644 index 0000000..20b7721 --- /dev/null +++ b/gocache/cachers/disk_notwindows.go @@ -0,0 +1,49 @@ +//go:build !windows + +package cachers + +import ( + "bytes" + "io" + "os" + "path/filepath" +) + +func writeActionFile(dest string, b []byte) error { + _, err := writeAtomic(dest, bytes.NewReader(b), 0o644) + return err +} + +// writeOutputFile writes output files with the executable bit set +// because they can be linked binaries (e.g. cached test binaries) +// that cmd/go executes directly from the cache directory. +func writeOutputFile(dest string, r io.Reader, _ int64, _ string) (int64, error) { + return writeAtomic(dest, r, 0o755) +} + +func writeAtomic(dest string, r io.Reader, perm os.FileMode) (int64, error) { + tf, err := os.CreateTemp(filepath.Dir(dest), filepath.Base(dest)+".*") + if err != nil { + return 0, err + } + size, err := io.Copy(tf, r) + if err != nil { + tf.Close() + os.Remove(tf.Name()) + return 0, err + } + if err := tf.Chmod(perm); err != nil { + tf.Close() + os.Remove(tf.Name()) + return 0, err + } + if err := tf.Close(); err != nil { + os.Remove(tf.Name()) + return 0, err + } + if err := os.Rename(tf.Name(), dest); err != nil { + os.Remove(tf.Name()) + return 0, err + } + return size, nil +} diff --git a/gocache/cachers/disk_notwindows_test.go b/gocache/cachers/disk_notwindows_test.go new file mode 100644 index 0000000..b242f4b --- /dev/null +++ b/gocache/cachers/disk_notwindows_test.go @@ -0,0 +1,31 @@ +//go:build !windows + +package cachers + +import ( + "context" + "os" + "strings" + "testing" +) + +// TestOutputFileExecutable verifies that output files are written with the +// executable bit set, as cache entries can be linked binaries that cmd/go +// executes directly from the cache directory. +func TestOutputFileExecutable(t *testing.T) { + dc := &DiskCache{Dir: t.TempDir()} + actionID := strings.Repeat("a", 64) + outputID := strings.Repeat("b", 64) + body := strings.NewReader("some output") + diskPath, err := dc.Put(context.Background(), actionID, outputID, int64(body.Len()), body) + if err != nil { + t.Fatal(err) + } + fi, err := os.Stat(diskPath) + if err != nil { + t.Fatal(err) + } + if got, want := fi.Mode().Perm(), os.FileMode(0o755); got != want { + t.Errorf("output file mode = %v; want %v", got, want) + } +} diff --git a/gocache/cachers/disk_windows.go b/gocache/cachers/disk_windows.go new file mode 100644 index 0000000..64b0b90 --- /dev/null +++ b/gocache/cachers/disk_windows.go @@ -0,0 +1,104 @@ +package cachers + +import ( + "crypto/sha256" + "errors" + "flag" + "fmt" + "io" + "os" +) + +// The functions in this file are based on Go's own cache in +// cmd/go/internal/cache/cache.go, particularly putIndexEntry and copyFile. + +// writeActionFile writes the indexEntry metadata for an ActionID to disk. It +// may be called for the same actionID concurrently from multiple processes, +// and the outputID for a specific actionID may change from time to time due +// to non-deterministic builds. It makes a best-effort to delete the file if +// anything goes wrong. +func writeActionFile(dest string, b []byte) (retErr error) { + f, err := os.OpenFile(dest, os.O_WRONLY|os.O_CREATE, 0o666) + if err != nil { + return err + } + defer func() { + cerr := f.Close() + if retErr != nil || cerr != nil { + retErr = errors.Join(retErr, cerr, os.Remove(dest)) + } + }() + + _, err = f.Write(b) + if err != nil { + return err + } + + // Truncate the file only *after* writing it. + // (This should be a no-op, but truncate just in case of previous corruption.) + // + // This differs from os.WriteFile, which truncates to 0 *before* writing + // via os.O_TRUNC. Truncating only after writing ensures that a second write + // of the same content to the same file is idempotent, and does not - even + // temporarily! - undo the effect of the first write. + return f.Truncate(int64(len(b))) +} + +// writeOutputFile writes content to be cached to disk. The outputID is the +// sha256 hash of the content, and each file should only be written ~once, +// assuming no sha256 hash collisions. It may be written multiple times if +// concurrent processes are both populating the same output. The file is opened +// with FILE_SHARE_READ|FILE_SHARE_WRITE, which means both processes can write +// the same contents concurrently without conflict. +// +// It makes a best effort to clean up if anything goes wrong, but the file may +// be left in an inconsistent state in the event of disk-related errors such as +// another process taking file locks, or power loss etc. +func writeOutputFile(dest string, r io.Reader, size int64, outputID string) (_ int64, retErr error) { + info, err := os.Stat(dest) + if err == nil && info.Size() == size { + // Already exists, check the hash. + if f, err := os.Open(dest); err == nil { + h := sha256.New() + io.Copy(h, f) + f.Close() + if fmt.Sprintf("%x", h.Sum(nil)) == outputID { + // Still drain the reader to ensure associated resources are released. + return io.Copy(io.Discard, r) + } + } + } + + // Didn't successfully find the pre-existing file, write it. + mode := os.O_WRONLY | os.O_CREATE + if err == nil && info.Size() > size { + mode |= os.O_TRUNC // Should never happen, but self-heal. + } + f, err := os.OpenFile(dest, mode, 0644) + if err != nil { + return 0, fmt.Errorf("failed to open output file %q: %w", dest, err) + } + defer func() { + cerr := f.Close() + if retErr != nil || cerr != nil { + retErr = errors.Join(retErr, cerr, os.Remove(dest)) + } + }() + + // Copy file to f, but also into h to double-check hash. + h := sha256.New() + w := io.MultiWriter(f, h) + n, err := io.Copy(w, r) + if err != nil { + return 0, err + } + if got := fmt.Sprintf("%x", h.Sum(nil)); got != outputID { + if len(got) != len(outputID) && flag.Lookup("test.v") != nil { + // In tests, tolerate fake outputIDs. + } else { + return 0, errors.New("file content changed underfoot") + } + } + + return n, nil +} diff --git a/gocache/cachers/http.go b/gocache/cachers/http.go new file mode 100644 index 0000000..bc11a90 --- /dev/null +++ b/gocache/cachers/http.go @@ -0,0 +1,381 @@ +package cachers + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "log" + "net/http" + "os" + "strconv" + "sync" + "sync/atomic" + "time" + + "github.com/pierrec/lz4/v4" +) + +// ActionValue is the JSON value returned by the cacher server for an GET /action request. +type ActionValue struct { + OutputID string `json:"outputID"` + Size int64 `json:"size"` +} + +type HTTPClient struct { + // BaseURL is the base URL of the cacher server, like "http://localhost:31364". + BaseURL string + + // Disk is where to write the output files to local disk, as required by the + // cache protocol. + Disk *DiskCache + + // HTTPClient optionally specifies the http.Client to use. + // If nil, http.DefaultClient is used. + HTTPClient *http.Client + + // Verbose optionally specifies whether to log verbose messages. + Verbose bool + + // AccessToken optionally specifies a Bearer access token to include + // in requests to the server. + AccessToken string + + // BestEffortHTTP, when true, makes all HTTP errors non-fatal. + // Get returns a cache miss and Put returns the local disk result, + // silently ignoring any HTTP failures (connection errors, server errors, etc.). + // It also makes the remote HTTP PUTs run asynchronously: Put returns as soon + // as the blob is on local disk, and the upload happens in a background goroutine. + BestEffortHTTP bool + + // AsyncPutTimeout, if non-nil, returns how long a background PUT of the + // given object size can run before its context is canceled. A zero or + // negative duration means no timeout. It only applies when BestEffortHTTP + // is set. + AsyncPutTimeout func(size int64) time.Duration + + // AsyncPutMaxConcurrent, if positive, caps the number of background PUTs + // running at once. New PUTs beyond the cap are queued until an existing PUT + // finishes. Disk writes are not affected, and it only applies when + // BestEffortHTTP is set. + AsyncPutMaxConcurrent int + + PutsTimedOut atomic.Int64 // The number of PUTs that timed out. + PutsCanceled atomic.Int64 // The number of PUTs canceled due to shutdown. + + asyncOnce sync.Once + asyncPutSem chan struct{} + inFlightWG sync.WaitGroup + baseCtx context.Context + baseCancel context.CancelFunc +} + +func (c *HTTPClient) ensureAsyncSetup() { + c.asyncOnce.Do(func() { + c.baseCtx, c.baseCancel = context.WithCancel(context.Background()) + if c.AsyncPutMaxConcurrent > 0 { + c.asyncPutSem = make(chan struct{}, c.AsyncPutMaxConcurrent) + } + }) +} + +// asyncPutSemaphore returns the buffered channel bounding concurrent background +// PUTs, or nil if AsyncPutMaxConcurrent is unset (unbounded). +func (c *HTTPClient) asyncPutSemaphore() chan struct{} { + c.ensureAsyncSetup() + return c.asyncPutSem +} + +// asyncPutContext returns the context shared by all background PUTs. Shutdown +// cancels it when the drain deadline elapses. +func (c *HTTPClient) asyncPutContext() context.Context { + c.ensureAsyncSetup() + return c.baseCtx +} + +// Shutdown blocks until all background PUTs have finished or ctx is done, +// whichever comes first. When ctx is done first, it cancels the context shared +// by any still-running PUTs and waits for them to unwind, so no background PUT +// outlives Shutdown. It reports whether all PUTs drained before ctx was done. +func (c *HTTPClient) Shutdown(ctx context.Context) (drained bool) { + done := make(chan struct{}) + go func() { + c.inFlightWG.Wait() + close(done) + }() + select { + case <-done: + return true + case <-ctx.Done(): + c.ensureAsyncSetup() + c.baseCancel() + <-done + return false + } +} + +func (c *HTTPClient) httpClient() *http.Client { + if c.HTTPClient != nil { + return c.HTTPClient + } + return http.DefaultClient +} + +// tryDrainResponse reads and throws away a small bounded amount of data from +// res.Body. This is a best-effort attempt to allow connection reuse. (Go's +// HTTP/1 Transport won't reuse a TCP connection unless you fully consume HTTP +// responses) +func tryDrainResponse(res *http.Response) { + io.CopyN(io.Discard, res.Body, 4<<10) +} + +func tryReadErrorMessage(res *http.Response) []byte { + msg, _ := io.ReadAll(io.LimitReader(res.Body, 4<<10)) + return msg +} + +// responseBody returns the response body and the uncompressed content length. +// If the response has Content-Encoding: lz4, the body is wrapped with an lz4 +// decompressor and the uncompressed length is read from X-Uncompressed-Length. +// For uncompressed responses, Content-Length is used directly. +func responseBody(res *http.Response) (body io.Reader, uncompressedLength int64, err error) { + if res.Header.Get("Content-Encoding") == "lz4" { + sizeStr := res.Header.Get("X-Uncompressed-Length") + if sizeStr == "" { + return nil, 0, fmt.Errorf("lz4-compressed response missing X-Uncompressed-Length header") + } + size, err := strconv.ParseInt(sizeStr, 10, 64) + if err != nil { + return nil, 0, fmt.Errorf("invalid X-Uncompressed-Length %q: %v", sizeStr, err) + } + return lz4.NewReader(res.Body), size, nil + } + if res.ContentLength == -1 { + return nil, 0, fmt.Errorf("no Content-Length from server") + } + return res.Body, res.ContentLength, nil +} + +func (c *HTTPClient) Get(ctx context.Context, actionID string) (outputID, diskPath string, err error) { + outputID, diskPath, err = c.Disk.Get(ctx, actionID) + if err == nil && outputID != "" { + return outputID, diskPath, nil + } + + req, _ := http.NewRequestWithContext(ctx, "GET", c.BaseURL+"/action/"+actionID, nil) + if c.AccessToken != "" { + req.Header.Set("Authorization", "Bearer "+c.AccessToken) + } + + req.Header.Set("Accept-Encoding", "lz4") + + // Set a header to indicate we want the object and metadata in one response. + // Prior to 2025-08-09, the protocol was two separate requests. Rather than + // change this repo's protocol and potentially break existing clients, + // we just add a header to indicate we want the object and then we support + // both the old and new response types. + req.Header.Set("Want-Object", "1") // opt in to new single roundtrip protocol + + res, err := c.httpClient().Do(req) + if err != nil { + if c.BestEffortHTTP { + return "", "", nil + } + return "", "", err + } + defer res.Body.Close() + defer tryDrainResponse(res) + if res.StatusCode == http.StatusNotFound { + return "", "", nil + } + if res.StatusCode != http.StatusOK { + msg := tryReadErrorMessage(res) + if c.BestEffortHTTP && res.StatusCode == http.StatusUnauthorized { + // Known error code that will repeatedly happen, so avoid + // filling the logs with these errors. + // + // 401: can happen when gocached restarts during a session, as it + // doesn't persist access tokens. + // TODO(tomhjp): make the client retry auth in the background. + } else { + log.Printf("error GET /action/%s: %v, %s", actionID, res.Status, msg) + } + if c.BestEffortHTTP { + return "", "", nil + } + return "", "", fmt.Errorf("unexpected GET /action/%s status %v", actionID, res.Status) + } + + switch res.Header.Get("Content-Type") { + default: + return "", "", fmt.Errorf("unexpected Content-Type %q from server", res.Header.Get("Content-Type")) + + case "application/octet-stream": // new single roundtrip protocol + outputID = res.Header.Get("Go-Output-Id") + if outputID == "" { + return "", "", fmt.Errorf("missing Go-Output-Id header in response") + } + var body io.Reader + var size int64 + body, size, err = responseBody(res) + if err != nil { + return "", "", err + } + diskPath, err = c.Disk.Put(ctx, actionID, outputID, size, body) + + case "application/json": // old two-hop protocol + var av ActionValue + if err := json.NewDecoder(res.Body).Decode(&av); err != nil { + return "", "", err + } + outputID = av.OutputID + + // If not on disk, download it to disk. + var putBody io.Reader + if av.Size == 0 { + putBody = bytes.NewReader(nil) + } else { + req, _ = http.NewRequestWithContext(ctx, "GET", c.BaseURL+"/output/"+outputID, nil) + req.Header.Set("Accept-Encoding", "lz4") + res, err = c.httpClient().Do(req) + if err != nil { + if c.BestEffortHTTP { + return "", "", nil + } + return "", "", err + } + defer res.Body.Close() + defer tryDrainResponse(res) + if res.StatusCode == http.StatusNotFound { + return "", "", nil + } + if res.StatusCode != http.StatusOK { + msg := tryReadErrorMessage(res) + log.Printf("error GET /output/%s: %v, %s", outputID, res.Status, msg) + if c.BestEffortHTTP { + return "", "", nil + } + return "", "", fmt.Errorf("unexpected GET /output/%s status %v", outputID, res.Status) + } + putBody, _, err = responseBody(res) + if err != nil { + return "", "", err + } + } + diskPath, err = c.Disk.Put(ctx, actionID, outputID, av.Size, putBody) + } + + return outputID, diskPath, err +} + +func (c *HTTPClient) Put(ctx context.Context, actionID, outputID string, size int64, body io.Reader) (diskPath string, _ error) { + // Write to disk first. + diskPath, err := c.Disk.Put(ctx, actionID, outputID, size, body) + if err != nil { + log.Printf("HTTPClient.Put local disk write error: %v", err) + return "", err + } + + if c.BestEffortHTTP { + sem := c.asyncPutSemaphore() + // Background PUTs use a client-owned context, not the per-request ctx, + // so they survive the RPC returning. + putCtx := c.asyncPutContext() + c.inFlightWG.Go(func() { + // Acquire a concurrency slot before opening the file, so a queued + // backlog doesn't hold an open file descriptor per pending upload. + if sem != nil { + select { + case sem <- struct{}{}: + case <-putCtx.Done(): + c.PutsCanceled.Add(1) + return + } + defer func() { + <-sem + }() + } + f, err := os.Open(diskPath) + if err != nil { + log.Printf("HTTPClient.Put local disk open after write error: %v", err) + return + } + c.putRemote(putCtx, actionID, outputID, size, f) + }) + return diskPath, nil + } + + f, err := os.Open(diskPath) + if err != nil { + log.Printf("HTTPClient.Put local disk open after write error: %v", err) + return "", err + } + if err := c.putRemote(ctx, actionID, outputID, size, f); err != nil { + return diskPath, err + } + + return diskPath, nil +} + +func (c *HTTPClient) putRemote(ctx context.Context, actionID, outputID string, size int64, body io.ReadCloser) error { + defer body.Close() + if c.BestEffortHTTP && c.AsyncPutTimeout != nil { + if d := c.AsyncPutTimeout(size); d > 0 { + var cancel context.CancelFunc + ctx, cancel = context.WithTimeout(ctx, d) + defer cancel() + } + } + if c.BestEffortHTTP { + defer func() { + switch { + case c.baseCtx.Err() != nil: + c.PutsCanceled.Add(1) + case ctx.Err() != nil: + c.PutsTimedOut.Add(1) + } + }() + } + // For a zero-length body, hand net/http NoBody so it sends an explicit + // Content-Length: 0 rather than switching to chunked transfer encoding, + // which the server rejects. + var reqBody io.Reader = body + if size == 0 { + reqBody = http.NoBody + } + req, _ := http.NewRequestWithContext(ctx, "PUT", c.BaseURL+"/"+actionID+"/"+outputID, reqBody) + req.ContentLength = size + if c.AccessToken != "" { + req.Header.Set("Authorization", "Bearer "+c.AccessToken) + } + + res, err := c.httpClient().Do(req) + if err != nil { + log.Printf("error PUT /%s/%s: %v", actionID, outputID, err) + return err + } + + defer res.Body.Close() + if res.StatusCode != http.StatusNoContent { + errMsg := fmt.Sprintf("error PUT /%s/%s: %s, %s", actionID, outputID, res.Status, tryReadErrorMessage(res)) + if c.BestEffortHTTP && (res.StatusCode == http.StatusUnauthorized || res.StatusCode == http.StatusForbidden) { + // Known error codes that will repeatedly happen, so avoid + // filling the logs with these errors. + // + // 401: can happen when gocached restarts during a session, as it + // doesn't persist access tokens. + // TODO(tomhjp): make the client retry auth in the background. + // + // 403: can happen when authed with a JWT that didn't grant global + // write permissions. + // TODO(tomhjp): support namespaces so all sessions can safely write. + } else { + log.Print(errMsg) + } + return errors.New(errMsg) + } + + return nil +} diff --git a/gocache/cachers/http_test.go b/gocache/cachers/http_test.go new file mode 100644 index 0000000..b730960 --- /dev/null +++ b/gocache/cachers/http_test.go @@ -0,0 +1,513 @@ +package cachers + +import ( + "bytes" + "context" + "encoding/binary" + "encoding/json" + "io" + "net" + "net/http" + "net/http/httptest" + "os" + "strconv" + "sync/atomic" + "testing" + "testing/synctest" + "time" + + "github.com/pierrec/lz4/v4" +) + +func lz4Compress(t *testing.T, data []byte) []byte { + t.Helper() + var buf bytes.Buffer + w := lz4.NewWriter(&buf) + if err := w.Apply(lz4.SizeOption(uint64(len(data)))); err != nil { + t.Fatal(err) + } + if _, err := w.Write(data); err != nil { + t.Fatal(err) + } + if err := w.Close(); err != nil { + t.Fatal(err) + } + compressed := buf.Bytes() + + // Verify the lz4 frame header contains the uncompressed content size + // at bytes [6:14] (little-endian uint64), after the 4-byte magic number + // and 2-byte FLG+BD descriptor. + if len(compressed) < 14 { + t.Fatalf("compressed output too short: %d bytes", len(compressed)) + } + gotSize := binary.LittleEndian.Uint64(compressed[6:14]) + if gotSize != uint64(len(data)) { + t.Fatalf("lz4 frame content size = %d, want %d", gotSize, len(data)) + } + + return compressed +} + +func TestHTTPClientGetLZ4(t *testing.T) { + const ( + testActionID = "aabbccdd" + testOutputID = "eeff0011" + ) + testData := []byte("hello, this is cached build output data for testing") + + tests := []struct { + name string + compress bool + oldProto bool + }{ + {"new_protocol_uncompressed", false, false}, + {"new_protocol_lz4", true, false}, + {"old_protocol_uncompressed", false, true}, + {"old_protocol_lz4", true, true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var gotAcceptEncoding []string + + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotAcceptEncoding = append(gotAcceptEncoding, r.Header.Get("Accept-Encoding")) + + switch { + case tt.oldProto && r.URL.Path == "/action/"+testActionID: + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(ActionValue{ + OutputID: testOutputID, + Size: int64(len(testData)), + }) + + case tt.oldProto && r.URL.Path == "/output/"+testOutputID: + body := testData + if tt.compress { + body = lz4Compress(t, testData) + w.Header().Set("Content-Encoding", "lz4") + w.Header().Set("X-Uncompressed-Length", strconv.Itoa(len(testData))) + } + w.Header().Set("Content-Length", strconv.Itoa(len(body))) + w.Write(body) + + case !tt.oldProto && r.URL.Path == "/action/"+testActionID: + w.Header().Set("Content-Type", "application/octet-stream") + w.Header().Set("Go-Output-Id", testOutputID) + body := testData + if tt.compress { + body = lz4Compress(t, testData) + w.Header().Set("Content-Encoding", "lz4") + w.Header().Set("X-Uncompressed-Length", strconv.Itoa(len(testData))) + } + w.Header().Set("Content-Length", strconv.Itoa(len(body))) + w.Write(body) + + default: + http.NotFound(w, r) + } + })) + defer ts.Close() + + hc := &HTTPClient{ + BaseURL: ts.URL, + Disk: &DiskCache{Dir: t.TempDir()}, + } + + outputID, diskPath, err := hc.Get(context.Background(), testActionID) + if err != nil { + t.Fatal(err) + } + if outputID != testOutputID { + t.Errorf("outputID = %q, want %q", outputID, testOutputID) + } + if diskPath == "" { + t.Fatal("diskPath is empty") + } + + got, err := os.ReadFile(diskPath) + if err != nil { + t.Fatal(err) + } + if !bytes.Equal(got, testData) { + t.Errorf("disk content = %q, want %q", got, testData) + } + + // Verify Accept-Encoding: lz4 was sent on all requests. + if len(gotAcceptEncoding) == 0 { + t.Fatal("no requests received by server") + } + for i, ae := range gotAcceptEncoding { + if ae != "lz4" { + t.Errorf("request %d: Accept-Encoding = %q, want %q", i, ae, "lz4") + } + } + if tt.oldProto && len(gotAcceptEncoding) != 2 { + t.Errorf("old protocol: got %d requests, want 2", len(gotAcceptEncoding)) + } + }) + } +} + +// TestHTTPClientPutServerRejectsBody verifies that Put still writes to disk +// when the server returns 403 without reading the request body. +func TestHTTPClientPutServerRejectsBody(t *testing.T) { + const ( + testActionID = "aabbccdd" + testOutputID = "eeff0011" + ) + + largeBody := bytes.Repeat([]byte("x"), 512<<10) + + tests := []struct { + name string + data []byte + }{ + {"empty", nil}, + {"small", []byte("hello, this is cached build output data for testing")}, + {"512k", largeBody}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + // Return 403 immediately without reading the request body. + http.Error(w, "forbidden", http.StatusForbidden) + })) + defer ts.Close() + + hc := &HTTPClient{ + BaseURL: ts.URL, + Disk: &DiskCache{Dir: t.TempDir()}, + } + + diskPath, err := hc.Put(context.Background(), testActionID, testOutputID, int64(len(tt.data)), bytes.NewReader(tt.data)) + if diskPath == "" { + t.Fatalf("diskPath is empty; err = %v", err) + } + + got, err := os.ReadFile(diskPath) + if err != nil { + t.Fatal(err) + } + if !bytes.Equal(got, tt.data) { + t.Errorf("disk content length = %d, want %d", len(got), len(tt.data)) + } + }) + } +} + +// blockingPutTransport is a fake http.RoundTripper whose PUTs block until +// release is closed. It counts how many round trips have started, letting tests +// observe in-flight uploads without a real network or wall-clock timing. It is +// safe to use inside a synctest bubble because it performs no real I/O. +type blockingPutTransport struct { + started atomic.Int64 + release chan struct{} +} + +func newBlockingPutTransport() *blockingPutTransport { + return &blockingPutTransport{ + release: make(chan struct{}), + } +} + +func (tr *blockingPutTransport) RoundTrip(req *http.Request) (*http.Response, error) { + tr.started.Add(1) + select { + case <-tr.release: + case <-req.Context().Done(): + return nil, req.Context().Err() + } + io.Copy(io.Discard, req.Body) + return &http.Response{ + StatusCode: http.StatusNoContent, + Body: http.NoBody, + Header: make(http.Header), + }, nil +} + +// TestHTTPClientAsyncPutConcurrencyLimit verifies that no more than +// AsyncPutMaxConcurrent uploads run at once, and that PUTs beyond the cap are +// queued. +func TestHTTPClientAsyncPutConcurrencyLimit(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + tr := newBlockingPutTransport() + hc := &HTTPClient{ + BaseURL: "http://fake", + Disk: &DiskCache{Dir: t.TempDir()}, + HTTPClient: &http.Client{Transport: tr}, + BestEffortHTTP: true, + AsyncPutMaxConcurrent: 2, + } + + data := []byte("some cached output") + for i := range 5 { + id := strconv.Itoa(i) + diskPath, err := hc.Put(t.Context(), "aaaa"+id, "bbbb"+id, int64(len(data)), bytes.NewReader(data)) + if err != nil { + t.Fatalf("Put %d returned error: %v", i, err) + } + if diskPath == "" { + t.Fatalf("Put %d returned empty diskPath", i) + } + } + + // Only 2 uploads may run at once; the other 3 wait for a slot. + synctest.Wait() + if got := tr.started.Load(); got != 2 { + t.Errorf("started %d uploads, want 2", got) + } + + // Releasing the blocked uploads lets the queued ones drain until all 5 + // have run. + close(tr.release) + synctest.Wait() + if got := tr.started.Load(); got != 5 { + t.Errorf("started %d uploads after release, want 5", got) + } + }) +} + +// TestHTTPClientShutdownDrains verifies Shutdown blocks until in-flight PUTs +// finish and returns true when they drain in time. +func TestHTTPClientShutdownDrains(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + tr := newBlockingPutTransport() + hc := &HTTPClient{ + BaseURL: "http://fake", + Disk: &DiskCache{Dir: t.TempDir()}, + HTTPClient: &http.Client{Transport: tr}, + BestEffortHTTP: true, + } + + data := []byte("some cached output") + if _, err := hc.Put(t.Context(), "aabbccdd", "eeff0011", int64(len(data)), bytes.NewReader(data)); err != nil { + t.Fatalf("Put returned error: %v", err) + } + synctest.Wait() + if got := tr.started.Load(); got != 1 { + t.Fatalf("started %d uploads, want 1", got) + } + + shutdownReturned := make(chan bool, 1) + go func() { + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + shutdownReturned <- hc.Shutdown(ctx) + }() + + // With the upload still blocked, Shutdown must not return. + synctest.Wait() + select { + case <-shutdownReturned: + t.Fatal("Shutdown returned before the in-flight PUT finished") + default: + } + + close(tr.release) + synctest.Wait() + select { + case drained := <-shutdownReturned: + if !drained { + t.Error("Shutdown returned drained=false, want true") + } + default: + t.Fatal("Shutdown did not return after the PUT was released") + } + }) +} + +// TestHTTPClientShutdownCancelsStragglers verifies that when the drain deadline +// elapses, Shutdown returns false and cancels the context shared by background +// PUTs, so an upload blocked past the deadline is abandoned rather than left to +// outlive Shutdown. +func TestHTTPClientShutdownCancelsStragglers(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + tr := newBlockingPutTransport() // never released + hc := &HTTPClient{ + BaseURL: "http://fake", + Disk: &DiskCache{Dir: t.TempDir()}, + HTTPClient: &http.Client{Transport: tr}, + BestEffortHTTP: true, + } + + data := []byte("some cached output") + if _, err := hc.Put(t.Context(), "aabbccdd", "eeff0011", int64(len(data)), bytes.NewReader(data)); err != nil { + t.Fatalf("Put returned error: %v", err) + } + synctest.Wait() + if got := tr.started.Load(); got != 1 { + t.Fatalf("started %d uploads, want 1", got) + } + + // The upload never unblocks, so Shutdown must hit its deadline, cancel + // the shared context, and report that it did not drain. + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + if drained := hc.Shutdown(ctx); drained { + t.Error("Shutdown returned drained=true, want false") + } + + // Canceling the shared context unblocks the straggler, so the + // background goroutine exits rather than leaking. + synctest.Wait() + if got := hc.PutsCanceled.Load(); got != 1 { + t.Errorf("PutsCanceled = %d, want 1", got) + } + if got := hc.PutsTimedOut.Load(); got != 0 { + t.Errorf("PutsTimedOut = %d, want 0", got) + } + }) +} + +// TestHTTPClientAsyncPutTimeout verifies that AsyncPutTimeout is applied per +// object size and that a PUT canceled by its own timeout is counted by +// PutsTimedOut. +func TestHTTPClientAsyncPutTimeout(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + tr := newBlockingPutTransport() // never released + var gotSize int64 + hc := &HTTPClient{ + BaseURL: "http://fake", + Disk: &DiskCache{Dir: t.TempDir()}, + HTTPClient: &http.Client{Transport: tr}, + BestEffortHTTP: true, + AsyncPutTimeout: func(size int64) time.Duration { + gotSize = size + return time.Second + }, + } + + data := []byte("some cached output") + if _, err := hc.Put(t.Context(), "aabbccdd", "eeff0011", int64(len(data)), bytes.NewReader(data)); err != nil { + t.Fatalf("Put returned error: %v", err) + } + + // The upload never unblocks, so its own timeout must fire and count it + // as timed out without any call to Shutdown. Sleeping past the timeout + // durably blocks this goroutine so the bubble's fake clock advances and + // the 1s timer fires. + time.Sleep(2 * time.Second) + synctest.Wait() + if gotSize != int64(len(data)) { + t.Errorf("AsyncPutTimeout called with size %d, want %d", gotSize, len(data)) + } + if got := hc.PutsTimedOut.Load(); got != 1 { + t.Errorf("PutsTimedOut = %d, want 1", got) + } + if got := hc.PutsCanceled.Load(); got != 0 { + t.Errorf("PutsCanceled = %d, want 0", got) + } + }) +} + +func TestHTTPClientBestEffort(t *testing.T) { + const ( + testActionID = "aabbccdd" + testOutputID = "eeff0011" + ) + testData := []byte("hello, this is cached build output data for testing") + + // closedPortURL returns a URL pointing at a TCP port that immediately + // refuses connections (RST). + closedPortURL := func(t *testing.T) string { + t.Helper() + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + addr := ln.Addr().String() + ln.Close() + return "http://" + addr + } + + tests := []struct { + name string + httpErr string // "rst" or "500" + }{ + {"connection_refused", "rst"}, + {"server_500", "500"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var baseURL string + if tt.httpErr == "rst" { + baseURL = closedPortURL(t) + } else { + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + http.Error(w, "internal server error", http.StatusInternalServerError) + })) + defer ts.Close() + baseURL = ts.URL + } + + t.Run("Get", func(t *testing.T) { + hc := &HTTPClient{ + BaseURL: baseURL, + Disk: &DiskCache{Dir: t.TempDir()}, + BestEffortHTTP: true, + } + + outputID, diskPath, err := hc.Get(context.Background(), testActionID) + if err != nil { + t.Fatalf("BestEffortHTTP Get returned error: %v", err) + } + if outputID != "" || diskPath != "" { + t.Errorf("expected cache miss, got outputID=%q, diskPath=%q", outputID, diskPath) + } + }) + + t.Run("Get_without_best_effort", func(t *testing.T) { + hc := &HTTPClient{ + BaseURL: baseURL, + Disk: &DiskCache{Dir: t.TempDir()}, + BestEffortHTTP: false, + } + + _, _, err := hc.Get(context.Background(), testActionID) + if err == nil { + t.Fatal("expected error from Get without BestEffortHTTP, got nil") + } + }) + + t.Run("Put", func(t *testing.T) { + hc := &HTTPClient{ + BaseURL: baseURL, + Disk: &DiskCache{Dir: t.TempDir()}, + BestEffortHTTP: true, + } + + diskPath, err := hc.Put(context.Background(), testActionID, testOutputID, int64(len(testData)), bytes.NewReader(testData)) + if err != nil { + t.Fatalf("BestEffortHTTP Put returned error: %v", err) + } + if diskPath == "" { + t.Fatal("diskPath is empty") + } + + got, err := os.ReadFile(diskPath) + if err != nil { + t.Fatal(err) + } + if !bytes.Equal(got, testData) { + t.Errorf("disk content = %q, want %q", got, testData) + } + }) + + t.Run("Put_without_best_effort", func(t *testing.T) { + hc := &HTTPClient{ + BaseURL: baseURL, + Disk: &DiskCache{Dir: t.TempDir()}, + BestEffortHTTP: false, + } + + _, err := hc.Put(context.Background(), testActionID, testOutputID, int64(len(testData)), bytes.NewReader(testData)) + if err == nil { + t.Fatal("expected error from Put without BestEffortHTTP, got nil") + } + }) + }) + } +} diff --git a/gocache/gocached/diskspace_other.go b/gocache/gocached/diskspace_other.go new file mode 100644 index 0000000..3554499 --- /dev/null +++ b/gocache/gocached/diskspace_other.go @@ -0,0 +1,11 @@ +//go:build !unix + +package gocached + +import "errors" + +// fsTotalBytes returns the total size in bytes of the filesystem holding dir. +// It is not implemented on this platform; callers fall back to a default. +func fsTotalBytes(dir string) (int64, error) { + return 0, errors.ErrUnsupported +} diff --git a/gocache/gocached/diskspace_unix.go b/gocache/gocached/diskspace_unix.go new file mode 100644 index 0000000..1d5de72 --- /dev/null +++ b/gocache/gocached/diskspace_unix.go @@ -0,0 +1,14 @@ +//go:build unix + +package gocached + +import "syscall" + +// fsTotalBytes returns the total size in bytes of the filesystem holding dir. +func fsTotalBytes(dir string) (int64, error) { + var st syscall.Statfs_t + if err := syscall.Statfs(dir, &st); err != nil { + return 0, err + } + return int64(st.Blocks) * int64(st.Bsize), nil +} diff --git a/gocache/gocached/gocached.go b/gocache/gocached/gocached.go new file mode 100644 index 0000000..2336ce4 --- /dev/null +++ b/gocache/gocached/gocached.go @@ -0,0 +1,3357 @@ +// Copyright (c) Tailscale Inc & AUTHORS +// SPDX-License-Identifier: BSD-3-Clause + +// The gocached package provides an HTTP server daemon that go-cacher can hit. +// It does cache tiering and evicts old large things from disk, and can fetch +// metadata and object contents from peer cache servers. +// +// It uses sqlite (the pure Go modernc.org/sqlite driver) to store metadata and +// indexes. +// +/* + +It speaks the same protocol as go-cacher-server, but requires +the "Want-Object: 1" header variant on the GET request. + + GET /action/ + Want-Object: 1 + + 200 OK + Content-Type: application/octet-stream + Content-Length: 1234 + Go-Output-Id: xxxxxxxxxxx + + + +And to insert an object: + + PUT // + Content-Length: 1234 + + + +*/ +package gocached + +import ( + "cmp" + "context" + "crypto/rand" + "crypto/sha256" + "database/sql" + "encoding/hex" + "encoding/json" + "errors" + "expvar" + "fmt" + "io" + "log" + "maps" + "math" + "net/http" + "net/http/pprof" + "os" + "path/filepath" + "reflect" + "runtime" + "slices" + "strconv" + "strings" + "sync" + "sync/atomic" + "time" + + "github.com/pierrec/lz4/v4" + "github.com/prometheus/client_golang/prometheus" + "github.com/prometheus/client_golang/prometheus/collectors" + "github.com/prometheus/client_golang/prometheus/promhttp" + dto "github.com/prometheus/client_model/go" + ijwt "github.com/tailscale/tb/gocache/internal/jwt" + "github.com/tailscale/tb/gocache/logger" + _ "modernc.org/sqlite" +) + +const ( + // smallObjectSize is the maximum size of an object that we store inline in the + // database, rather than on disk. Empirically, about half of objects are 1KB or + // smaller. + smallObjectSize = 1 << 10 + + // lz4CompressThreshold is the minimum blob size at which we lz4-compress + // before storing on disk. Intentionally two bytes above smallObjectSize so + // there's a real gap (disk blobs that aren't lz4-compressed), preventing + // assumptions that the two thresholds are equal while keeping them close. + lz4CompressThreshold = smallObjectSize + 2 + + // tokenPrefix is the prefix for all gocached access tokens. + tokenPrefix = "gocached-token-" + + // gocachedAudience is the audience we require JWTs to have. Could be + // configurable in future, but for now just needs to be specific to gocached. + gocachedAudience = "gocached" +) + +const schemaVersion = 4 + +// walJournalSizeLimit caps the on-disk WAL size after a successful +// checkpoint. It is set as a per-connection PRAGMA so that even SQLite's +// built-in PASSIVE autocheckpoint, which would otherwise leave the file at +// its high-water mark, truncates the WAL down to this size. Without it the +// WAL can grow without bound under continuous write traffic and never +// shrink, even when frames are being checkpointed in place. +const walJournalSizeLimit = 1 << 30 // 1 GiB + +// checkpointInterval is how often [Server.runCheckpointLoop] runs a TRUNCATE +// checkpoint in the background. SQLite's autocheckpoint only runs PASSIVE +// checkpoints (which reuse WAL space in place but never shrink the file on +// disk past walJournalSizeLimit); a periodic explicit TRUNCATE is what +// actually keeps the file small in steady state. +const checkpointInterval = time.Minute + +// dbSizeMetricsInterval is how often [Server.runDBSizeMetricsLoop] re-stats +// the SQLite files to update the size gauges. It is intentionally shorter +// than checkpointInterval so the gauge sees the WAL grow between checkpoints, +// not just snap back to zero each minute. +const dbSizeMetricsInterval = 15 * time.Second + +// defaultShardPrefixLen is the [Server.shardPrefixLen] used when the option +// is not set (zero value). 2 yields 256 shards ("00".."ff"), which keeps +// individual scans fast even at hundreds of millions of rows. +const defaultShardPrefixLen = 2 + +// minShardPrefixLen is the minimum [Server.shardPrefixLen]. 1 yields 16 +// shards. Going to 0 (1 shard) would defeat the purpose of sharded +// statistics, so we reject it. +const minShardPrefixLen = 1 + +// maxShardPrefixLen is the maximum [Server.shardPrefixLen]. 4 hex chars -> +// 65,536 shards is more than we ever expect to need; higher values would +// make the stats loop spend most of its time on shards with little or no +// data, and would also blow up the size of the BlobShardStats table. +const maxShardPrefixLen = 4 + +// shardStatsMinInterval sets a defensive floor on how long [Server.runShardStatsLoop] will +// sleep between iterations to ensure it never spins in a tight loop. [shardStalenessTarget] +// dictates the interval in steady state. +const shardStatsMinInterval = 2 * time.Second + +// cleanupTickInterval is how often [Server.runCleanLoop] wakes up to check +// whether the cache is over its limits and, if so, evict one batch of +// Actions. It's tight on purpose: each tick does at most cleanupBatchSize +// deletes so the write transaction is short. +const cleanupTickInterval = 1 * time.Second + +// cleanupBatchSize is the maximum number of Actions deleted in a single +// [Server.evictOldestActions] call. Small enough that the write +// transaction (including any file removals) holds the SQLite write lock +// for only a brief moment, large enough that steady-state eviction keeps +// up with realistic write traffic at ~cleanupBatchSize/cleanupTickInterval. +const cleanupBatchSize = 200 + +// shardStalenessTarget bounds how old any single shard's cohort histogram is +// allowed to get. The stats loop sleeps until the oldest shard reaches this +// age, then scans it; on a quiet cache that means scans naturally settle to +// ~numShards/shardStalenessTarget per second instead of running continuously. +// 10m drift on the 1d cohort is well under 1% error, which is negligible +// compared to the day-granularity cohort cutoffs we record. +const shardStalenessTarget = 10 * time.Minute + +const schema = ` +PRAGMA journal_mode=WAL; +CREATE TABLE IF NOT EXISTS Actions ( + NamespaceID INTEGER NOT NULL, -- 0 for global trusted namespace (generation 1; see WithGlobalGeneration) + ActionID TEXT NOT NULL, + BlobID INTEGER NOT NULL, + AltOutputID TEXT NOT NULL DEFAULT '', -- if non-empty, the alternate object ID to use for this action; NULL means the blob's sha256 + CreateTime INTEGER NOT NULL, -- unix sec when inserted (locally or on a peer) + AccessTime INTEGER NOT NULL, -- unix sec of last access + + PRIMARY KEY (NamespaceID, ActionID), + + CHECK (ActionID = lower(ActionID)), + CHECK (ActionID GLOB '[0-9a-f]*'), + CHECK (CreateTime >= 0), + CHECK (AccessTime >= 0) +) STRICT; + +CREATE INDEX IF NOT EXISTS idx_actions_access ON Actions(AccessTime); +CREATE INDEX IF NOT EXISTS idx_actions_blobid ON Actions(BlobID); + +CREATE TABLE IF NOT EXISTS Blobs ( + BlobID INTEGER PRIMARY KEY AUTOINCREMENT, + SHA256 TEXT NOT NULL, + StoredSize INTEGER NOT NULL, -- bytes stored: len(SmallData) if inline, file size on disk (possibly lz4) + UncompressedSize INTEGER NOT NULL, -- original uncompressed content size + SmallData BLOB, + + CHECK (SmallData IS NULL OR length(SmallData) = StoredSize) +) STRICT; + +CREATE UNIQUE INDEX IF NOT EXISTS idx_blobs_sha256 ON Blobs(SHA256); + +CREATE TABLE IF NOT EXISTS Namespaces ( + NamespaceID INTEGER PRIMARY KEY AUTOINCREMENT, + Namespace TEXT NOT NULL UNIQUE +) STRICT; + +-- BlobShardStats persists per-shard usage histograms across restarts so +-- the server can serve PUTs with accurate global usage from right at startup, +-- without blocking for minutes computing stats before serving can begin. +-- There's one row per SHA256 prefix shard, refreshes one at a time in the background. +CREATE TABLE IF NOT EXISTS BlobShardStats ( + Prefix TEXT PRIMARY KEY, -- "01".."ff" by default for shardPrefixLen=2; must be lowercase hex + ScannedAt INTEGER NOT NULL, -- unix seconds + StatsJSON TEXT NOT NULL -- JSON-encoded [usageStats] of that shard +) STRICT; +` + +func openDB(dbDir string) (*sql.DB, error) { + dbPath := filepath.Join(dbDir, fmt.Sprintf("gocached-v%d.db", schemaVersion)) + + // If the v4 DB doesn't exist but v3 does, migrate it. + if schemaVersion == 4 { + if _, err := os.Stat(dbPath); os.IsNotExist(err) { + v3Path := filepath.Join(dbDir, "gocached-v3.db") + if _, err := os.Stat(v3Path); err == nil { + if err := migrateV3ToV4(dbDir, v3Path, dbPath); err != nil { + return nil, fmt.Errorf("migrating v3→v4: %w", err) + } + } + } + } + + dsn := fmt.Sprintf("file:%s?_pragma=busy_timeout(5000)&_pragma=journal_size_limit(%d)", + dbPath, walJournalSizeLimit) + db, err := sql.Open("sqlite", dsn) + if err != nil { + return nil, err + } + numConns := min(runtime.NumCPU(), 4) + db.SetMaxOpenConns(numConns) + db.SetMaxIdleConns(numConns) + db.SetConnMaxLifetime(0) // no limit + var ddl string + err = db.QueryRow(`SELECT sql FROM sqlite_master WHERE type='table' AND name='Namespaces'`).Scan(&ddl) + if err == nil && strings.Contains(ddl, "lower(Namespace)") { + // Drop the Namespaces table _only_ if it has the lowercase constraint so it + // can be recreated without it. + if _, err := db.Exec(`DROP TABLE Namespaces`); err != nil { + return nil, fmt.Errorf("dropping Namespaces lowercase constraint: %w", err) + } + } + if _, err := db.Exec(schema); err != nil { + return nil, err + } + return db, nil +} + +// migrateV3ToV4 copies the v3 database to a temp file, runs the v3→v4 +// migration on it, then atomically renames it to the v4 path. This is +// crash-safe: if anything fails, the original v3 DB is untouched. +func migrateV3ToV4(dbDir, v3Path, v4Path string) error { + // Checkpoint the v3 WAL so the main .db file is self-contained before + // we copy it. Without this, data sitting in gocached-v3.db-wal would + // be silently lost. + v3DB, err := sql.Open("sqlite", "file:"+v3Path+"?_pragma=busy_timeout(5000)") + if err != nil { + return fmt.Errorf("opening v3 for checkpoint: %w", err) + } + _, err = v3DB.Exec("PRAGMA wal_checkpoint(TRUNCATE)") + v3DB.Close() + if err != nil { + return fmt.Errorf("checkpointing v3 WAL: %w", err) + } + + // Copy v3 to a temp file in the same directory (for same-filesystem rename). + tmpFile, err := os.CreateTemp(dbDir, "gocached-v4-migrate-*.db") + if err != nil { + return err + } + tmpPath := tmpFile.Name() + defer func() { + // Clean up temp file on failure. + os.Remove(tmpPath) + }() + + src, err := os.Open(v3Path) + if err != nil { + tmpFile.Close() + return err + } + if _, err := io.Copy(tmpFile, src); err != nil { + src.Close() + tmpFile.Close() + return err + } + src.Close() + if err := tmpFile.Close(); err != nil { + return err + } + + // Open the temp DB and run the migration. + tmpDB, err := sql.Open("sqlite", "file:"+tmpPath+"?_pragma=busy_timeout(5000)") + if err != nil { + return err + } + for _, stmt := range []string{ + `ALTER TABLE Blobs RENAME COLUMN BlobSize TO StoredSize`, + `ALTER TABLE Blobs ADD COLUMN UncompressedSize INTEGER NOT NULL DEFAULT 0`, + `UPDATE Blobs SET UncompressedSize = StoredSize WHERE UncompressedSize = 0`, + `PRAGMA wal_checkpoint(TRUNCATE)`, + } { + if _, err := tmpDB.Exec(stmt); err != nil { + tmpDB.Close() + return fmt.Errorf("migration statement %q: %w", stmt, err) + } + } + if err := tmpDB.Close(); err != nil { + return err + } + + // Atomic rename on same filesystem. + if err := os.Rename(tmpPath, v4Path); err != nil { + return err + } + return nil +} + +// start initializes the server, including defaults and background goroutines. +func (srv *Server) start() error { + if srv.dir == "" { + d, err := os.UserCacheDir() + if err != nil { + return fmt.Errorf("getting user cache dir: %w", err) + } + srv.dir = filepath.Join(d, "gocached") + srv.logf("Defaulting to cache dir %v ...", srv.dir) + } + if err := os.MkdirAll(srv.dir, 0750); err != nil { + return fmt.Errorf("creating cache dir: %w", err) + } + if srv.sqliteDir == "" { + srv.sqliteDir = srv.dir + } else if err := os.MkdirAll(srv.sqliteDir, 0750); err != nil { + return fmt.Errorf("creating sqlite dir: %w", err) + } + if srv.hotDir != "" { + if srv.hotDir == srv.dir { + return fmt.Errorf("hot dir %q must differ from cache dir", srv.hotDir) + } + if srv.hotCap <= 0 { + return fmt.Errorf("hot capacity must be positive when hot dir is set") + } + if err := os.MkdirAll(srv.hotDir, 0750); err != nil { + return fmt.Errorf("creating hot dir: %w", err) + } + srv.hot = newHotIndex() + } + + // Pending PUT metadata is memory-only, so any spool files left over + // from a previous process are unreachable; delete them. This must + // happen before the hot tier scan starts walking hotDir, so the scan + // never sees the directory mid-removal. + queueDir := filepath.Join(srv.dir, putQueueDirName) + if srv.hotDir != "" { + queueDir = filepath.Join(srv.hotDir, putQueueDirName) + } + if err := os.RemoveAll(queueDir); err != nil { + return fmt.Errorf("cleaning put-queue dir: %w", err) + } + if err := os.MkdirAll(queueDir, 0750); err != nil { + return fmt.Errorf("creating put-queue dir: %w", err) + } + // The cleanup intent directory, unlike the spool, is deliberately NOT + // wiped: intents surviving a crash are what let the sweeper find blobs + // that reached the main blob directory without SQLite accounting. + if err := os.MkdirAll(filepath.Join(srv.dir, cleanupDirName), 0750); err != nil { + return fmt.Errorf("creating cleanup dir: %w", err) + } + srv.putq = newPutQueue(srv, queueDir) + + if srv.hot != nil { + // Populate the index asynchronously: scanning millions of files can + // take a while, and reads can be served from existing hot files right + // away. Until the scan completes and marks the index ready, nothing + // new is written into the hot tier. + go func() { + if err := srv.hot.scan(srv.shutdownCtx, srv.hotDir, srv.logf); err != nil { + srv.logf("hot tier: scan failed; hot tier writes remain disabled: %v", err) + return + } + srv.evictHotIfOver() + }() + } + db, err := openDB(srv.sqliteDir) + if err != nil { + return fmt.Errorf("openDB: %w", err) + } + srv.db = db + + // Run a TRUNCATE checkpoint up front, before any other reader can pin a + // snapshot. If the WAL on disk is large (e.g. from a prior version of + // gocached that lacked the periodic checkpointer), this is what actually + // shrinks it. + ckCtx, ckCancel := context.WithTimeout(srv.shutdownCtx, 2*time.Minute) + if busy, log, ckpt, err := srv.checkpointTruncate(ckCtx); err != nil { + srv.logf("startup wal_checkpoint(TRUNCATE) error: %v", err) + } else { + srv.logf("startup wal_checkpoint(TRUNCATE): busy=%d log=%d ckpt=%d", busy, log, ckpt) + } + ckCancel() + + if srv.shardPrefixLen == 0 { + srv.shardPrefixLen = defaultShardPrefixLen + } + if srv.shardPrefixLen < minShardPrefixLen || srv.shardPrefixLen > maxShardPrefixLen { + return fmt.Errorf("shardPrefixLen %d out of range [%d, %d]", + srv.shardPrefixLen, minShardPrefixLen, maxShardPrefixLen) + } + + // Compute the cohort cutoffs once so every shard scan agrees. maxAge is + // added to the set if it's not already a standard duration so cleanup can + // read its bucket directly out of the aggregate. + srv.durs = computeDurs(srv.maxAge) + + // Allocate the dead-reckoning delta state. One entry per shard, indexed + // by the integer value of the SHA256 hex prefix. + srv.shardDeltas = make([]shardDelta, srv.numShards()) + + // Create the shardScanDuration histogram before loadShardStats so any + // QueryDuration values persisted by a previous run get replayed into + // the histogram immediately on restart, rather than waiting for the + // stats loop to repopulate it from fresh scans. + srv.shardScanDuration = prometheus.NewHistogram(prometheus.HistogramOpts{ + Name: "gocached_shard_scan_duration_seconds", + Help: "wall time of each per-shard usage-stats SQL scan", + Buckets: prometheus.DefBuckets, + }) + + // Buckets span roughly 1ms to 2s; cache hits should land in the low + // milliseconds, large PUTs and disk-backed GETs land higher up. + reqLatencyBuckets := []float64{0.001, 0.002, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1, 2} + srv.getDuration = prometheus.NewHistogramVec(prometheus.HistogramOpts{ + Name: "gocached_get_duration_seconds", + Help: "wall time of each cache get request, labeled by storage path: hot (hot tier disk), disk, inline, pending (put-queue), none (HEAD request), or error; and type: get (hit), miss (404), or error", + Buckets: reqLatencyBuckets, + }, []string{"storage", "type"}) + srv.putDuration = prometheus.NewHistogramVec(prometheus.HistogramOpts{ + Name: "gocached_put_duration_seconds", + Help: "wall time of each cache put request, labeled by storage path: disk, inline, or error; and type: put (success), dup, or error", + Buckets: reqLatencyBuckets, + }, []string{"storage", "type"}) + // Copies into the main blob directory and metadata flushes are + // background work against possibly remote storage; buckets span 1ms + // to about a minute. + pipelineBuckets := prometheus.ExponentialBuckets(0.001, 2, 16) + srv.putCopyDuration = prometheus.NewHistogramVec(prometheus.HistogramOpts{ + Name: "gocached_put_queue_copy_duration_seconds", + Help: "wall time of each attempt to write a spooled blob's cleanup intent and copy it into the main blob directory, labeled by result: ok or error", + Buckets: pipelineBuckets, + }, []string{"result"}) + srv.putFlushDuration = prometheus.NewHistogram(prometheus.HistogramOpts{ + Name: "gocached_put_queue_flush_duration_seconds", + Help: "wall time of each put-queue metadata batch transaction, including waiting for the SQLite write lock; divide by gocached_put_queue_flushed_items per flush to see whether SQLite commit throughput is what bounds the flusher", + Buckets: pipelineBuckets, + }) + blobSizeBuckets := []float64{1} + for i := 6; i <= 30; i += 2 { + bucket := 1 << i + blobSizeBuckets = append(blobSizeBuckets, float64(bucket)) + } + srv.blobSize = prometheus.NewHistogramVec(prometheus.HistogramOpts{ + Name: "gocached_blob_size_bytes", + Help: "object size in bytes transmitted on the wire (whether compressed or uncompressed) for successful GETs and PUTs, labeled by storage: hot, disk, inline, or error; and type: get, put, or dup", + Buckets: blobSizeBuckets, + }, []string{"storage", "type"}) + + // Fill the namespace ID cache. + rows, err := srv.db.Query("SELECT NamespaceID, Namespace FROM Namespaces") + if err != nil { + return err + } + defer rows.Close() + for rows.Next() { + var id int64 + var ns string + if err := rows.Scan(&id, &ns); err != nil { + return err + } + srv.namespaces[Namespace(ns)] = id + } + if err := rows.Err(); err != nil { + return err + } + + if srv.globalGeneration < 0 { + return fmt.Errorf("global generation %d must not be negative", srv.globalGeneration) + } + if srv.globalGeneration > 1 { + // Later generations of the global namespace live under a synthetic + // name in the Namespaces table. The "!" prefix is outside the + // character set validateNamespace permits, so no mapping-provided + // namespace can collide with it. + id, err := srv.resolveNamespaceID(Namespace(fmt.Sprintf("!global:%d", srv.globalGeneration))) + if err != nil { + return fmt.Errorf("resolving global generation namespace: %w", err) + } + srv.globalNamespaceID = id + } + + reg := prometheus.NewRegistry() + reg.MustRegister( + collectors.NewGoCollector(), + collectors.NewProcessCollector(collectors.ProcessCollectorOpts{}), + collectors.NewBuildInfoCollector(), + ) + srv.registerMetrics(reg) + reg.MustRegister(srv.shardScanDuration, srv.getDuration, srv.putDuration, srv.blobSize, + srv.putCopyDuration, srv.putFlushDuration) + + // Per-scrape gauges for shard stats loop health. GaugeFunc recomputes on + // every Prometheus scrape, so the values stay fresh between scans + // (an expvar.Int gauge would only update once per shard scan, and + // "oldest age" would lag the wall clock between scans). + reg.MustRegister(prometheus.NewGaugeFunc( + prometheus.GaugeOpts{ + Name: "gocached_shard_stats_unscanned", + Help: "number of shards with no cached usage stats; nonzero means the stats loop has not yet covered the full shard space (e.g. just after first deploy of this code)", + }, + func() float64 { + srv.shardStatsMu.Lock() + defer srv.shardStatsMu.Unlock() + return float64(srv.numShards() - len(srv.shardStats)) + }, + )) + reg.MustRegister(prometheus.NewGaugeFunc( + prometheus.GaugeOpts{ + Name: "gocached_shard_stats_oldest_age_seconds", + Help: "age in seconds of the oldest cached shard scan, or 0 if no shard has been scanned yet; alert when this exceeds shardStalenessTarget by some margin to catch the stats loop falling behind", + }, + func() float64 { + _, oldest, _ := srv.shardFreshness(srv.now()) + if oldest.IsZero() { + return 0 + } + return srv.now().Sub(oldest).Seconds() + }, + )) + // Dead-reckoning gauges: PUTs/evictions adjust these between shard + // scans, so they show how much the live state has drifted from the + // persisted aggregate. Signed: a sustained negative delta means evictions + // are running ahead of writes, which is fine; a sustained large positive + // delta hints the stats loop is falling behind write traffic. + reg.MustRegister(prometheus.NewGaugeFunc( + prometheus.GaugeOpts{ + Name: "gocached_pending_blob_count", + Help: "signed pending Action-count delta since the last shard scan; combined with gocached_blob_count this gives the live total", + }, + func() float64 { + c, _ := srv.sumShardDeltas() + return float64(c) + }, + )) + reg.MustRegister(prometheus.NewGaugeFunc( + prometheus.GaugeOpts{ + Name: "gocached_pending_blob_bytes", + Help: "signed pending aggregate-bytes delta since the last shard scan; combined with gocached_blob_bytes this gives the live total used to gate maxSize cleanup", + }, + func() float64 { + _, b := srv.sumShardDeltas() + return float64(b) + }, + )) + + reg.MustRegister(prometheus.NewGaugeFunc( + prometheus.GaugeOpts{ + Name: "gocached_put_queue_pending", + Help: "PUTs accepted but whose metadata hasn't been committed to SQLite yet; should return to zero shortly after a write burst", + }, + func() float64 { + c, _ := srv.putq.pendingStats() + return float64(c) + }, + )) + reg.MustRegister(prometheus.NewGaugeFunc( + prometheus.GaugeOpts{ + Name: "gocached_put_queue_pending_bytes", + Help: "sum of the reserved sizes of pending spooled PUTs; bounds spool disk usage and reflects backpressure toward its cap", + }, + func() float64 { + _, b := srv.putq.pendingStats() + return float64(b) + }, + )) + // The caps and the stage backlogs are exported so dashboards and alerts + // can show utilization and locate the bottleneck without knowing the + // server's constants. + for _, g := range []struct { + name, help string + fn func() float64 + }{ + {"gocached_put_queue_pending_cap", "capacity of the spooled lane in pending PUTs; compare with gocached_put_queue_pending minus gocached_put_queue_pending_inline", + func() float64 { return putQueuePendingCountCap }}, + {"gocached_put_queue_pending_bytes_cap", "capacity of the spooled lane in reserved bytes; compare with gocached_put_queue_pending_bytes", + func() float64 { return float64(srv.putq.spoolCap) }}, + {"gocached_put_queue_pending_inline", "subset of gocached_put_queue_pending that are inline (small) PUTs in the inline lane, which waits only on the metadata flusher and never on the movers", + func() float64 { return float64(srv.putq.pendingInline()) }}, + {"gocached_put_queue_pending_inline_cap", "capacity of the inline lane in pending PUTs; compare with gocached_put_queue_pending_inline", + func() float64 { return putQueuePendingInlineCap }}, + {"gocached_put_queue_mover_backlog", "spooled blobs waiting for a mover to start copying them into the main blob directory; nonzero while the movers are the constraint", + func() float64 { return float64(len(srv.putq.moverCh)) }}, + {"gocached_put_queue_flush_backlog", "settled PUTs waiting for the metadata flusher; nonzero while SQLite commits are the constraint", + func() float64 { return float64(len(srv.putq.flushCh)) }}, + } { + reg.MustRegister(prometheus.NewGaugeFunc(prometheus.GaugeOpts{Name: g.name, Help: g.help}, g.fn)) + } + + if srv.hot != nil { + reg.MustRegister(prometheus.NewGaugeFunc( + prometheus.GaugeOpts{ + Name: "gocached_hot_bytes", + Help: "bytes currently stored in the hot tier directory; bounded by the configured hot capacity", + }, + func() float64 { return float64(srv.hot.usageBytes()) }, + )) + reg.MustRegister(prometheus.NewGaugeFunc( + prometheus.GaugeOpts{ + Name: "gocached_hot_files", + Help: "number of blob files currently stored in the hot tier directory", + }, + func() float64 { return float64(srv.hot.len()) }, + )) + reg.MustRegister(prometheus.NewGaugeFunc( + prometheus.GaugeOpts{ + Name: "gocached_hot_ready", + Help: "1 once the hot tier startup scan has completed and new blobs may be written into the hot tier, else 0; alert if this stays 0 well past startup", + }, + func() float64 { + if srv.hot.ready.Load() { + return 1 + } + return 0 + }, + )) + reg.MustRegister(prometheus.NewGaugeFunc( + prometheus.GaugeOpts{ + Name: "gocached_hot_scan_running_seconds", + Help: "how long the hot tier startup scan has currently been running, in seconds; 0 when the scan is not running", + }, + func() float64 { + ns := srv.hot.scanStart.Load() + if ns == 0 { + return 0 + } + return time.Since(time.Unix(0, ns)).Seconds() + }, + )) + } + + srv.metricsHandler = promhttp.HandlerFor(reg, promhttp.HandlerOpts{ + ErrorLog: log.Default(), + }) + + // Seed shardStats + the aggregate from BlobShardStats so usageStats is + // accurate from the first HTTP request after restart, without a full-table + // scan. Stale shards are refreshed by runShardStatsLoop in the + // background; the cleanup loop already runs periodically and will catch + // up any over-limit state once enough shards have been rescanned. + if err := srv.loadShardStats(srv.shutdownCtx); err != nil { + srv.logf("loading persisted shard stats: %v", err) + } + srv.logf("gocached: %d/%d shards loaded; usage: %v of limit %v", + len(srv.shardStats), srv.numShards(), srv.lastUsage.Load().All(), bytesFmt(srv.maxSize)) + + if len(srv.jwtIssuers) > 0 { + srv.jwtValidator = ijwt.NewJWTValidator(srv.logf, gocachedAudience, srv.jwtIssuers) + if err := srv.jwtValidator.RunUpdateJWKSLoop(srv.shutdownCtx); err != nil { + return fmt.Errorf("failed to fetch JWKS for JWT validator: %w", err) + } + + for _, iss := range srv.jwtIssuers { + srv.logf("gocached: using JWT issuer %q", iss) + } + + go srv.runCleanSessionsLoop() + } + + if !srv.disableBackgroundLoops { + srv.putq.start(srv.shutdownCtx) + go srv.runCleanLoop() + go srv.runCheckpointLoop() + go srv.runDBSizeMetricsLoop() + go srv.runShardStatsLoop() + } + + return nil +} + +func (srv *Server) registerMetrics(reg *prometheus.Registry) { + rv := reflect.ValueOf(&srv.m).Elem() + t := reflect.TypeOf(&srv.m).Elem() + for i := 0; i < t.NumField(); i++ { + sf := t.Field(i) + if sf.Type == reflect.TypeFor[expvar.Int]() { + expvarInt := rv.Field(i).Addr().Interface().(*expvar.Int) + typ := sf.Tag.Get("type") + name := sf.Tag.Get("name") + if typ == "" { + panic("missing type tag for " + sf.Name) + } + if name == "" { + panic("missing name tag for " + sf.Name) + } + help := sf.Tag.Get("help") + metricName := "gocached_" + name + + if tag := sf.Tag.Get("type"); tag != "" { + if tag == "gauge" { + reg.MustRegister(singleMetricCollector{&expvarGaugeMetric{ + desc: prometheus.NewDesc(metricName, help, nil, nil), + v: expvarInt, + }}) + } else if tag == "counter" { + reg.MustRegister(singleMetricCollector{&expvarCounterMetric{ + desc: prometheus.NewDesc(metricName, help, nil, nil), + v: expvarInt, + }}) + } + } + } + } +} + +// ServerOption configures a gocached Server. +type ServerOption func(*Server) + +// WithShutdownCtx sets the context used to signal server shutdown. Defaults to +// context.Background(). +func WithShutdownCtx(ctx context.Context) ServerOption { + return func(srv *Server) { + srv.shutdownCtx = ctx + } +} + +// WithDir sets the directory where the server stores its data. Defaults to the +// OS user cache directory under $XDG_CACHE_HOME/gocached or equivalent. +func WithDir(dir string) ServerOption { + return func(srv *Server) { + srv.dir = dir + } +} + +// WithSQLiteDir sets the directory where the server stores its SQLite +// database. Defaults to the blob directory (see [WithDir]). +func WithSQLiteDir(dir string) ServerOption { + return func(srv *Server) { + srv.sqliteDir = dir + } +} + +// WithHotDir enables storage tiering, using dir as a fast storage tier (e.g. +// local NVMe) holding a bounded copy of recently used blobs. The main blob +// directory (see [WithDir]) remains the source of truth containing all blobs; +// files in the hot directory can be deleted at any time. If set, +// [WithHotCapacity] must also be set. Defaults to empty, meaning no tiering. +func WithHotDir(dir string) ServerOption { + return func(srv *Server) { + srv.hotDir = dir + } +} + +// WithHotCapacity sets the maximum number of bytes stored in the hot tier +// directory (see [WithHotDir]). When usage exceeds this capacity, the least +// recently used hot files are deleted. +func WithHotCapacity(bytes int64) ServerOption { + return func(srv *Server) { + srv.hotCap = bytes + } +} + +// WithPutSpoolCapacity sets the byte capacity of the put queue's spooled +// lane: the total declared size of accepted PUTs whose blobs may be waiting +// on local disk for a copy into the main blob directory. PUTs beyond it wait +// (backpressure). The spool lives on the hot tier's filesystem when tiering +// is enabled (see [WithHotDir]) and must fit there alongside the hot +// capacity. Zero, the default, uses one eighth of the spool filesystem's +// size, or 16 GiB if that can't be determined. +func WithPutSpoolCapacity(bytes int64) ServerOption { + return func(srv *Server) { + srv.putSpoolCap = bytes + } +} + +// WithVerbose enables verbose logging for the server. Defaults to false. +func WithVerbose(verbose bool) ServerOption { + return func(srv *Server) { + srv.verbose = verbose + } +} + +// WithLogf sets a custom logging function for the server. Defaults to +// [log.Printf]. +func WithLogf(logf logger.Logf) ServerOption { + return func(srv *Server) { + srv.logf = logf + } +} + +// WithMaxSize sets the maximum size of the cache in bytes. Defaults to 0, which +// means no limit. +func WithMaxSize(maxSize int64) ServerOption { + return func(srv *Server) { + srv.maxSize = maxSize + } +} + +// WithMaxAge sets the maximum age of objects in the cache. Objects older than +// this duration will be cleaned periodically. Defaults to 0, which means no +// limit. +func WithMaxAge(maxAge time.Duration) ServerOption { + return func(srv *Server) { + srv.maxAge = maxAge + } +} + +// WithShardPrefixLen sets the number of hex characters used as each shard key +// in BlobShardStats. The shard count is 16^n. Valid values are 1..4 inclusive +// (16, 256, 4096, or 65536 shards). Passing 0 means "use the default" (2; +// 256 shards), which keeps individual shard scans fast at hundreds of +// millions of rows. Tune up only if a single shard scan becomes too slow +// under continued growth; any other value causes [NewServer] to fail. +func WithShardPrefixLen(n int) ServerOption { + return func(srv *Server) { + srv.shardPrefixLen = n + } +} + +// WithGlobalGeneration sets the generation number of the global namespace. +// Generations 0 and 1 both mean the original global namespace; incrementing +// the generation effectively wipes the global namespace without any O(n) +// deletion work: reads and writes move to a fresh namespace, and the previous +// generation's entries stop being accessed, so they age out of the cache via +// the normal LRU eviction (which requires [WithMaxSize] or [WithMaxAge] to be +// set). Decrementing back to a previous generation makes that generation's +// surviving entries visible again. +func WithGlobalGeneration(gen int) ServerOption { + return func(srv *Server) { + srv.globalGeneration = gen + } +} + +// Namespace identifies a logical partition of the cache where each peer is +// equally trusted. A [WithNamespaceMapping] function decides which single +// namespace (if any) a session may write to and which additional namespace (if +// any) it may read from, on top of the always-readable [GlobalNamespace]. It +// may only contain characters from the set [a-zA-Z0-9._~:/@+|=-]. +type Namespace string + +// GlobalNamespace is the namespace shared by all peers. Every authenticated +// session may read from it; a session may write to it only if its mapping +// function grants it as the write namespace. When JWT auth is enabled without a +// mapping function, every session both reads and writes GlobalNamespace. +// Which storage namespace it maps to is controlled by [WithGlobalGeneration]. +const GlobalNamespace Namespace = "" + +// WithJWTAuth enables JWT-based authentication for the server. Each issuer +// must be a reachable HTTP(S) server that serves its JWKS via a URL +// discoverable at /.well-known/openid-configuration. JWTs presented for token +// exchange must pass the standard signature/issuer/audience/expiry checks +// against one of these issuers. If [WithNamespaceMapping] is provided, then +// it may still be rejected if the mapping function returns an error for its +// claims. No requests other than token exchange are allowed without +// authentication. It may be called multiple times; issuers accumulate. +func WithJWTAuth(issuers ...string) ServerOption { + return func(srv *Server) { + srv.jwtIssuers = append(srv.jwtIssuers, issuers...) + } +} + +// NamespaceGrant is the access decision returned by a [WithNamespaceMapping] +// function. Every authenticated session may read from [GlobalNamespace]; these +// fields grant access beyond that. +type NamespaceGrant struct { + // WriteNamespace, if non-nil, is the single namespace the session may write + // to. A nil WriteNamespace yields a read-only session whose PUT requests are + // rejected with 403. + WriteNamespace *Namespace + // ExtraReadNamespace, if non-nil, is one additional namespace the session may + // read from, beyond the always-readable [GlobalNamespace]. + ExtraReadNamespace *Namespace +} + +// WithNamespaceMapping sets the function that makes policy decisions based on +// a JWT's claims. It is called once per token exchange after the JWT's +// signature and standard claims have been validated. It should return an error +// if the claims are not authorized; the token exchange then fails with 401. +// +// Otherwise it returns a [NamespaceGrant] describing the session's access. See +// [Namespace] for character set constraints and [NamespaceGrant] for the +// meaning of its fields. Check claims["iss"] to switch on per-issuer rules. +// +// If JWT auth is enabled but no mapping function is provided, every session +// both reads and writes [GlobalNamespace]. +func WithNamespaceMapping(fn func(ctx context.Context, claims map[string]any) (*NamespaceGrant, error)) ServerOption { + return func(srv *Server) { + srv.namespaceMapping = fn + } +} + +// NewServer creates and starts a new gocached [Server] that is ready to serve +// requests. It defaults to requiring no authentication and storing its data in +// the OS user cache directory under $XDG_CACHE_HOME/gocached or equivalent. +func NewServer(opts ...ServerOption) (*Server, error) { + srv := &Server{ + shutdownCtx: context.Background(), + logf: log.Printf, + sessions: make(map[string]*sessionData), + namespaces: make(map[Namespace]int64), + clock: time.Now, + } + for _, opt := range opts { + opt(srv) + } + + if len(srv.jwtIssuers) > 0 && srv.namespaceMapping == nil { + // If JWT auth is enabled, but not namespace mapping, every session reads + // and writes the global namespace. + srv.namespaceMapping = func(ctx context.Context, claims map[string]any) (*NamespaceGrant, error) { + return &NamespaceGrant{ + WriteNamespace: new(GlobalNamespace), + }, nil + } + } + + srv.shutdownCtx, srv.shutdownCancel = context.WithCancel(srv.shutdownCtx) + err := srv.start() + if err != nil { + // start may have already opened the SQLite database, and leaking that handle + // keeps the file open. + srv.Close() + return nil, err + } + + return srv, nil +} + +// Close shuts down the server, stopping background goroutines and closing the +// database. +func (srv *Server) Close() error { + srv.shutdownCancel() + + // Settle any PUTs the background pipeline hadn't finished: their + // metadata exists only in memory until flushed. + var err error + if drainErr := srv.drainPendingPuts(); drainErr != nil { + err = fmt.Errorf("draining pending puts: %w", drainErr) + } + + srv.mu.Lock() + defer srv.mu.Unlock() + + if srv.updateAccessTimeStmt != nil { + err = errors.Join(err, srv.updateAccessTimeStmt.Close()) + } + if srv.writeConn != nil { + err = errors.Join(err, srv.writeConn.Close()) + } + + if srv.db == nil { + return err + } + + // Final TRUNCATE checkpoint so the WAL doesn't linger on disk past + // shutdown. Use context.Background because srv.shutdownCtx has already + // been canceled. + ckCtx, ckCancel := context.WithTimeout(context.Background(), 2*time.Minute) + if _, _, _, ckErr := srv.checkpointTruncate(ckCtx); ckErr != nil { + err = errors.Join(err, fmt.Errorf("final wal_checkpoint: %w", ckErr)) + } + ckCancel() + + return errors.Join(err, srv.db.Close()) +} + +// Server implements a gocached server. Use [NewServer] to create and start a +// valid instance. +// +// PUT requests are acknowledged as soon as their bytes are spooled to the +// local fast disk (or memory, for inline-sized objects); a background +// pipeline (see [putQueue]) then copies blobs into the main blob directory +// and commits metadata to SQLite in batched transactions. GETs consult the +// pending PUTs before the database, so acknowledged writes are always +// readable; a crash loses only unflushed PUTs, which is acceptable for a +// cache. +type Server struct { + db *sql.DB + dir string // for large blobs (and the SQLite DB, unless sqliteDir is set) + sqliteDir string // for the SQLite DB; if empty, defaults to dir + hotDir string // if non-empty, fast storage tier for a bounded copy of hot blobs + hotCap int64 // maximum bytes in hotDir; must be positive if hotDir is set + hot *hotIndex + putq *putQueue + putSpoolCap int64 // byte capacity of the put queue's spooled lane; 0 means derive from the spool filesystem + verbose bool + logf logger.Logf + clock func() time.Time // if non-nil, alternate time.Now for testing + metricsHandler http.Handler + maxSize int64 // maximum size of the cache in bytes; 0 means no limit + maxAge time.Duration // maximum age of objects; 0 means no limit + shutdownCtx context.Context + shutdownCancel context.CancelFunc + + // globalGeneration is the generation number of the global namespace, set + // via [WithGlobalGeneration]. Values 0 and 1 both mean the original + // global namespace. + globalGeneration int + + // globalNamespaceID is the NamespaceID that [GlobalNamespace] resolves + // to. It is 0 for globalGeneration <= 1; for later generations it is the + // ID of a synthetic per-generation row in the Namespaces table, computed + // once in start. + globalNamespaceID int64 + + jwtValidator *ijwt.Validator // nil unless jwtIssuers is non-empty + jwtIssuers []string // accepted issuer URLs + namespaceMapping func(ctx context.Context, claims map[string]any) (*NamespaceGrant, error) // required when jwtIssuers is non-empty + + mu sync.RWMutex // guards following fields in this block + sessions map[string]*sessionData // maps access token -> session data. + accessDirty map[actionKey]int64 // action -> accessTime + accessFlushTimer *time.Timer // nil if no flush is scheduled + namespaces map[Namespace]int64 // cached namespace string -> NamespaceID + + // sqliteWriteMu serializes access to SQLite. In theory the SQLite driver + // should serialize access with our 5000ms busy timeout, but empirically we + // sometimes seen DB busy errors. Just serialize it explicitly out of + // laziness for now. + // + // Lock ordering: sqliteWriteMu before mu. putQueue.mu is a leaf: + // nothing else is acquired while holding it. + sqliteWriteMu sync.Mutex + writeConn *sql.Conn // nil until first used; single connection for writes + updateAccessTimeStmt *sql.Stmt // nil until first used; for updating access times on writeConn + + // lastUsage holds the most recent aggregate of per-shard stats, recomputed + // after every shard scan. It is the source of truth for /usage, + // BlobCount/BlobBytes gauges, and cleanup decisions; nothing on the hot + // path runs a full-table scan. + lastUsage atomic.Pointer[usageStats] + + // durs is the set of cohort cutoffs (in standardDurs order, plus maxAge if + // set and not already standard) used by all per-shard scans. Computed once + // in start. Persisted shard stats whose cohort set doesn't match these are + // discarded at load time and re-scanned by the stats loop. + durs []time.Duration + + // shardPrefixLen is the number of hex characters in each shard key. Set + // via [WithShardPrefixLen]; defaults to [defaultShardPrefixLen] when 0, + // and must be in [minShardPrefixLen, maxShardPrefixLen] inclusive once + // [Server.start] has run. The derived shard count is 16^shardPrefixLen, + // exposed via [Server.numShards]. + shardPrefixLen int + + // disableBackgroundLoops, when true, skips the start of every periodic + // goroutine (shard stats, cleanup, checkpoint, DB-size metrics). + // Test-only via withoutBackgroundLoops because the mocked clock collides + // scannedAt across concurrent scans; production never sets this. + disableBackgroundLoops bool + + // shardStatsMu guards shardStats and serializes recomputeAggregateLocked. + shardStatsMu sync.Mutex + shardStats map[shardPrefix]*shardSnapshot + + // shardDeltas is dead-reckoning state: each entry tracks (count, bytes) + // added or removed for that shard since its last persisted scan, so + // cleanupTick can react to write bursts faster than shardStalenessTarget. + // PUTs increment; evictions decrement; scanAndPersistShard subtracts the + // (newStats - oldStats) diff so the delta only carries changes that + // aren't already in the persisted shard. Allocated once in start + // (len == numShards()) so per-shard updates are an addressable struct + // with its own mutex; no allocation per PUT. + shardDeltas []shardDelta + + // shardScanDuration observes the wall time of each per-shard SQL scan. + // Combined with the count of shards, the histogram shows whether any one + // shard scan is pathologically slow (e.g. a hot range with millions of + // entries while others are nearly empty). Registered manually in start + // because it isn't an expvar.Int and so doesn't fit the m-struct + // reflection. + shardScanDuration prometheus.Histogram + + // getDuration, putDuration, and blobSize observe the wall time and transfer + // size of each /action GET and PUT request handler, labeled by storage and + // type. type is one of "get" (hit), "miss" (404), "put", "dup", or "error". + // storage is one of "disk", "inline", "none", or "error". + getDuration *prometheus.HistogramVec + putDuration *prometheus.HistogramVec + blobSize *prometheus.HistogramVec + + // putCopyDuration and putFlushDuration observe the put queue's two + // background stages: copying a spooled blob into the main blob + // directory (labeled by result, ok or error) and committing a metadata + // batch to SQLite. They are nil in tests that build a bare Server. + putCopyDuration *prometheus.HistogramVec + putFlushDuration prometheus.Histogram + + // Metrics. Exported fields for reflection, but within a private struct + // field to control the gocached Server API surface. + m struct { + ActiveGets expvar.Int `type:"gauge" name:"active_gets" help:"currently pending get requests; should usually be zero"` + ActivePuts expvar.Int `type:"gauge" name:"active_puts" help:"currently pending put requests; should usually be zero"` + Gets expvar.Int `type:"counter" name:"gets" help:"total number of gocache get requests"` // gets = getHits + getErrs + implicit misses + GetBytes expvar.Int `type:"counter" name:"get_bytes" help:"total bytes fetched from gocache gets that were cache hits"` + GetHits expvar.Int `type:"counter" name:"get_hits" help:"total number of successful gocache get requests"` + GetAccessBumps expvar.Int `type:"counter" name:"get_access_bumps" help:"number of times a get request updated the access time of object"` + GetHitsInline expvar.Int `type:"counter" name:"get_hits_inline" help:"cache hits served from inline database storage (small objects)"` + GetErrs expvar.Int `type:"counter" name:"get_errs" help:"number of gocache get request errors"` + Puts expvar.Int `type:"counter" name:"puts" help:"total number of gocache put requests"` + PutErrs expvar.Int `type:"counter" name:"put_errs" help:"number of gocache put request errors"` + PutsDup expvar.Int `type:"counter" name:"puts_dup" help:"total number of gocache put requests that are duplicates of a mapping we already had"` + PutsBytes expvar.Int `type:"counter" name:"put_bytes" help:"total bytes added from gocache puts"` + PutsInline expvar.Int `type:"counter" name:"put_inline" help:"subset of gocached_puts that were stored inline (small objects)"` + BlobCount expvar.Int `type:"gauge" name:"blob_count" help:"number of blobs currently stored in the cache"` + BlobBytes expvar.Int `type:"gauge" name:"blob_bytes" help:"sum of blob sizes currently stored in the cache"` + EvictedActions expvar.Int `type:"counter" name:"evicted_actions" help:"number of Actions evicted from the cache by the cleanup loop"` + EvictedBlobs expvar.Int `type:"counter" name:"evicted_blobs" help:"number of Blobs evicted from the cache; a Blob is evicted when its last referencing Action is evicted"` + EvictedBytes expvar.Int `type:"counter" name:"evicted_bytes" help:"number of bytes reclaimed by evicting Blobs from the cache"` + Sessions expvar.Int `type:"gauge" name:"sessions" help:"number of active authenticated sessions"` + Auths expvar.Int `type:"counter" name:"auth_attempts" help:"number of successful token exchanges"` + AuthErrs expvar.Int `type:"counter" name:"auth_errs" help:"number of failed token exchanges"` + + SQLiteDataBytes expvar.Int `type:"gauge" name:"sqlite_data_bytes" help:"size in bytes of the SQLite main database file on disk"` + SQLiteWALBytes expvar.Int `type:"gauge" name:"sqlite_wal_bytes" help:"size in bytes of the SQLite WAL file on disk; should stay bounded near walJournalSizeLimit"` + + HotHits expvar.Int `type:"counter" name:"hot_hits" help:"GETs served from the hot tier"` + HotMisses expvar.Int `type:"counter" name:"hot_misses" help:"GETs that were not served from the hot tier"` + HotReadErrs expvar.Int `type:"counter" name:"hot_read_errs" help:"hot tier opens that failed for a reason other than the file not existing; reads fall back to the main blob directory"` + HotPromotions expvar.Int `type:"counter" name:"hot_promotions" help:"blobs copied into the hot tier after a hot tier miss"` + HotPromotionErrs expvar.Int `type:"counter" name:"hot_promotion_errs" help:"failed attempts to copy a blob into the hot tier"` + HotWriteErrs expvar.Int `type:"counter" name:"hot_write_errs" help:"failed hot tier writes during PUT; the PUT itself still succeeds"` + HotEvicted expvar.Int `type:"counter" name:"hot_evicted" help:"files evicted from the hot tier to stay under its capacity"` + HotEvictedBytes expvar.Int `type:"counter" name:"hot_evicted_bytes" help:"bytes reclaimed by evicting files from the hot tier"` + + PutQueueBlocked expvar.Int `type:"counter" name:"put_queue_blocked" help:"PUT requests that had to wait for put-queue backpressure before being admitted; see the per-lane put_queue_blocked_* counters for which cap bound"` + PutQueueBlockedInline expvar.Int `type:"counter" name:"put_queue_blocked_inline" help:"PUT requests that waited on the inline lane's count cap; the lane drains only through the metadata flusher, so this growing means SQLite commit throughput is the constraint"` + PutQueueBlockedSpoolBytes expvar.Int `type:"counter" name:"put_queue_blocked_spool_bytes" help:"PUT requests that waited on the spooled lane's byte cap; raise the spool capacity if the disk has room, or look at the movers if pending never drains"` + PutQueueBlockedSpoolCount expvar.Int `type:"counter" name:"put_queue_blocked_spool_count" help:"PUT requests that waited on the spooled lane's count cap; many small blobs are piling up faster than the movers copy them"` + PutQueueCopyErrs expvar.Int `type:"counter" name:"put_queue_copy_errs" help:"failed attempts to copy a spooled blob into the main blob directory; retried before the PUT is dropped"` + PutQueueDropped expvar.Int `type:"counter" name:"put_queue_dropped" help:"pending PUTs abandoned after repeated copy or flush failures; the client saw success but the object was lost"` + PutQueueFlushes expvar.Int `type:"counter" name:"put_queue_flushes" help:"metadata batch transactions committed by the put-queue flusher"` + PutQueueFlushedItems expvar.Int `type:"counter" name:"put_queue_flushed_items" help:"pending PUTs whose metadata was committed by the put-queue flusher"` + PutQueueFlushDups expvar.Int `type:"counter" name:"put_queue_flush_dups" help:"subset of put_queue_flushed_items that were duplicates of an already-stored action"` + PutQueueOrphansSwept expvar.Int `type:"counter" name:"put_queue_orphans_swept" help:"blobs deleted from the main blob directory by the cleanup sweeper because a crashed or dropped PUT left them without SQLite accounting"` + } +} + +// sessionData corresponds to a specific access token, and is only used if JWT +// auth is enabled. +type sessionData struct { + expiry time.Time // Session valid until. + + canWrite bool // Whether this session may write at all. + writeNamespaceID int64 // NamespaceID this session writes to (the current global generation's ID for GlobalNamespace). It is only meaningful if canWrite is true. + extraReadNamespaceID int64 // One namespace this session may read from besides the always-readable global namespace. 0 means only global. + + // These fields are for debug display purposes only. + writeNamespace Namespace // The Namespace this session writes to, if any. + extraReadNamespace Namespace // The extra Namespace this session reads from, if any. + claims map[string]any // Claims from the JWT used to create this session. + + mu sync.Mutex // Guards stats. + stats stats +} + +// scope returns the OAuth 2.0 scope string (RFC 8693 §2.2.1 / RFC 6749 §3.3: +// space-delimited, case-sensitive) advertising the operations this session may +// perform so a client can skip requests that would be denied. It is advisory: +// the server still enforces access on every request, so a client that ignores +// it is only less efficient, never wrong. +// +// action and object permissions are currently always assigned together, but +// may be split in future. Object writes can be self-verified as they are +// content-addressed, so require less trust in the client. +func (d *sessionData) scope() string { + s := []string{"action:read", "object:read"} + if d.canWrite { + s = append(s, "action:write", "object:write") + } + return strings.Join(s, " ") +} + +// stats holds per-request or per-session stats which get rolled up into server +// stats. See [server] struct for detailed definitions. +type stats struct { + LastUsed time.Time // Only applies to session stats. Last time the access token for this session was used. + Gets int64 + GetBytes int64 + GetHits int64 + GetAccessBumps int64 + GetHitsInline int64 + GetNanos int64 + GetErrs int64 + Puts int64 + PutErrs int64 + PutsDup int64 + PutsBytes int64 + PutsInline int64 + PutsNanos int64 +} + +func (srv *Server) now() time.Time { + return srv.clock() +} + +// ServeHTTPDebug serves debug HTTP endpoints. It is unauthenticated, so should +// only be used on a separate debug listener. +func (srv *Server) ServeHTTPDebug(w http.ResponseWriter, r *http.Request) { + if srv.verbose { + srv.logf("ServeHTTPDebug: %s %s", r.Method, r.RequestURI) + } + switch { + case r.URL.Path == "/": + w.Header().Set("Content-Type", "text/html; charset=utf-8") + io.WriteString(w, "

gocached

") + io.WriteString(w, "

This is a shared Go build cache server, hit by GOCACHEPROG clients.

") + io.WriteString(w, "

See /usage for usage stats.

") + io.WriteString(w, "

See /sessions for session data

") + io.WriteString(w, "

See /metrics for Prometheus metrics.

") + io.WriteString(w, "

See /debug/pprof/ for pprof

") + io.WriteString(w, "

See /debug/pprof/goroutine?debug=2 - full goroutines

") + case r.URL.Path == "/usage": + srv.serveUsage(w, r) + case r.URL.Path == "/sessions": + srv.serveSessions(w, r) + case r.URL.Path == "/metrics": + srv.metricsHandler.ServeHTTP(w, r) + case strings.HasPrefix(r.URL.Path, "/debug/pprof/profile"): + pprof.Profile(w, r) + case strings.HasPrefix(r.URL.Path, "/debug/pprof/cmdline"): + pprof.Cmdline(w, r) + case strings.HasPrefix(r.URL.Path, "/debug/pprof/symbol"): + pprof.Symbol(w, r) + case strings.HasPrefix(r.URL.Path, "/debug/pprof/trace"): + pprof.Trace(w, r) + case strings.HasPrefix(r.URL.Path, "/debug/pprof/"): + pprof.Index(w, r) + default: + http.Error(w, "not found", http.StatusNotFound) + } +} + +// ServeHTTP implements gocached's API via [http.Handler]. +func (srv *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) { + if srv.verbose { + srv.logf("ServeHTTP: %s %s", r.Method, r.RequestURI) + } + + var sessionData *sessionData // remains nil for unauthenticated requests. + reqStats := &stats{} + defer func() { + // Call inside func to capture maybe-updated sessionData pointer. + srv.processRequestStats(reqStats, sessionData) + }() + + // Handle session auth first if enabled. + if srv.jwtValidator != nil { + // If JWT auth enabled, this is the only unauthenticated (non-debug) endpoint. + if r.Method == "POST" && r.URL.Path == "/auth/exchange-token" { + srv.handleTokenExchange(w, r) + return + } + + // Check for session data and error if none. + token := strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer ") + if !strings.HasPrefix(token, tokenPrefix) { + http.Error(w, "unauthorized", http.StatusUnauthorized) + return + } + + var ok bool + sessionData, ok = srv.getSessionData(token) + if !ok || srv.now().After(sessionData.expiry) { + if srv.verbose { + reason := fmt.Sprintf("exists: %v", ok) + if sessionData != nil { + reason += fmt.Sprintf(", expiry: %v", sessionData.expiry) + } + srv.logf("unauthorized; %s", reason) + } + http.Error(w, "unauthorized", http.StatusUnauthorized) + return + } + } + + if r.Method == "PUT" { + writeNS := srv.globalNamespaceID + if sessionData != nil { + if !sessionData.canWrite { + http.Error(w, "forbidden: read-only session", http.StatusForbidden) + return + } + writeNS = sessionData.writeNamespaceID + } + srv.handlePut(w, r, reqStats, writeNS) + return + } + if r.Method != "GET" && r.Method != "HEAD" { + http.Error(w, "bad method", http.StatusBadRequest) + return + } + if strings.HasPrefix(r.URL.Path, "/action/") { + srv.handleGetAction(w, r, reqStats, sessionData) + return + } + if sessionData != nil && r.URL.Path == "/session/stats" { + srv.handleSessionStats(w, sessionData) + return + } + http.Error(w, "not found", http.StatusNotFound) +} + +func (srv *Server) getSessionData(token string) (*sessionData, bool) { + srv.mu.Lock() + defer srv.mu.Unlock() + sessionData, ok := srv.sessions[token] + return sessionData, ok +} + +func (srv *Server) addSessionData(token string, sessionData *sessionData) { + srv.mu.Lock() + defer srv.mu.Unlock() + srv.sessions[token] = sessionData + srv.m.Sessions.Add(1) +} + +func (srv *Server) processRequestStats(req *stats, sessionData *sessionData) { + srv.m.Gets.Add(req.Gets) + srv.m.GetBytes.Add(req.GetBytes) + srv.m.GetHits.Add(req.GetHits) + srv.m.GetAccessBumps.Add(req.GetAccessBumps) + srv.m.GetHitsInline.Add(req.GetHitsInline) + srv.m.GetErrs.Add(req.GetErrs) + srv.m.Puts.Add(req.Puts) + srv.m.PutErrs.Add(req.PutErrs) + srv.m.PutsDup.Add(req.PutsDup) + srv.m.PutsBytes.Add(req.PutsBytes) + srv.m.PutsInline.Add(req.PutsInline) + + if sessionData != nil { + sessionData.mu.Lock() + defer sessionData.mu.Unlock() + + sessionData.stats.LastUsed = srv.now().UTC() + sessionData.stats.Gets += req.Gets + sessionData.stats.GetBytes += req.GetBytes + sessionData.stats.GetHits += req.GetHits + sessionData.stats.GetAccessBumps += req.GetAccessBumps + sessionData.stats.GetHitsInline += req.GetHitsInline + sessionData.stats.GetErrs += req.GetErrs + sessionData.stats.GetNanos += req.GetNanos + sessionData.stats.Puts += req.Puts + sessionData.stats.PutErrs += req.PutErrs + sessionData.stats.PutsDup += req.PutsDup + sessionData.stats.PutsBytes += req.PutsBytes + sessionData.stats.PutsInline += req.PutsInline + sessionData.stats.PutsNanos += req.PutsNanos + } +} + +func getHexSuffix(r *http.Request, prefix string) (hexSuffix string, ok bool) { + hexSuffix, _ = strings.CutPrefix(r.RequestURI, prefix) + if !validHex(hexSuffix) { + return "", false + } + return hexSuffix, true +} + +// actionKey is the comparable value type for the (NamespaceID, ActionID) +// primary key tuple used in the SQLite Actions table. +type actionKey struct { + NamespaceID int64 // Server.globalNamespaceID for global + ActionID string +} + +func validHex(x string) bool { + if len(x) < 4 || len(x) > 1000 || len(x)%2 == 1 { + return false + } + for i := range x { + b := x[i] + if b >= '0' && b <= '9' || b >= 'a' && b <= 'f' { + continue + } + return false + } + return true +} + +// relAtimeSeconds is how old an access time needs to be before +// we do a DB write to update it. +const relAtimeSeconds = 60 * 60 * 24 // 1 day + +// getFromNamespace fetches an action row from the global namespace or the +// session's one extra read namespace, preferring a hit in global so the shared +// cache mtime is bumped with higher priority than namespaced entries. The +// parameters are (globalID, extraID, actionID, globalID); pass the global ID +// again as the extra namespace when there is none, so the IN clause collapses +// to global only. +const getFromNamespace = ` +SELECT b.SHA256, b.StoredSize, b.UncompressedSize, b.SmallData, a.AltOutputID, a.AccessTime, a.NamespaceID +FROM Actions a, Blobs b +WHERE a.NamespaceID IN (?, ?) + AND a.ActionID = ? + AND a.BlobID = b.BlobID +ORDER BY CASE a.NamespaceID WHEN ? THEN 0 ELSE 1 END +LIMIT 1 +` + +func (srv *Server) handleGetAction(w http.ResponseWriter, r *http.Request, stats *stats, sessionData *sessionData) { + srv.m.ActiveGets.Add(1) + defer srv.m.ActiveGets.Add(-1) + + start := srv.now() + labels := writeObjectResponseLabels{storage: "error", result: "error"} + defer func() { + d := srv.now().Sub(start) + stats.GetNanos += d.Nanoseconds() + srv.getDuration.WithLabelValues(labels.storage, labels.result).Observe(d.Seconds()) + }() + stats.Gets++ + ctx := r.Context() + + httpErr := func(msg string, code int) { + http.Error(w, msg, code) + stats.GetErrs++ + } + + actionID, ok := getHexSuffix(r, "/action/") + if !ok { + httpErr("bad request", http.StatusBadRequest) + return + } + if r.Header.Get("Want-Object") != "1" { + httpErr("bad request: missing Want-Object header", http.StatusBadRequest) + return + } + + // A session may read from its one extra namespace in addition to the + // always-readable global namespace. Unauthenticated sessions (or those + // with no extra grant) read global only, expressed by using the global ID + // as the extra ID so the SQL IN clause collapses to global. + extraReadID := srv.globalNamespaceID + if sessionData != nil && sessionData.extraReadNamespaceID != 0 { + extraReadID = sessionData.extraReadNamespaceID + } + + // Serve from the put-queue first: a just-written PUT is readable here + // before its metadata reaches SQLite. Like the SQL query below, global is + // preferred over the extra namespace. + pp, havePending := srv.putq.lookup(actionKey{NamespaceID: srv.globalNamespaceID, ActionID: actionID}) + if !havePending && extraReadID != srv.globalNamespaceID { + pp, havePending = srv.putq.lookup(actionKey{NamespaceID: extraReadID, ActionID: actionID}) + } + if havePending { + // Open the spool file before committing to this path: if the entry + // was retired between the lookup and now, its file is gone but its + // metadata is committed, so fall through to the SQL path below. + var rc io.ReadCloser + ok := true + if pp.smallData == nil { + f, err := os.Open(pp.queueFile) + if err != nil { + ok = false + } else { + rc = f + } + } + if ok { + stats.GetHits++ + opened := false + defer func() { + if !opened && rc != nil { + rc.Close() + } + }() + labels = srv.writeObjectResponse(w, r, stats, objectSource{ + sha256hex: pp.sha256hex, + storedSize: pp.storedSize, + uncompressedSize: pp.uncompressedSize, + altOutputID: pp.altOutputID, + inline: pp.smallData != nil, + smallData: pp.smallData, + open: func() (io.ReadCloser, bool, error) { + opened = true + return rc, false, nil + }, + }) + labels.markPending() + return + } + } + + var ( + sha256hex string + storedSize, uncompressedSize int64 + smallData sql.NullString + altObjectID string + accessTime int64 + err error + actionKey = actionKey{ + ActionID: actionID, + } + ) + err = srv.db.QueryRow(getFromNamespace, srv.globalNamespaceID, extraReadID, actionID, srv.globalNamespaceID).Scan( + &sha256hex, &storedSize, &uncompressedSize, &smallData, &altObjectID, &accessTime, &actionKey.NamespaceID) + if err != nil { + if errors.Is(err, sql.ErrNoRows) { + labels = writeObjectResponseLabels{storage: "none", result: "miss"} + http.Error(w, "not found", http.StatusNotFound) + return + } + srv.logf("QueryRow error: %v", err) + httpErr("QueryRow error", http.StatusInternalServerError) + return + } + + if srv.maybeBumpAccessTime(actionKey, accessTime) { + stats.GetAccessBumps++ + } + + stats.GetHits++ + + var sd []byte + if smallData.Valid { + sd = []byte(smallData.String) + } + isLZ4 := storedSize != uncompressedSize + labels = srv.writeObjectResponse(w, r, stats, objectSource{ + sha256hex: sha256hex, + storedSize: storedSize, + uncompressedSize: uncompressedSize, + altOutputID: altObjectID, + inline: smallData.Valid, + smallData: sd, + open: func() (io.ReadCloser, bool, error) { + return srv.getObjectFromDiskOrPeer(ctx, sha256hex, isLZ4) + }, + }) +} + +// objectSource describes a cached object to be written as an HTTP response, +// independent of whether its metadata came from SQLite or elsewhere. +type objectSource struct { + sha256hex string + storedSize int64 // bytes stored (possibly lz4-compressed) + uncompressedSize int64 + altOutputID string // "" if the output ID equals sha256hex + inline bool // whether smallData holds the object bytes + smallData []byte // object bytes if inline + + // open opens a non-inline object for reading. + // A nil ReadCloser with nil error means the object was not found. + open func() (rc io.ReadCloser, fromHot bool, err error) +} + +// writeObjectResponseLabels describes how a [Server.writeObjectResponse] +// call was served. Its fields are the label values reported to the get +// latency metric. +type writeObjectResponseLabels struct { + // storage says where the object bytes came from: "inline" (the + // database), "disk" (the main blob directory), or "hot" (the hot + // tier); "none" for a response without a body (a HEAD request, an + // empty object, or a miss); or "error". Callers serving from the + // put-queue relabel hits as "pending" via markPending. + storage string + + // result is the request outcome: "get" (a hit), "miss" (a 404), or + // "error". + result string +} + +// markPending relabels a hit that carried a body as served from "pending" +// storage, for responses served from the put-queue rather than committed +// storage. Bodyless responses, misses, and errors keep their labels. +func (l *writeObjectResponseLabels) markPending() { + if l.result == "get" && l.storage != "none" { + l.storage = "pending" + } +} + +// writeObjectResponse writes src to w, negotiating lz4 content encoding and +// serving inline data directly. It returns the labels for the get latency +// metric. +func (srv *Server) writeObjectResponse(w http.ResponseWriter, r *http.Request, stats *stats, src objectSource) (labels writeObjectResponseLabels) { + labels = writeObjectResponseLabels{storage: "error", result: "get"} + + outputID := cmp.Or(src.altOutputID, src.sha256hex) + isLZ4 := src.storedSize != src.uncompressedSize + clientAcceptsLZ4 := requestAcceptsEncoding(r, "lz4") + + w.Header().Set("Content-Type", "application/octet-stream") + w.Header().Set("Go-Output-Id", outputID) + + if src.inline || !isLZ4 { + // Inline, empty, or legacy uncompressed: Content-Length is the stored size. + w.Header().Set("Content-Length", fmt.Sprint(src.storedSize)) + } else if isLZ4 && clientAcceptsLZ4 { + // Disk lz4 + client accepts lz4: stream raw compressed file. + w.Header().Set("Content-Encoding", "lz4") + w.Header().Set("Content-Length", fmt.Sprint(src.storedSize)) + w.Header().Set("X-Uncompressed-Length", fmt.Sprint(src.uncompressedSize)) + } else { + // Disk lz4 + client doesn't accept lz4: decompress for client. + w.Header().Set("Content-Length", fmt.Sprint(src.uncompressedSize)) + } + + if r.Method == "HEAD" || (src.storedSize == 0 && src.uncompressedSize == 0) { + labels.storage = "none" + return + } + + if src.inline { + // For small outputs stored inline in the database, we can return them directly. + stats.GetHitsInline++ + stats.GetBytes += src.storedSize + labels.storage = "inline" + srv.blobSize.WithLabelValues(labels.storage, labels.result).Observe(float64(src.storedSize)) + w.Write(src.smallData) + return + } + + // Otherwise, for large objects that we know about, we can try to get them + // from our local disk or a peer. + + rc, fromHot, err := src.open() + if err != nil { + srv.logf("Get object error: %v", err) + stats.GetErrs++ + http.Error(w, "Get object error", http.StatusInternalServerError) + labels = writeObjectResponseLabels{storage: "error", result: "error"} + return + } + if rc == nil { + // Our database suggested we should've had this object, + // but maybe somebody delete it by hand from the filesystem. + // Just treat it as a cache miss. The background cleanup + // will eventually remove the Action row from the DB + // after identifying it as a dangling reference. + http.Error(w, "not found", http.StatusNotFound) + labels = writeObjectResponseLabels{storage: "none", result: "miss"} + return + } + defer rc.Close() + labels.storage = "disk" + if fromHot { + labels.storage = "hot" + } + + if isLZ4 && !clientAcceptsLZ4 { + // Client doesn't accept lz4; decompress on the fly. + stats.GetBytes += src.uncompressedSize + srv.blobSize.WithLabelValues(labels.storage, labels.result).Observe(float64(src.uncompressedSize)) + io.Copy(w, lz4.NewReader(rc)) + } else { + stats.GetBytes += src.storedSize + srv.blobSize.WithLabelValues(labels.storage, labels.result).Observe(float64(src.storedSize)) + io.Copy(w, rc) + } + return +} + +// maybeBumpAccessTime reports whether it enqueued an access time bump for the +// given actionKey, if the provided prior access time (unix seconds) is old +// enough to warrant an update. +func (srv *Server) maybeBumpAccessTime(actionKey actionKey, priorAccessTimeUnixSec int64) (didBump bool) { + now := srv.now().Unix() + if priorAccessTimeUnixSec > now-relAtimeSeconds { + return false + } + // If it's been more than a day since the last access, update the access time. + // This is similar to the Linux "relatime" behavior. + return srv.enqueueAccessTimeBump(actionKey) +} + +// accessBatchSizeSoftLimit is the size of the accessDirty map at which +// we stop kicking the timer can down the road and let the thing flush, +// considering it large enough. +const accessBatchSizeSoftLimit = 1000 + +func (srv *Server) enqueueAccessTimeBump(action actionKey) (changed bool) { + srv.mu.Lock() + defer srv.mu.Unlock() + if _, ok := srv.accessDirty[action]; ok { + // Another caller already bumped this from old (over realTimeSeconds) + // to recent, so let that caller win, even if it's a few seconds old. + // Those few seconds don't matter much for cleaning purposes, and it's + // better for stats to only have one caller report the bump. + return false + } + if srv.accessDirty == nil { + srv.accessDirty = make(map[actionKey]int64) + } + srv.accessDirty[action] = srv.now().Unix() + if srv.accessFlushTimer == nil { + srv.accessFlushTimer = time.AfterFunc(5*time.Second, srv.flushAccessTimeBumps) + } else if len(srv.accessDirty) < accessBatchSizeSoftLimit { + srv.accessFlushTimer.Reset(5 * time.Second) + } + return true +} + +// flushAccessTimeBumps writes any pending access time updates to the database. +// +// It doesn't return an error to it can be easily used as a time.AfterFunc +// callback. +func (srv *Server) flushAccessTimeBumps() { + srv.flushAccessTimeBumpsWithErr() +} + +// flushAccessTimeBumpsWithErr is the same as flushAccessTimeBumps but returns +// any error that's used only for the internal rescheduling defer logic. +func (srv *Server) flushAccessTimeBumpsWithErr() (ret error) { + srv.sqliteWriteMu.Lock() // before srv.mu + defer srv.sqliteWriteMu.Unlock() + + srv.mu.Lock() + defer srv.mu.Unlock() + + defer func() { + if ret != nil { + srv.logf("flushAccessTimeBumps error: %v", ret) + srv.accessFlushTimer = time.AfterFunc(5*time.Second, srv.flushAccessTimeBumps) + } + }() + + srv.accessFlushTimer = nil + + if len(srv.accessDirty) == 0 { + return nil + } + + if srv.writeConn == nil { + var err error + srv.writeConn, err = srv.db.Conn(context.Background()) + if err != nil { + return fmt.Errorf("getting writeConn: %w", err) + } + } + + if srv.updateAccessTimeStmt == nil { + var err error + srv.updateAccessTimeStmt, err = srv.writeConn.PrepareContext(context.Background(), "UPDATE Actions SET AccessTime = ? WHERE NamespaceID = ? AND ActionID = ?") + if err != nil { + return fmt.Errorf("prepare updateAccessTimeStmt: %w", err) + } + } + + tx, err := srv.writeConn.BeginTx(context.Background(), nil) + if err != nil { + return fmt.Errorf("begin Tx: %w", err) + } + defer tx.Rollback() + + txStmt := tx.Stmt(srv.updateAccessTimeStmt) + + for action, accessTime := range srv.accessDirty { + _, err := txStmt.Exec(accessTime, action.NamespaceID, action.ActionID) + if err != nil { + return fmt.Errorf("updating access time for ns=%d,action=%q: %w", action.NamespaceID, action.ActionID, err) + } + } + if err := tx.Commit(); err != nil { + return fmt.Errorf("commit: %w", err) + } + + srv.accessDirty = nil + return nil +} + +// getObjectFromDiskOrPeer retrieves the object for the given actionID, either +// from disk or a peer. This is used after a local DB lookup discovers the content +// exists but is not stored in SQLite. +// +// If isLZ4 is true, the .lz4 suffixed path is opened; otherwise the plain path. +// +// With tiering enabled, the hot tier is consulted first; fromHot reports +// whether the returned reader is backed by the hot tier. A hot tier miss that +// hits the main directory kicks off an asynchronous promotion into the hot +// tier. +// +// It returns (nil, false, nil) on miss. +func (srv *Server) getObjectFromDiskOrPeer(_ context.Context, sha256hex string, isLZ4 bool) (rc io.ReadCloser, fromHot bool, err error) { + if len(sha256hex) != sha256.Size*2 { + return nil, false, fmt.Errorf("invalid sha256hex %q", sha256hex) + } + name := sha256hex + if isLZ4 { + name += ".lz4" + } + if srv.hot != nil { + f, err := os.Open(srv.hotFilepath(name)) + if err == nil { + srv.hot.touch(name) + srv.m.HotHits.Add(1) + return f, true, nil + } + srv.m.HotMisses.Add(1) + if !os.IsNotExist(err) { + srv.m.HotReadErrs.Add(1) + srv.logf("hot tier: opening %v: %v", name, err) + } + } + diskPath := filepath.Join(srv.dir, sha256hex[:2], name) + f, err := os.Open(diskPath) + if err != nil { + if os.IsNotExist(err) { + // TODO(bradfitz): search peers, S3, etc. + // For now, just return a miss. + return nil, false, nil + } + return nil, false, err + } + if srv.hot != nil { + // The index may still hold an entry if the hot file vanished out from + // under us (e.g. a wiped ephemeral disk); drop it so the stale entry + // doesn't block re-promotion and usage accounting stays honest. + srv.hot.remove(name) + if srv.hot.tryStartPromotion(name) { + go srv.promoteToHot(name) + } + } + return f, false, nil +} + +// promoteToHot copies the blob file with the given base filename from the +// main blob directory into the hot tier. It runs in its own goroutine after a +// hot tier miss; the caller must have registered the promotion via +// [hotIndex.tryStartPromotion]. +func (srv *Server) promoteToHot(name string) { + defer srv.hot.endPromotion(name) + if srv.shutdownCtx.Err() != nil { + return + } + size, err := srv.copyToHot(name) + if err != nil { + srv.logf("hot tier: promoting %v: %v", name, err) + srv.m.HotPromotionErrs.Add(1) + return + } + srv.hot.add(name, size) + srv.m.HotPromotions.Add(1) + srv.evictHotIfOver() +} + +// copyToHot copies the named blob file from the main blob directory to a temp +// file in the hot dir and atomically renames it into place, returning the +// file size. The bytes are copied verbatim (in already-compressed form, if +// applicable), so no re-hashing or re-compression is needed. +func (srv *Server) copyToHot(name string) (size int64, err error) { + src, err := os.Open(filepath.Join(srv.dir, name[:2], name)) + if err != nil { + return 0, err + } + defer src.Close() + + tf, err := os.CreateTemp(srv.hotDir, fmt.Sprintf("upload-%d-*", srv.now().Unix())) + if err != nil { + return 0, err + } + defer func() { + if err != nil { + tf.Close() + os.Remove(tf.Name()) + } + }() + + size, err = io.Copy(tf, src) + if err != nil { + return 0, err + } + if err := tf.Close(); err != nil { + return 0, err + } + target := srv.hotFilepath(name) + if err := os.MkdirAll(filepath.Dir(target), 0750); err != nil { + return 0, err + } + if err := os.Rename(tf.Name(), target); err != nil { + return 0, err + } + return size, nil +} + +// hotEvictTargetPct is the percentage of the hot tier capacity that eviction +// shrinks usage down to once the capacity is exceeded. Evicting slightly past +// the limit reduces churn from evicting one file at a time on every write. +const hotEvictTargetPct = 95 + +// evictHotIfOver deletes the least recently used hot tier files if the hot +// tier is over its capacity. It's a no-op when tiering is disabled, the +// startup scan hasn't yet measured usage, or usage is within bounds. Hot +// files are just copies, so deletion never loses data. +func (srv *Server) evictHotIfOver() { + if srv.hot == nil || !srv.hot.ready.Load() || srv.hot.usageBytes() <= srv.hotCap { + return + } + for _, ent := range srv.hot.evictLRU(srv.hotCap * hotEvictTargetPct / 100) { + if err := os.Remove(srv.hotFilepath(ent.name)); err != nil && !os.IsNotExist(err) { + srv.logf("hot tier: evicting %v: %v", ent.name, err) + } + srv.m.HotEvicted.Add(1) + srv.m.HotEvictedBytes.Add(ent.size) + } +} + +// removeFromHot deletes both possible hot tier files (plain and .lz4) for the +// given SHA256 hex string and drops them from the hot index. It is called +// when a blob is evicted from the cache entirely. The files are removed even +// if the index doesn't know them (e.g. the startup scan hasn't reached them +// yet). Errors are logged only: hot files are disposable copies, so a +// leftover is at worst wasted space that eviction later reclaims. +func (srv *Server) removeFromHot(sha256Hex string) { + if srv.hot == nil { + return + } + for _, name := range []string{sha256Hex, sha256Hex + ".lz4"} { + srv.hot.remove(name) + if err := os.Remove(srv.hotFilepath(name)); err != nil && !os.IsNotExist(err) { + srv.logf("hot tier: removing %v: %v", name, err) + } + } +} + +const ( + // putMinUploadRate is the minimum acceptable PUT upload rate, in bytes + // per second, used to scale each PUT's body read deadline by its + // declared size. Clients are same-region, same-VPC AWS instances that + // normally upload orders of magnitude faster, but they're guest VMs + // speaking TLS and may be CPU over-subscribed, so this is deliberately + // slack: the point is to have some bound reaping dead or crawling + // senders (which would otherwise pin their put-queue reservations for + // the multi-minute TCP keepalive timeout), not a tight one. + putMinUploadRate = 1 << 20 + + // putUploadGrace is the flat allowance added to every PUT's body read + // deadline, covering per-request overhead independent of size. + putUploadGrace = 30 * time.Second + + // putMaxBodyReadTime caps a PUT's body read deadline, keeping the + // duration math sane for absurd declared sizes. At putMinUploadRate it + // permits uploads into the tens of GB, far past what a Go build cache + // stores. + putMaxBodyReadTime = time.Hour +) + +// putBodyReadDeadline returns the wall-clock deadline for reading a PUT body +// of the declared size, arriving no slower than putMinUploadRate after an +// initial putUploadGrace. +func putBodyReadDeadline(now time.Time, contentLength int64) time.Time { + // Cap in units of seconds, before the Duration multiply can overflow + // for absurd declared sizes. + secs := contentLength / putMinUploadRate + if secs > int64(putMaxBodyReadTime/time.Second) { + return now.Add(putMaxBodyReadTime) + } + return now.Add(min(putUploadGrace+time.Duration(secs)*time.Second, putMaxBodyReadTime)) +} + +func (s *Server) handlePut(w http.ResponseWriter, r *http.Request, stats *stats, namespaceID int64) { + s.m.ActivePuts.Add(1) + defer s.m.ActivePuts.Add(-1) + + start := s.now() + storage := "error" + result := "error" + defer func() { + d := s.now().Sub(start) + stats.PutsNanos += d.Nanoseconds() + s.putDuration.WithLabelValues(storage, result).Observe(d.Seconds()) + }() + + if r.Method != "PUT" { + http.Error(w, "bad method", http.StatusMethodNotAllowed) + return + } + actionID, outputID, ok := strings.Cut(r.RequestURI[len("/"):], "/") + if !ok || !validHex(actionID) || !validHex(outputID) { + http.Error(w, "bad URI", http.StatusBadRequest) + return + } + if r.ContentLength == -1 { + http.Error(w, "missing Content-Length", http.StatusBadRequest) + return + } + + // Backpressure: reserve queue room for the declared size before reading + // any of the body. This blocks when the background pipeline is behind, + // and aborts if the client goes away while waiting. + reserved, err := s.putq.reserve(r.Context(), r.ContentLength) + if err != nil { + stats.PutErrs++ + http.Error(w, "canceled while awaiting queue room", http.StatusServiceUnavailable) + return + } + // The reservation is handed off to the pending entry on enqueue; until + // then, any early return must give it back. + handedOff := false + defer func() { + if !handedOff { + s.putq.unreserve(reserved) + } + }() + + // Bound the body read (best effort) so a dead or crawling client can't + // pin its reservation: set after the backpressure wait so queueing time + // doesn't count against the upload, on the wall clock (not s.now, whose + // test clock isn't meaningful to the network stack), and cleared on + // return so the deadline doesn't outlive this request on a reused + // connection. + rc := http.NewResponseController(w) + rc.SetReadDeadline(putBodyReadDeadline(time.Now(), r.ContentLength)) + defer rc.SetReadDeadline(time.Time{}) + + hasher := sha256.New() + hashingBody := io.TeeReader(r.Body, hasher) + + storedSize := r.ContentLength + uncompressedSize := r.ContentLength + + var smallData []byte + var queueFile string + if r.ContentLength <= smallObjectSize { + // Small objects are held in memory and eventually stored inline in + // the database. + smallData, err = io.ReadAll(hashingBody) + if err != nil { + s.logf("Read content error: %v", err) + http.Error(w, "Read content error", http.StatusInternalServerError) + return + } + if int64(len(smallData)) != r.ContentLength { + // This check is redundant with net/http's validation, but + // for extra clarity. + http.Error(w, "bad content length", http.StatusInternalServerError) + return + } + } else { + // Larger objects are spooled (lz4 compressed) to the put-queue + // directory on the local disk; the background movers copy them + // into the main blob directory. + diskSize, path, err := s.putq.spoolBlob(r.ContentLength, hashingBody) + if err != nil { + s.logf("Spool blob error: %v", err) + stats.PutErrs++ + http.Error(w, "Spool blob error", http.StatusInternalServerError) + return + } + storedSize = diskSize + queueFile = path + } + + sha256hex := fmt.Sprintf("%x", hasher.Sum(nil)) + + altOutputID := "" + if sha256hex != outputID { + altOutputID = outputID + } + p := &pendingPut{ + key: actionKey{NamespaceID: namespaceID, ActionID: actionID}, + sha256hex: sha256hex, + storedSize: storedSize, + uncompressedSize: uncompressedSize, + altOutputID: altOutputID, + createTime: s.now().Unix(), + smallData: smallData, + queueFile: queueFile, + reservation: reserved, + } + + if s.putq.enqueue(p) { + // An entry for this action is already pending; the first PUT wins, + // as it would at insert time in the database. Duplicates that are + // already committed to SQLite aren't detected here; the flusher + // discovers and counts those later. + if queueFile != "" { + os.Remove(queueFile) + } + stats.PutsDup++ + result = "dup" + } else { + handedOff = true + result = "put" + } + + stats.Puts++ + stats.PutsBytes += r.ContentLength + if smallData != nil { + stats.PutsInline++ + storage = "inline" + } else { + storage = "disk" + } + s.blobSize.WithLabelValues(storage, result).Observe(float64(r.ContentLength)) + + w.WriteHeader(http.StatusNoContent) +} + +// handleTokenExchange handles POST /auth/exchange-token requests to exchange +// a JWT for an access token. Each access token represents a session that will +// last for one hour and have cache stats associated with it. +func (srv *Server) handleTokenExchange(w http.ResponseWriter, r *http.Request) { + var req struct { + JWT string `json:"jwt"` + } + // JWTs are often sent in HTTP headers, so 4KiB should ~always be enough. + if err := json.NewDecoder(http.MaxBytesReader(w, r.Body, 4<<10)).Decode(&req); err != nil { + srv.m.AuthErrs.Add(1) + http.Error(w, "bad request: "+err.Error(), http.StatusBadRequest) + return + } + + jwtClaims, err := srv.jwtValidator.Validate(r.Context(), req.JWT) + if err != nil { + srv.m.AuthErrs.Add(1) + if srv.verbose { + srv.logf("token exchange: JWT validation error: %v", err) + } + http.Error(w, "unauthorized", http.StatusUnauthorized) + return + } + + grant, err := srv.namespaceMapping(r.Context(), jwtClaims) + if err != nil || grant == nil { + srv.m.AuthErrs.Add(1) + if srv.verbose { + srv.logf("token exchange: namespace func error: %v", err) + } + http.Error(w, "unauthorized", http.StatusUnauthorized) + return + } + + for _, ns := range []*Namespace{grant.WriteNamespace, grant.ExtraReadNamespace} { + if ns == nil { + continue + } + if err := validateNamespace(*ns); err != nil { + srv.m.AuthErrs.Add(1) + // This is a bug in the namespace mapping function, so always log errors. + srv.logf("token exchange: invalid namespace from claims: %v", err) + http.Error(w, "unauthorized", http.StatusUnauthorized) + return + } + } + + const ttl = time.Hour + sd := &sessionData{ + expiry: srv.now().UTC().Add(ttl), + claims: jwtClaims, + } + if grant.WriteNamespace != nil { + sd.canWrite = true + sd.writeNamespace = *grant.WriteNamespace + if sd.writeNamespaceID, err = srv.resolveNamespaceID(*grant.WriteNamespace); err != nil { + srv.m.AuthErrs.Add(1) + srv.logf("token exchange: %v", err) + http.Error(w, "internal error", http.StatusInternalServerError) + return + } + } + if grant.ExtraReadNamespace != nil { + sd.extraReadNamespace = *grant.ExtraReadNamespace + if sd.extraReadNamespaceID, err = srv.resolveNamespaceID(*grant.ExtraReadNamespace); err != nil { + srv.m.AuthErrs.Add(1) + srv.logf("token exchange: %v", err) + http.Error(w, "internal error", http.StatusInternalServerError) + return + } + } + + // 52 base32 characters, 256 bits of entropy. + accessToken := tokenPrefix + strings.ToLower(rand.Text()+rand.Text()) + srv.addSessionData(accessToken, sd) + + resp := map[string]any{ + "access_token": accessToken, + "token_type": "Bearer", + "expires_in": ttl.Seconds(), + "scope": sd.scope(), + } + w.Header().Set("Content-Type", "application/json") + if err := json.NewEncoder(w).Encode(resp); err != nil { + srv.m.AuthErrs.Add(1) + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } + + srv.m.Auths.Add(1) +} + +// namespaceAllowedBytes is the set of non-alphanumeric bytes permitted in a +// namespace. It is chosen to cover characters common in JWT identity claims: +// issuer URLs, emails, and provider-structured "sub" values such as +// "repo:org/repo:environment:prod", "auth0|abc123", and dependabot[bot]. It is +// deliberately ASCII-only so SQLite BINARY comparison matches Go string equality +// byte for byte, avoiding Unicode normalization divergence. +const namespaceAllowedBytes = "._~:/@+|=-[]" + +func validateNamespace(ns Namespace) error { + if ns == GlobalNamespace { + return nil + } + for i := 0; i < len(ns); i++ { + c := ns[i] + switch { + case c >= 'A' && c <= 'Z', c >= 'a' && c <= 'z', c >= '0' && c <= '9': + case strings.IndexByte(namespaceAllowedBytes, c) >= 0: + default: + return fmt.Errorf("namespace contains disallowed byte %#x at index %d", c, i) + } + } + return nil +} + +// resolveNamespaceID returns the integer ID for the given Namespace, +// inserting a row in the Namespaces table if one doesn't already exist. +// GlobalNamespace always resolves to [Server.globalNamespaceID] (0 unless +// [WithGlobalGeneration] is 2+) without touching the database. +func (srv *Server) resolveNamespaceID(ns Namespace) (int64, error) { + if ns == GlobalNamespace { + return srv.globalNamespaceID, nil + } + + // If it's not a new namespace, we only need to consult our cache of IDs. + srv.mu.Lock() + id, ok := srv.namespaces[ns] + srv.mu.Unlock() + if ok { + return id, nil + } + + srv.sqliteWriteMu.Lock() + defer srv.sqliteWriteMu.Unlock() + + srv.mu.Lock() + defer srv.mu.Unlock() + + // Check if we lost a race now that we have both locks. + if id, ok = srv.namespaces[ns]; ok { + return id, nil + } + + err := srv.db.QueryRow(`INSERT INTO Namespaces (Namespace) VALUES (?) + RETURNING NamespaceID;`, ns).Scan(&id) + if err != nil { + return 0, fmt.Errorf("resolving namespace %q: %w", ns, err) + } + + srv.namespaces[ns] = id + + return id, nil +} + +func (srv *Server) handleSessionStats(w http.ResponseWriter, sessionData *sessionData) { + w.Header().Set("Content-Type", "application/json") + if err := json.NewEncoder(w).Encode(sessionData.stats); err != nil { + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } +} + +func (s *Server) sha256Filepath(hash [sha256.Size]byte) string { + hex := fmt.Sprintf("%x", hash) + return filepath.Join(s.dir, hex[:2], hex) +} + +// hotFilepath returns the hot tier path for the given base filename, which is +// either "" or ".lz4". +func (s *Server) hotFilepath(name string) string { + return filepath.Join(s.hotDir, name[:2], name) +} + +type countAndSize struct { + Count int64 // number of actions + Size int64 // total stored size in bytes (after compression, if any) of all actions' blobs, even if shared by other actions +} + +func (cs countAndSize) String() string { + if cs.Count == 0 { + return "0 objects, 0 bytes" + } + return fmt.Sprintf("%d objects, %s", cs.Count, bytesFmt(cs.Size)) +} + +type usageStats struct { + // ActionsLE is a histogram of the actions in the DB by their access time. + // Sizes reflect stored size on disk (after any compression). + // + // The key is a Prometheus-style histogram "less than" value. That is, if + // there are map keys for 24h and 48h, the latter includes the sum of the + // 24h values as well. + // + // The map keys are day-granularity, as the access time is only updated once + // it's over a day old. + // + // So the map keys are 24h, 48h, 96h, 168h (7d), 336h (14d), 720h + // (30d), and 2160h (90d) and math.MaxInt64 for infinity. + ActionsLE map[time.Duration]countAndSize + + // MissingBlobRows is the number of rows in the Actions table that + // reference a BlobID that doesn't exist in the Blobs table. + // This should always be zero in a healthy system. + MissingBlobRows int + + // QueryDuration is how long the per-shard SQL query took. It is set by + // [Server.scanShard] for shard snapshots and is left at the zero value + // in the aggregate returned by [Server.usageStats]. Persisted in JSON + // so /usage can surface slow shards across restarts. + QueryDuration time.Duration +} + +func (us *usageStats) All() countAndSize { return us.ActionsLE[math.MaxInt64] } + +const day = 24 * time.Hour + +var standardDurs = []time.Duration{ + 1 * day, + 2 * day, + 4 * day, + 7 * day, + 14 * day, + 30 * day, + 90 * day, + math.MaxInt64, +} + +// computeDurs calculates histogram buckets for usage stats, given the +// configured maxAge. The math.MaxInt64 sentinel is always included as the +// "no upper bound" bucket. +func computeDurs(maxAge time.Duration) []time.Duration { + if maxAge == 0 || slices.Contains(standardDurs, maxAge) { + return slices.Clone(standardDurs) + } + durs := []time.Duration{maxAge, math.MaxInt64} + for _, d := range standardDurs { + if d < maxAge { + durs = append(durs, d) + } + } + slices.Sort(durs) + return durs +} + +// shardPrefix is the lowercase hex prefix that identifies one SHA256 shard +// (e.g. "00", "ab"). The named type keeps shard-key strings from being +// passed where any string would do — code that means "shard key" reads as +// shardPrefix; SQL parameters and other free-form strings stay plain string. +type shardPrefix string + +// shardSnapshot is the cached state for a single SHA256-prefix shard. +type shardSnapshot struct { + stats *usageStats + scannedAt time.Time +} + +// shardDelta is the dead-reckoning state for one shard. count and bytes are +// signed: PUTs add positive values, evictions add negative ones. +// scanAndPersistShard subtracts the change in the persisted view of this +// shard (newStats - oldStats) so anything that just landed in the aggregate +// leaves the delta. +// +// The per-shard mu serializes the count+bytes pair so a reader can't catch +// a PUT mid-update with count incremented but bytes not (or vice versa). +// Storing the values inline avoids allocating a new struct on every PUT. +type shardDelta struct { + mu sync.Mutex + count int64 + bytes int64 +} + +// addBlobDelta adjusts the shardDelta for the shard owning sha256hex by the +// given signed (count, bytes) values. Callers pass +1, +storedSize on a +// successful new Action insert, and -1, -storedSize when evicting an Action. +// Bogus sha256hex (too short or non-hex prefix) is silently ignored. +func (srv *Server) addBlobDelta(sha256hex string, count, bytes int64) { + if srv.shardPrefixLen <= 0 || len(sha256hex) < srv.shardPrefixLen { + return + } + n, err := strconv.ParseUint(sha256hex[:srv.shardPrefixLen], 16, 64) + if err != nil || int(n) >= len(srv.shardDeltas) { + return + } + sd := &srv.shardDeltas[n] + sd.mu.Lock() + sd.count += count + sd.bytes += bytes + sd.mu.Unlock() +} + +// sumShardDeltas returns the sum of every shard's pending delta. The caller +// adds this to the persisted aggregate to get a live estimate of total cache +// occupancy between scans. +func (srv *Server) sumShardDeltas() (count, bytes int64) { + for i := range srv.shardDeltas { + sd := &srv.shardDeltas[i] + sd.mu.Lock() + count += sd.count + bytes += sd.bytes + sd.mu.Unlock() + } + return +} + +// numShards returns 16^shardPrefixLen, the total number of SHA256-prefix +// shards used to partition usage statistics. +func (srv *Server) numShards() int { + return 1 << (4 * srv.shardPrefixLen) +} + +// shardPrefix returns the hex prefix for shard index i, zero-padded to +// shardPrefixLen characters (e.g. 0 -> "00", 255 -> "ff"). +func (srv *Server) shardPrefix(i int) shardPrefix { + return shardPrefix(fmt.Sprintf("%0*x", srv.shardPrefixLen, i)) +} + +// shardRange returns the [lo, hi) SHA256 hex range covered by shard index i. +// For the final shard, hi is a sentinel string of 'g' characters: 'g' sorts +// after every valid hex digit, so "SHA256 < hi" excludes nothing in that +// range. Using a sentinel keeps the SQL uniform (always two parameters). +// Returns plain strings since both ends bind directly into SQL. +func (srv *Server) shardRange(i int) (lo, hi string) { + lo = string(srv.shardPrefix(i)) + if i+1 == srv.numShards() { + hi = strings.Repeat("g", srv.shardPrefixLen) + } else { + hi = string(srv.shardPrefix(i + 1)) + } + return +} + +// usageStats returns the most recent aggregate of per-shard stats. It does no +// SQL work: the aggregate is maintained by the shard stats loop and seeded +// from the BlobShardStats table at startup, so callers see usage immediately +// after process restart rather than after a multi-minute full-table scan. The +// returned value is never nil; ActionsLE may be empty if no shard has been +// scanned and no persisted stats exist. +func (s *Server) usageStats() (*usageStats, error) { + if us := s.lastUsage.Load(); us != nil { + return us, nil + } + return &usageStats{ActionsLE: make(map[time.Duration]countAndSize)}, nil +} + +// scanShard runs a single SQL-aggregated query that computes the cohort +// histogram for Actions whose Blob's SHA256 falls in shard idx's range. +// The aggregation happens entirely in SQL via SUM(CASE WHEN ...) clauses, so +// the driver returns one row regardless of how many Actions fall in the shard. +// This is what makes a scan cheap enough to pin a reader snapshot for only +// milliseconds per shard rather than minutes for the whole table. +func (srv *Server) scanShard(ctx context.Context, idx int) (*usageStats, error) { + durs := srv.durs + now := srv.now().Unix() + + var sb strings.Builder + args := make([]any, 0, 2*len(durs)+2) + sb.WriteString("SELECT ") + finiteCount := 0 + for _, d := range durs { + if d == math.MaxInt64 { + continue + } + cutoff := now - int64(d/time.Second) + if finiteCount > 0 { + sb.WriteByte(',') + } + sb.WriteString("COALESCE(SUM(CASE WHEN a.AccessTime > ? THEN 1 ELSE 0 END),0),") + sb.WriteString("COALESCE(SUM(CASE WHEN a.AccessTime > ? THEN b.StoredSize ELSE 0 END),0)") + args = append(args, cutoff, cutoff) + finiteCount++ + } + sb.WriteString(",COUNT(*),COALESCE(SUM(b.StoredSize),0)") + sb.WriteString(" FROM Blobs b JOIN Actions a ON a.BlobID = b.BlobID") + sb.WriteString(" WHERE b.SHA256 >= ? AND b.SHA256 < ?") + lo, hi := srv.shardRange(idx) + args = append(args, lo, hi) + + vals := make([]int64, finiteCount*2+2) + dest := make([]any, len(vals)) + for i := range vals { + dest[i] = &vals[i] + } + start := time.Now() + if err := srv.db.QueryRowContext(ctx, sb.String(), args...).Scan(dest...); err != nil { + return nil, err + } + queryDuration := time.Since(start) + if srv.shardScanDuration != nil { + // Nil-check: scanShard is reachable from tests that build a bare + // Server without calling start (e.g. TestShardPrefix). + srv.shardScanDuration.Observe(queryDuration.Seconds()) + } + + st := &usageStats{ + ActionsLE: make(map[time.Duration]countAndSize, len(durs)), + QueryDuration: queryDuration, + } + j := 0 + for _, d := range durs { + if d == math.MaxInt64 { + continue + } + st.ActionsLE[d] = countAndSize{Count: vals[j*2], Size: vals[j*2+1]} + j++ + } + st.ActionsLE[math.MaxInt64] = countAndSize{Count: vals[finiteCount*2], Size: vals[finiteCount*2+1]} + return st, nil +} + +// shardCohortsMatch reports whether the persisted ActionsLE map keys exactly +// match the current srv.durs set. A mismatch means the persisted shard was +// written with a different maxAge (or a different code version), so we +// discard it and force the stats loop to re-scan that shard. +func shardCohortsMatch(le map[time.Duration]countAndSize, durs []time.Duration) bool { + if len(le) != len(durs) { + return false + } + for _, d := range durs { + if _, ok := le[d]; !ok { + return false + } + } + return true +} + +// loadShardStats reads every BlobShardStats row at startup, validates that +// the persisted cohort set matches srv.durs, and seeds shardStats + the +// aggregate. Rows with mismatched cohorts are dropped so the stats loop +// re-scans them. The aggregate becomes the source of truth for /usage and the +// cleanup loop until the first new shard scan completes. +func (srv *Server) loadShardStats(ctx context.Context) error { + rows, err := srv.db.QueryContext(ctx, "SELECT Prefix, ScannedAt, StatsJSON FROM BlobShardStats") + if err != nil { + return err + } + defer rows.Close() + srv.shardStatsMu.Lock() + defer srv.shardStatsMu.Unlock() + srv.shardStats = make(map[shardPrefix]*shardSnapshot, srv.numShards()) + for rows.Next() { + var prefix, statsJSON string + var scannedAt int64 + if err := rows.Scan(&prefix, &scannedAt, &statsJSON); err != nil { + return err + } + if len(prefix) != srv.shardPrefixLen { + // A previous run used a different shardPrefixLen. The stats loop + // won't write to these keys, so they'd otherwise sit in the + // table forever contributing stale data to the aggregate. + srv.logf("BlobShardStats[%q] prefix length %d != current %d; will rescan", + prefix, len(prefix), srv.shardPrefixLen) + continue + } + var st usageStats + if err := json.Unmarshal([]byte(statsJSON), &st); err != nil { + srv.logf("BlobShardStats[%q] JSON decode: %v; will rescan", prefix, err) + continue + } + if !shardCohortsMatch(st.ActionsLE, srv.durs) { + srv.logf("BlobShardStats[%q] cohort mismatch; will rescan", prefix) + continue + } + srv.shardStats[shardPrefix(prefix)] = &shardSnapshot{ + stats: &st, + scannedAt: time.Unix(scannedAt, 0), + } + // Replay the persisted scan time into the histogram so the + // /metrics endpoint shows reasonable percentiles immediately after + // restart instead of waiting a full pass over every shard. The observation is + // older than "now" but it's the most accurate signal we have for + // this shard until the stats loop refreshes it. + if srv.shardScanDuration != nil && st.QueryDuration > 0 { + srv.shardScanDuration.Observe(st.QueryDuration.Seconds()) + } + } + if err := rows.Err(); err != nil { + return err + } + srv.recomputeAggregateLocked() + return nil +} + +// upsertShardStats persists one shard's scan result. The row is keyed on +// Prefix so subsequent scans of the same shard overwrite in place. +func (srv *Server) upsertShardStats(ctx context.Context, prefix shardPrefix, scannedAt time.Time, st *usageStats) error { + statsJSON, err := json.Marshal(st) + if err != nil { + return err + } + srv.sqliteWriteMu.Lock() + defer srv.sqliteWriteMu.Unlock() + _, err = srv.db.ExecContext(ctx, + `INSERT INTO BlobShardStats (Prefix, ScannedAt, StatsJSON) VALUES (?, ?, ?) + ON CONFLICT(Prefix) DO UPDATE SET ScannedAt=excluded.ScannedAt, StatsJSON=excluded.StatsJSON`, + prefix, scannedAt.Unix(), string(statsJSON)) + return err +} + +// recomputeAggregateLocked sums the cohort histograms across every cached +// shard, stores the result in lastUsage, and updates the BlobCount/BlobBytes +// gauges. The caller must hold shardStatsMu. +func (srv *Server) recomputeAggregateLocked() { + agg := &usageStats{ActionsLE: make(map[time.Duration]countAndSize, len(srv.durs))} + for _, d := range srv.durs { + agg.ActionsLE[d] = countAndSize{} + } + for _, sh := range srv.shardStats { + for d, cs := range sh.stats.ActionsLE { + was := agg.ActionsLE[d] + was.Count += cs.Count + was.Size += cs.Size + agg.ActionsLE[d] = was + } + } + srv.lastUsage.Store(agg) + total := agg.All() + srv.m.BlobCount.Set(total.Count) + srv.m.BlobBytes.Set(total.Size) +} + +// pickOldestShard returns the shard index whose ScannedAt is oldest, along +// with that ScannedAt. A shard that has never been scanned (no entry in +// shardStats) wins immediately, and is returned with a zero time so the +// caller treats it as infinitely stale (scan it now, no sleep). +func (srv *Server) pickOldestShard() (idx int, scannedAt time.Time) { + srv.shardStatsMu.Lock() + defer srv.shardStatsMu.Unlock() + oldestIdx := -1 + var oldestTime time.Time + for i := range srv.numShards() { + sh, ok := srv.shardStats[srv.shardPrefix(i)] + if !ok { + return i, time.Time{} + } + if oldestIdx == -1 || sh.scannedAt.Before(oldestTime) { + oldestIdx = i + oldestTime = sh.scannedAt + } + } + return oldestIdx, oldestTime +} + +// scanAndPersistShard scans shard idx, writes the result back to +// BlobShardStats, and updates the in-memory cache + aggregate. +func (srv *Server) scanAndPersistShard(ctx context.Context, idx int) error { + prefix := srv.shardPrefix(idx) + scannedAt := srv.now() + stats, err := srv.scanShard(ctx, idx) + if err != nil { + return fmt.Errorf("scan shard %q: %w", prefix, err) + } + if err := srv.upsertShardStats(ctx, prefix, scannedAt, stats); err != nil { + return fmt.Errorf("persist shard %q: %w", prefix, err) + } + + srv.shardStatsMu.Lock() + var oldAll countAndSize + if old, ok := srv.shardStats[prefix]; ok { + oldAll = old.stats.All() + } + srv.shardStats[prefix] = &shardSnapshot{stats: stats, scannedAt: scannedAt} + srv.recomputeAggregateLocked() + srv.shardStatsMu.Unlock() + + // Subtract from the delta the exact change that just landed in the + // persisted view: (newStats - oldStats). Anything in delta whose Action + // commit happened before scanShard's snapshot is now in newStats, so we + // subtract it out; anything that committed after the snapshot isn't in + // newStats, so the diff doesn't touch its delta contribution. PUTs whose + // addBlobDelta hasn't been called yet temporarily push delta negative; + // when the addBlobDelta runs it cancels out. + // + // First-scan caveat: when oldStats was missing (no prior scan for this + // shard AND no persisted row from a previous process run), oldAll is + // zero. If the DB had pre-existing rows that delta never saw (e.g. first + // deploy of this code on a populated DB whose BlobShardStats table was + // empty), subtracting newStats from delta will leave delta with a + // negative offset equal to the pre-existing row count. The aggregate + // becomes correct from then on; the negative offset stays in the delta + // until process restart (which reloads aggregate from the now-populated + // BlobShardStats and starts delta at zero). Operators deploying onto a + // populated DB can avoid this by nuking the DB first. + newAll := stats.All() + diffCount := newAll.Count - oldAll.Count + diffBytes := newAll.Size - oldAll.Size + sd := &srv.shardDeltas[idx] + sd.mu.Lock() + sd.count -= diffCount + sd.bytes -= diffBytes + sd.mu.Unlock() + return nil +} + +// runShardStatsLoop keeps usage stats fresh by rescanning the +// oldest-scanned shard whenever it crosses [shardStalenessTarget]. The +// per-scan cadence in steady state works out to about +// shardStalenessTarget/numShards (~2.3s at defaults) +func (srv *Server) runShardStatsLoop() { + for { + idx, scannedAt := srv.pickOldestShard() + var sleep time.Duration + if !scannedAt.IsZero() { + sleep = max(shardStalenessTarget-srv.now().Sub(scannedAt), shardStatsMinInterval) + } + if sleep > 0 { + select { + case <-srv.shutdownCtx.Done(): + return + case <-time.After(sleep): + } + } + if err := srv.scanAndPersistShard(srv.shutdownCtx, idx); err != nil { + if errors.Is(err, context.Canceled) { + return + } + srv.logf("shard stats loop: %v", err) + } + } +} + +// scanAllShards synchronously rescans every shard. Used by tests to get a +// deterministic aggregate, and by POST /usage as a manual "refresh now" +// trigger when an operator doesn't want to wait for the stats loop. +func (srv *Server) scanAllShards(ctx context.Context) error { + for i := range srv.numShards() { + if err := srv.scanAndPersistShard(ctx, i); err != nil { + return err + } + } + return nil +} + +// shardScanDurationPercentiles returns p25/p50/p90 of the QueryDuration +// across every cached shard, and the count of shards that contributed (i.e. +// have a non-zero QueryDuration). Returns zeros if no shard has been scanned +// yet. Linear interpolation is overkill given numShards is 16/256/etc.; a +// rank-index pick is good enough for a /usage page. +func (srv *Server) shardScanDurationPercentiles() (p25, p50, p90 time.Duration, n int) { + srv.shardStatsMu.Lock() + durs := make([]time.Duration, 0, len(srv.shardStats)) + for _, sh := range srv.shardStats { + if sh.stats.QueryDuration > 0 { + durs = append(durs, sh.stats.QueryDuration) + } + } + srv.shardStatsMu.Unlock() + n = len(durs) + if n == 0 { + return + } + slices.Sort(durs) + pick := func(p int) time.Duration { + i := (n * p) / 100 + if i >= n { + i = n - 1 + } + return durs[i] + } + return pick(25), pick(50), pick(90), n +} + +// shardFreshnessBuckets are the "scanned within last X" age cohorts shown on +// the /usage page. The largest cohort is intentionally larger than +// [shardStalenessTarget] so the page makes it obvious when the stats loop +// has fallen behind (the smaller cohorts drop to zero before the largest does). +var shardFreshnessBuckets = []time.Duration{ + 1 * time.Minute, + 5 * time.Minute, + 15 * time.Minute, + 1 * time.Hour, + 6 * time.Hour, + 24 * time.Hour, +} + +// shardFreshness counts how many cached shards were last scanned within each +// of [shardFreshnessBuckets], and returns the absolute time of the oldest +// scan. The buckets are cumulative ("within last X"), matching the cohort +// histogram style elsewhere in this package. +func (srv *Server) shardFreshness(now time.Time) (counts map[time.Duration]int, oldest time.Time, total int) { + counts = make(map[time.Duration]int, len(shardFreshnessBuckets)) + srv.shardStatsMu.Lock() + defer srv.shardStatsMu.Unlock() + total = len(srv.shardStats) + for _, sh := range srv.shardStats { + age := now.Sub(sh.scannedAt) + for _, b := range shardFreshnessBuckets { + if age <= b { + counts[b]++ + } + } + if oldest.IsZero() || sh.scannedAt.Before(oldest) { + oldest = sh.scannedAt + } + } + return counts, oldest, total +} + +// evictionCandidateQuery is the SQL used by [Server.evictOldestActions] to +// pick which Actions to delete next. INDEXED BY is the documented +// SQLite mechanism (https://sqlite.org/lang_indexedby.html) for locking +// down a plan so a future schema change can't silently regress this query +// from an O(log N) index range scan into an O(N) table scan. The plan is +// verified in TestEvictionQueryPlan. +const evictionCandidateQuery = ` +SELECT NamespaceID, ActionID, BlobID +FROM Actions INDEXED BY idx_actions_access +WHERE AccessTime <= ? +ORDER BY AccessTime ASC +LIMIT ?` + +// evictOldestActions deletes up to maxCount Actions whose AccessTime is at +// or before cutoff, oldest first, in a single short write transaction. For +// each deleted Action, if its Blob has no other referencing Actions +// remaining, the Blob row and its on-disk file are also removed. The batch +// stops early once at least maxBytes have been reclaimed (use math.MaxInt64 +// for no byte budget — e.g., maxAge cleanup, where we want every stale +// Action gone regardless of size). +// +// This is Action-LRU: a stale Action is evicted even if its Blob is shared +// with newer Actions (the Blob then stays alive via those newer refs). That +// differs from the prior Blob-LRU policy, where a Blob was only evicted when +// MAX(AccessTime) of all its Actions was past the cutoff. The cache is +// mostly 1:1 Action-to-Blob so the two policies are equivalent for the +// common case; on shared Blobs, Action-LRU is the stricter LRU semantics. +func (srv *Server) evictOldestActions(ctx context.Context, cutoff int64, maxCount int, maxBytes int64) (countAndSize, error) { + var ret countAndSize + + rows, err := srv.db.QueryContext(ctx, evictionCandidateQuery, cutoff, maxCount) + if err != nil { + return ret, fmt.Errorf("query eviction candidates: %w", err) + } + type candidate struct { + NamespaceID int64 + ActionID string + BlobID int64 + } + var cands []candidate + for rows.Next() { + var c candidate + if err := rows.Scan(&c.NamespaceID, &c.ActionID, &c.BlobID); err != nil { + rows.Close() + return ret, fmt.Errorf("scan candidate: %w", err) + } + cands = append(cands, c) + } + if err := rows.Close(); err != nil { + return ret, fmt.Errorf("close candidate rows: %w", err) + } + if len(cands) == 0 { + return ret, nil + } + + srv.sqliteWriteMu.Lock() + defer srv.sqliteWriteMu.Unlock() + + tx, err := srv.db.BeginTx(ctx, nil) + if err != nil { + return ret, fmt.Errorf("eviction Begin: %w", err) + } + defer tx.Rollback() + + // pendingDelta queues addBlobDelta updates for after Commit. Holding + // them until commit means a rolled-back tx (rare) doesn't poison the + // dead-reckoning state. + type pendingDelta struct { + sha256Hex string + storedSize int64 + } + var pending []pendingDelta + + var evictedBlobs int64 + var evictedBlobBytes int64 // bytes reclaimed (disk space + EvictedBytes metric) + var evictedActionBytes int64 // aggregate-side bytes reduction (drives maxBytes check) + var evictedActions int64 + for _, c := range cands { + // Fetch SHA256 + StoredSize up front so we can both (a) decrement + // the dead-reckoning delta for this Action and (b) remove the disk + // file if the Blob ends up orphaned. A LEFT-JOIN here is just an + // extra query; the row is keyed on the primary key so it's cheap. + var sha256Hex string + var storedSize int64 + blobExists := true + switch err := tx.QueryRowContext(ctx, + "SELECT SHA256, StoredSize FROM Blobs WHERE BlobID = ?", + c.BlobID).Scan(&sha256Hex, &storedSize); { + case err == nil: + case errors.Is(err, sql.ErrNoRows): + blobExists = false + default: + return ret, fmt.Errorf("fetch Blob row: %w", err) + } + + if _, err := tx.ExecContext(ctx, + "DELETE FROM Actions WHERE NamespaceID = ? AND ActionID = ?", + c.NamespaceID, c.ActionID); err != nil { + return ret, fmt.Errorf("delete action: %w", err) + } + evictedActions++ + if !blobExists { + // Orphan Action pointing at a missing Blob; delete the row and + // move on. No delta, no Blob row, no file. + continue + } + // Each Action contributes 1 count + storedSize bytes to the aggregate + // (see scanShard); removing it reduces both by exactly that. Decrement + // deferred until after Commit so a rollback doesn't leave the delta + // off. + pending = append(pending, pendingDelta{sha256Hex, storedSize}) + evictedActionBytes += storedSize + + // Is this Blob now orphaned? idx_actions_blobid makes this O(log N). + var stillReferenced int + if err := tx.QueryRowContext(ctx, + "SELECT EXISTS(SELECT 1 FROM Actions WHERE BlobID = ?)", + c.BlobID).Scan(&stillReferenced); err != nil { + return ret, fmt.Errorf("check orphan: %w", err) + } + if stillReferenced == 1 { + continue + } + if _, err := tx.ExecContext(ctx, "DELETE FROM Blobs WHERE BlobID = ?", c.BlobID); err != nil { + return ret, fmt.Errorf("delete Blob: %w", err) + } + evictedBlobs++ + evictedBlobBytes += storedSize + var hash [sha256.Size]byte + if _, err := hex.Decode(hash[:], []byte(sha256Hex)); err == nil { + base := srv.sha256Filepath(hash) + // Try removing both lz4 and plain paths; one or neither may exist. + if err := os.Remove(base + ".lz4"); err != nil && !os.IsNotExist(err) { + return ret, fmt.Errorf("removing disk file: %w", err) + } + if err := os.Remove(base); err != nil && !os.IsNotExist(err) { + return ret, fmt.Errorf("removing disk file: %w", err) + } + srv.removeFromHot(sha256Hex) + } + // Stop once we've freed enough aggregate bytes. For maxAge cleanup + // callers pass math.MaxInt64 so this check never triggers. + if evictedActionBytes >= maxBytes { + break + } + } + if err := tx.Commit(); err != nil { + return ret, fmt.Errorf("eviction Commit: %w", err) + } + + // Apply the queued delta decrements now that the tx is durably committed. + for _, p := range pending { + srv.addBlobDelta(p.sha256Hex, -1, -p.storedSize) + } + + srv.m.EvictedActions.Add(evictedActions) + srv.m.EvictedBlobs.Add(evictedBlobs) + srv.m.EvictedBytes.Add(evictedBlobBytes) + ret.Count = evictedBlobs + ret.Size = evictedBlobBytes + return ret, nil +} + +// cleanupTick decides whether the cache is over its configured maxAge or +// maxSize and, if so, runs one [Server.evictOldestActions] batch. The +// aggregate from the shard stats loop drives the decision; the eviction itself +// goes straight at idx_actions_access without consulting the shards. +func (srv *Server) cleanupTick(ctx context.Context) (countAndSize, error) { + srv.evictHotIfOver() + var ret countAndSize + us := srv.lastUsage.Load() + if us == nil { + return ret, nil + } + all := us.All() + // Include the dead-reckoning delta in the size pressure check so a + // burst of PUTs between scans actually triggers cleanup. The maxAge + // check intentionally ignores the delta: fresh PUTs are by definition + // within maxAge, so they neither add to nor subtract from the count of + // over-age Actions. + dCount, dBytes := srv.sumShardDeltas() + all.Count += dCount + all.Size += dBytes + + // Default cutoff to "any age" for size-only cleanup. evictOldestActions + // still uses idx_actions_access (the WHERE/ORDER BY drive the plan), so + // the very-large cutoff just means "no upper bound". + cutoff := int64(math.MaxInt64) + overAge := false + if srv.maxAge > 0 { + cutoff = srv.now().Add(-srv.maxAge).Unix() + overAge = us.All().Count-us.ActionsLE[srv.maxAge].Count > 0 + } + overSize := srv.maxSize > 0 && all.Size > srv.maxSize + if !overAge && !overSize { + return ret, nil + } + // maxAge cleanup runs to count limit (every stale Action must go); + // size-only cleanup runs until we've reclaimed enough bytes. + maxBytes := int64(math.MaxInt64) + if !overAge && overSize { + maxBytes = all.Size - srv.maxSize + } + return srv.evictOldestActions(ctx, cutoff, cleanupBatchSize, maxBytes) +} + +// checkpointTruncate runs PRAGMA wal_checkpoint(TRUNCATE) and returns SQLite's +// three result columns. A fully-applied checkpoint returns busy=0 and +// logFrames==ckptFrames; otherwise some frames remain in the WAL because of a +// concurrent reader pinning an older snapshot. +func (srv *Server) checkpointTruncate(ctx context.Context) (busy, logFrames, ckptFrames int, err error) { + err = srv.db.QueryRowContext(ctx, "PRAGMA wal_checkpoint(TRUNCATE)").Scan(&busy, &logFrames, &ckptFrames) + return busy, logFrames, ckptFrames, err +} + +// runCheckpointLoop periodically runs a TRUNCATE checkpoint to keep the WAL +// bounded on disk. SQLite's autocheckpoint only runs PASSIVE checkpoints, which +// reuse WAL space in place but never shrink the file; without this loop the WAL +// can grow without bound under continuous traffic. +func (srv *Server) runCheckpointLoop() { + t := time.NewTicker(checkpointInterval) + defer t.Stop() + for { + select { + case <-srv.shutdownCtx.Done(): + return + case <-t.C: + } + ctx, cancel := context.WithTimeout(srv.shutdownCtx, 2*time.Minute) + busy, logFrames, ckptFrames, err := srv.checkpointTruncate(ctx) + cancel() + if err != nil { + if errors.Is(err, context.Canceled) { + return + } + srv.logf("wal_checkpoint(TRUNCATE) error: %v", err) + continue + } + if busy != 0 || logFrames != ckptFrames { + // A reader is pinning frames; we'll catch up next tick. Logged + // because persistent partial checkpoints mean walJournalSizeLimit + // is the only thing keeping the file bounded, and we'd want to + // investigate. + srv.logf("wal_checkpoint(TRUNCATE) partial: busy=%d log=%d ckpt=%d", busy, logFrames, ckptFrames) + } else if srv.verbose { + srv.logf("wal_checkpoint(TRUNCATE): log=%d ckpt=%d", logFrames, ckptFrames) + } + // Refresh the size gauges immediately so dashboards see the + // post-truncate values without waiting for the next sampler tick. + srv.updateDBSizeMetrics() + } +} + +// dbPath returns the on-disk path of the SQLite main database file. +// The WAL file is at dbPath() + "-wal". +func (srv *Server) dbPath() string { + return filepath.Join(srv.sqliteDir, fmt.Sprintf("gocached-v%d.db", schemaVersion)) +} + +// updateDBSizeMetrics re-stats the SQLite files and updates the size gauges. +// A missing WAL file (e.g. on a fresh DB before the first write flushes) is +// reported as zero bytes. Other stat errors are logged but don't update the +// gauge, so a transient filesystem hiccup leaves the last-known value visible. +func (srv *Server) updateDBSizeMetrics() { + dbPath := srv.dbPath() + if fi, err := os.Stat(dbPath); err == nil { + srv.m.SQLiteDataBytes.Set(fi.Size()) + } else { + srv.logf("stat %s: %v", dbPath, err) + } + walPath := dbPath + "-wal" + switch fi, err := os.Stat(walPath); { + case err == nil: + srv.m.SQLiteWALBytes.Set(fi.Size()) + case errors.Is(err, os.ErrNotExist): + srv.m.SQLiteWALBytes.Set(0) + default: + srv.logf("stat %s: %v", walPath, err) + } +} + +// runDBSizeMetricsLoop samples the SQLite file sizes more frequently than the +// checkpoint loop runs, so the WAL gauge captures inter-checkpoint growth +// rather than only the post-truncate values. +func (srv *Server) runDBSizeMetricsLoop() { + t := time.NewTicker(dbSizeMetricsInterval) + defer t.Stop() + srv.updateDBSizeMetrics() // seed an initial sample at startup + for { + select { + case <-srv.shutdownCtx.Done(): + return + case <-t.C: + } + srv.updateDBSizeMetrics() + } +} + +func (srv *Server) runCleanLoop() { + for { + select { + case <-srv.shutdownCtx.Done(): + return + case <-time.After(cleanupTickInterval): + } + if _, err := srv.cleanupTick(srv.shutdownCtx); err != nil { + if errors.Is(err, context.Canceled) { + return + } + srv.logf("cleanup: %v", err) + } + } +} + +func (srv *Server) runCleanSessionsLoop() { + for { + select { + case <-srv.shutdownCtx.Done(): + return + case <-time.After(time.Hour): + } + + srv.mu.Lock() + count := len(srv.sessions) + var deleted int + for token, metadata := range srv.sessions { + if time.Now().After(metadata.expiry) { + delete(srv.sessions, token) + deleted++ + srv.m.Sessions.Add(-1) + } + } + srv.mu.Unlock() + if count > 0 { + srv.logf("cleaned up %d/%d access tokens", deleted, count) + } + } +} + +func durFmt(d time.Duration) string { + days := int(d.Hours() / 24) + if days > 0 { + return fmt.Sprintf("%dd", days) + } + return d.String() +} + +// requestAcceptsEncoding reports whether r's Accept-Encoding header includes +// the given encoding. It handles optional quality parameters (e.g. +// "gzip;q=1.0, lz4;q=0.5") by stripping them before comparison. +func requestAcceptsEncoding(r *http.Request, encoding string) bool { + for part := range strings.SplitSeq(r.Header.Get("Accept-Encoding"), ",") { + name, _, _ := strings.Cut(strings.TrimSpace(part), ";") + if strings.EqualFold(strings.TrimSpace(name), encoding) { + return true + } + } + return false +} + +func bytesFmt(n int64) string { + if n >= 1<<30 { + return fmt.Sprintf("%.1f GiB", float64(n)/(1<<30)) + } + if n >= 1<<20 { + return fmt.Sprintf("%.1f MiB", float64(n)/(1<<20)) + } + if n >= 1<<10 { + return fmt.Sprintf("%.1f KiB", float64(n)/(1<<10)) + } + return fmt.Sprintf("%d bytes", n) +} + +func (srv *Server) serveUsage(w http.ResponseWriter, r *http.Request) { + if r.Method == "POST" { + // Manual "refresh now" trigger: a synchronous scan of every shard. + // Slow on a large DB (numShards * per-shard query time), but + // operators may want an immediate aggregate after a config change. + if err := srv.scanAllShards(r.Context()); err != nil { + http.Error(w, "scanAllShards: "+err.Error(), http.StatusInternalServerError) + return + } + } + + us := srv.lastUsage.Load() + if us == nil { + http.Error(w, "no usage stats available", http.StatusInternalServerError) + return + } + + // Dead-reckon the totals so PUT/eviction bursts since the last shard + // scan show up on this page immediately. The persisted vs pending split + // is implementation noise for /usage readers; if anyone needs the + // breakdown they can read gocached_{blob,pending_blob}_{count,bytes} + // from /metrics. + dCount, dBytes := srv.sumShardDeltas() + live := us.All() + live.Count += dCount + live.Size += dBytes + + w.Header().Set("Content-Type", "text/html; charset=utf-8") + fmt.Fprintf(w, "

gocached usage stats

\n") + fmt.Fprintf(w, "

Current usage: %v of limit %v

\n", + live, bytesFmt(srv.maxSize)) + + // Cohort histogram of stored Actions, by access-time age. + fmt.Fprintf(w, "

Actions by access-time age

\n") + fmt.Fprintf(w, "\n") + fmt.Fprintf(w, "\n") + for _, d := range slices.Sorted(maps.Keys(us.ActionsLE)) { + var title string + if d == math.MaxInt64 { + title = "all" + } else { + title = "<= " + durFmt(d) + } + c := us.ActionsLE[d] + fmt.Fprintf(w, "\n", + title, c.Count, bytesFmt(c.Size)) + } + fmt.Fprintf(w, "
AgeCountSize
%s%d%s
\n") + + // Shard scan freshness — how many shards were scanned within each + // cohort. Watching the smaller-bucket counts drop is the easiest way to + // spot the stats loop falling behind. + now := srv.now() + freshness, oldest, totalShards := srv.shardFreshness(now) + fmt.Fprintf(w, "

Shard scan freshness

\n") + fmt.Fprintf(w, "

%d of %d shards have a cached scan; staleness target %v.

\n", + totalShards, srv.numShards(), shardStalenessTarget) + if !oldest.IsZero() { + fmt.Fprintf(w, "

Oldest scan: %v ago (%v).

\n", + durFmt(now.Sub(oldest).Round(time.Second)), oldest.UTC().Format(time.RFC3339)) + } + fmt.Fprintf(w, "\n") + fmt.Fprintf(w, "\n") + for _, b := range shardFreshnessBuckets { + fmt.Fprintf(w, "\n", durFmt(b), freshness[b]) + } + fmt.Fprintf(w, "
Scanned withinShards
<= %v%d
\n") + + // Per-shard query duration percentiles, snapshotted from the cached + // QueryDuration of every persisted shard. The Prometheus histogram + // gocached_shard_scan_duration_seconds carries the same data over time. + p25, p50, p90, nDur := srv.shardScanDurationPercentiles() + fmt.Fprintf(w, "

Per-shard scan duration

\n") + if nDur == 0 { + fmt.Fprintf(w, "

No shard has been scanned yet.

\n") + } else { + fmt.Fprintf(w, "\n") + fmt.Fprintf(w, "\n") + fmt.Fprintf(w, "\n", nDur, p25.Round(time.Millisecond)) + fmt.Fprintf(w, "\n", p50.Round(time.Millisecond)) + fmt.Fprintf(w, "\n", p90.Round(time.Millisecond)) + fmt.Fprintf(w, "
PercentileDuration
p25 (n=%d)%v
p50%v
p90%v
\n") + } +} + +func (srv *Server) serveSessions(w http.ResponseWriter, r *http.Request) { + if r.Method != "GET" { + http.Error(w, "bad method", http.StatusMethodNotAllowed) + return + } + + srv.mu.RLock() + // Make a copy of all session data (excluding the mutex). + sessions := make([]*sessionData, 0, len(srv.sessions)) + for _, v := range srv.sessions { + v.mu.Lock() + sessions = append(sessions, &sessionData{ + expiry: v.expiry, + canWrite: v.canWrite, + writeNamespaceID: v.writeNamespaceID, + writeNamespace: v.writeNamespace, + extraReadNamespaceID: v.extraReadNamespaceID, + extraReadNamespace: v.extraReadNamespace, + claims: v.claims, + stats: v.stats, + }) + v.mu.Unlock() + } + srv.mu.RUnlock() + + w.Header().Set("Content-Type", "text/html; charset=utf-8") + fmt.Fprintf(w, "

gocached sessions

\n") + for _, iss := range srv.jwtIssuers { + fmt.Fprintf(w, "

JWT issuer: %s

\n", iss) + } + fmt.Fprintf(w, "

Number of sessions: %d

\n", len(sessions)) + + fmt.Fprintf(w, "\n") + fmt.Fprintf(w, "\n") + slices.SortFunc(sessions, func(a, b *sessionData) int { + return a.expiry.Compare(b.expiry) + }) + nsLabel := func(ns Namespace) string { + if ns == GlobalNamespace { + return "(global)" + } + return fmt.Sprintf("%q", ns) + } + for _, d := range slices.Backward(sessions) { + lastUsed := "never" + if !d.stats.LastUsed.IsZero() { + lastUsed = durFmt(time.Since(d.stats.LastUsed)) + " ago" + } + writeLabel := "(read-only)" + if d.canWrite { + writeLabel = nsLabel(d.writeNamespace) + } + readLabels := []string{nsLabel(GlobalNamespace)} + if d.extraReadNamespaceID != 0 { + readLabels = append(readLabels, nsLabel(d.extraReadNamespace)) + } + statsJSON, _ := json.MarshalIndent(d.stats, "", " ") + claimsJSON, _ := json.MarshalIndent(d.claims, "", " ") + fmt.Fprintf(w, "\n", + lastUsed, d.expiry.Format(time.RFC3339), writeLabel, strings.Join(readLabels, ", "), statsJSON, claimsJSON) + } + fmt.Fprintf(w, "
Last usedExpiry timeWriteReadsStatsClaims
%s%s%s%s
%s
%s
\n") +} + +// expvarCounterMetric is a Prometheus counter metric backed by an expvar.Int. +type expvarCounterMetric struct { + desc *prometheus.Desc + v *expvar.Int +} + +var _ prometheus.Metric = (*expvarCounterMetric)(nil) + +func (m *expvarCounterMetric) Desc() *prometheus.Desc { return m.desc } + +func (m *expvarCounterMetric) Write(out *dto.Metric) error { + val := float64(m.v.Value()) + out.Counter = &dto.Counter{Value: &val} + return nil +} + +// expvarGaugeMetric is a Prometheus gauge metric backed by an expvar.Int. +type expvarGaugeMetric struct { + desc *prometheus.Desc + v *expvar.Int +} + +var _ prometheus.Metric = (*expvarGaugeMetric)(nil) + +func (m *expvarGaugeMetric) Desc() *prometheus.Desc { return m.desc } + +func (m *expvarGaugeMetric) Write(out *dto.Metric) error { + val := float64(m.v.Value()) + out.Gauge = &dto.Gauge{Value: &val} + return nil +} + +// singleMetricCollector is a Prometheus collector that collects a single metric. +type singleMetricCollector struct { + metric prometheus.Metric +} + +var _ prometheus.Collector = singleMetricCollector{} + +func (c singleMetricCollector) Describe(ch chan<- *prometheus.Desc) { ch <- c.metric.Desc() } +func (c singleMetricCollector) Collect(ch chan<- prometheus.Metric) { ch <- c.metric } diff --git a/gocache/gocached/gocached_perf_test.go b/gocache/gocached/gocached_perf_test.go new file mode 100644 index 0000000..bbe2911 --- /dev/null +++ b/gocache/gocached/gocached_perf_test.go @@ -0,0 +1,275 @@ +// Copyright (c) Tailscale Inc & AUTHORS +// SPDX-License-Identifier: BSD-3-Clause + +package gocached + +// This file holds the perf regression test. It always runs in CI (so +// refactors that change query shape are caught by the assertions immediately), +// but defaults to a tiny row count so normal CI takes a fraction of a +// second. Operators investigating real-world scale set +// -performance-test-rows to e.g. 250_000_000 and -performance-test-dir to a +// path on a fast disk; the seed is reused across runs. + +import ( + "context" + "crypto/sha256" + "database/sql" + "encoding/hex" + "errors" + "flag" + "fmt" + "math" + "os" + "strconv" + "strings" + "testing" + "time" +) + +var ( + perfTestRows = flag.Int("performance-test-rows", 100, + "target number of (Blob, Action) pairs to populate before measuring; default 100 so plain `go test` finishes in a fraction of a second. Bump to e.g. 250_000_000 for real-scale measurements; only the deficit is inserted on each run, so re-runs with the same value are instant.") + perfTestDir = flag.String("performance-test-dir", "", + "directory holding the pre-seeded perf test DB. If empty (the default), a fresh per-run t.TempDir() is used. Set to a persistent path (with $GOCACHED_PERFTEST_DIR as an alternative) when -performance-test-rows is large enough that you don't want to re-seed every run.") +) + +// perfSeedBatchSize is the number of rows inserted per transaction during +// seeding. 5000 keeps each multi-row INSERT comfortably below SQLite's +// default 32766 host-parameter limit (5000 * 4 = 20000 < 32766) and gives a +// reasonable bytes-per-syscall ratio for fsync-free seed inserts. +const perfSeedBatchSize = 5000 + +func TestPerfQueries(t *testing.T) { + dir := *perfTestDir + if dir == "" { + if v := os.Getenv("GOCACHED_PERFTEST_DIR"); v != "" { + dir = v + } else { + // Default: per-test temp dir so a normal `go test` doesn't + // leave a seeded DB hanging around. Operators running real-scale + // measurements set -performance-test-dir explicitly to get + // seed reuse across runs. + dir = t.TempDir() + } + } + if err := os.MkdirAll(dir, 0o750); err != nil { + t.Fatal(err) + } + t.Logf("perf test dir: %s (rows=%d)", dir, *perfTestRows) + + srv, err := NewServer( + WithDir(dir), + WithLogf(t.Logf), + withoutBackgroundLoops(), + WithShardPrefixLen(2), + ) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { srv.Close() }) + + n := *perfTestRows + if err := perfSeed(t, srv, n); err != nil { + t.Fatalf("seed: %v", err) + } + + // Confirm the eviction query plan still uses idx_actions_access even on + // a 250M-row table — that's the whole point of INDEXED BY, but it's + // worth re-asserting at scale because the planner picks differently as + // table cardinality grows. + t.Run("EvictionQueryPlan_atScale", func(t *testing.T) { + row := srv.db.QueryRow("EXPLAIN QUERY PLAN "+evictionCandidateQuery, 0, cleanupBatchSize) + var id, parent, notused int + var detail string + if err := row.Scan(&id, &parent, ¬used, &detail); err != nil { + t.Fatalf("EXPLAIN QUERY PLAN: %v", err) + } + t.Logf("plan: %s", detail) + if !strings.Contains(detail, "idx_actions_access") { + t.Errorf("plan does not use idx_actions_access: %q", detail) + } + if strings.Contains(detail, "SCAN Actions") || strings.Contains(detail, "SCAN TABLE Actions") { + t.Errorf("plan does a table scan of Actions: %q", detail) + } + }) + + timed := func(name string, f func() error) time.Duration { + t.Helper() + start := time.Now() + if err := f(); err != nil { + t.Errorf("%s failed: %v", name, err) + return 0 + } + d := time.Since(start) + t.Logf("%-44s %v", name, d) + return d + } + + // usageStats is a plain in-memory aggregate read; it doesn't touch + // SQLite at all. The benchmark is here to lock in that property: any + // future regression that re-introduces a per-call DB scan would blow + // past the sub-millisecond threshold. + aggRead := timed("usageStats (in-memory aggregate read)", func() error { + _, err := srv.usageStats() + return err + }) + + perShard := timed(fmt.Sprintf("scanShard (1 shard, ~%dM rows)", n/256/1_000_000), func() error { + _, err := srv.scanShard(context.Background(), 0) + return err + }) + + cleanup := timed(fmt.Sprintf("evictOldestActions (%d-row batch)", cleanupBatchSize), func() error { + _, err := srv.evictOldestActions(context.Background(), math.MaxInt64, cleanupBatchSize, math.MaxInt64) + return err + }) + + // Legacy world-scoped GROUP BY: the query gocached used to run every + // 5 minutes. It's here as a baseline so reviewers can see how much + // faster the per-shard approach is and as an early-warning if SQLite + // ever gets dramatically faster at this shape. + legacy := timed(fmt.Sprintf("legacy GROUP BY cleanCandidates (%d-row LIMIT)", cleanupBatchSize), func() error { + rows, err := srv.db.Query(` + SELECT b.BlobID, MAX(a.AccessTime), b.StoredSize + FROM Blobs b LEFT JOIN Actions a ON b.BlobID = a.BlobID + GROUP BY b.BlobID + HAVING MAX(a.AccessTime) <= ? + ORDER BY MAX(a.AccessTime) + LIMIT ?`, int64(math.MaxInt64), cleanupBatchSize) + if err != nil { + return err + } + defer rows.Close() + var blobID, accessTime, storedSize int64 + for rows.Next() { + if err := rows.Scan(&blobID, &accessTime, &storedSize); err != nil { + return err + } + } + return rows.Err() + }) + + if cleanup > 0 { + t.Logf("legacy/incremental cleanup ratio: %.1fx slower", legacy.Seconds()/cleanup.Seconds()) + } + + // Wall-clock regression thresholds. These are deliberately loose for + // laptop/dev hardware; tighten on the CI box once we have a baseline + // run there. The intent is "scream loudly if any of these queries + // quietly grows by an order of magnitude," not "lock in a tight SLO." + if aggRead > 1*time.Millisecond { + t.Errorf("usageStats took %v; want sub-millisecond (it should not hit SQLite)", aggRead) + } + if perShard > 30*time.Second { + t.Errorf("scanShard took %v on ~%dM rows; want < 30s", perShard, n/256/1_000_000) + } + if cleanup > 2*time.Second { + t.Errorf("evictOldestActions took %v on %d-row batch; want < 2s (idx_actions_access range scan)", cleanup, cleanupBatchSize) + } +} + +// perfSeed brings the DB up to at least n (Blob, Action) pairs. If the DB +// already has that many it returns immediately so re-runs are fast; if it's +// short, only the deficit is inserted. Per-batch work is wrapped in one +// transaction so a Ctrl-C resumes cleanly at the last complete batch. +func perfSeed(t *testing.T, srv *Server, n int) error { + t.Helper() + + var seq sql.NullInt64 + err := srv.db.QueryRow(`SELECT seq FROM sqlite_sequence WHERE name = 'Blobs'`).Scan(&seq) + if err != nil && !errors.Is(err, sql.ErrNoRows) { + return fmt.Errorf("query sqlite_sequence: %w", err) + } + cur := int(seq.Int64) + if cur >= n { + t.Logf("DB already has %d Blobs (>= %d target); reusing seed", cur, n) + return nil + } + + t.Logf("seeding rows %d..%d", cur, n) + // PRAGMA synchronous=OFF makes the seed dramatically faster at the + // cost of crash safety. Acceptable for a synthetic test DB; the + // process explicitly runs with no other writers (background loops are + // disabled via withoutBackgroundLoops). The pragma resets on Close. + if _, err := srv.db.Exec(`PRAGMA synchronous = OFF`); err != nil { + return fmt.Errorf("PRAGMA synchronous: %w", err) + } + // 512 MiB page cache so the bulk inserts thrash less. + if _, err := srv.db.Exec(`PRAGMA cache_size = -524288`); err != nil { + return fmt.Errorf("PRAGMA cache_size: %w", err) + } + + start := time.Now() + lastLog := start + for i := cur; i < n; i += perfSeedBatchSize { + end := min(i+perfSeedBatchSize, n) + if err := perfSeedBatch(srv, i, end); err != nil { + return fmt.Errorf("seed batch [%d, %d): %w", i, end, err) + } + if time.Since(lastLog) > 30*time.Second { + elapsed := time.Since(start) + rate := float64(end-cur) / elapsed.Seconds() + remaining := time.Duration(float64(n-end) / rate * float64(time.Second)) + t.Logf("seed progress: %d / %d (%.1f%%, elapsed %v, ETA %v)", + end, n, 100*float64(end)/float64(n), + elapsed.Round(time.Second), remaining.Round(time.Second)) + lastLog = time.Now() + } + } + t.Logf("seeded %d rows in %v", n-cur, time.Since(start).Round(time.Second)) + return nil +} + +// perfSeedBatch inserts (Blob, Action) pairs for indices [start, end) in a +// single transaction. Synthetic SHA256s come from sha256(itoa(i)) so the +// resulting hashes are uniformly distributed across all shards. BlobIDs are +// the AUTOINCREMENT-assigned sequential values; we know them in advance +// because seeding always continues from sqlite_sequence's max. +func perfSeedBatch(srv *Server, start, end int) error { + tx, err := srv.db.Begin() + if err != nil { + return err + } + defer tx.Rollback() + + rowCount := end - start + + var sb strings.Builder + sb.Grow(64 + rowCount*8) + sb.WriteString("INSERT INTO Blobs (SHA256, StoredSize, UncompressedSize) VALUES ") + args := make([]any, 0, rowCount*3) + for i := start; i < end; i++ { + if i > start { + sb.WriteByte(',') + } + sb.WriteString("(?,?,?)") + h := sha256.Sum256([]byte(strconv.Itoa(i))) + args = append(args, hex.EncodeToString(h[:]), int64(100), int64(100)) + } + if _, err := tx.Exec(sb.String(), args...); err != nil { + return fmt.Errorf("insert Blobs: %w", err) + } + + sb.Reset() + sb.Grow(96 + rowCount*16) + sb.WriteString("INSERT INTO Actions (NamespaceID, ActionID, BlobID, AltOutputID, CreateTime, AccessTime) VALUES ") + args = args[:0] + for i := start; i < end; i++ { + if i > start { + sb.WriteByte(',') + } + sb.WriteString("(0,?,?,'',?,?)") + h := sha256.Sum256([]byte(strconv.Itoa(i))) + args = append(args, + hex.EncodeToString(h[:]), + int64(i+1), // BlobID; matches the AUTOINCREMENT sequence + int64(1234+i), // CreateTime + int64(1234+i), // AccessTime — monotonically increasing makes cleanup deterministic + ) + } + if _, err := tx.Exec(sb.String(), args...); err != nil { + return fmt.Errorf("insert Actions: %w", err) + } + + return tx.Commit() +} diff --git a/gocache/gocached/gocached_test.go b/gocache/gocached/gocached_test.go new file mode 100644 index 0000000..30ee848 --- /dev/null +++ b/gocache/gocached/gocached_test.go @@ -0,0 +1,2117 @@ +// Copyright (c) Tailscale Inc & AUTHORS +// SPDX-License-Identifier: BSD-3-Clause + +package gocached + +import ( + "bytes" + "context" + "crypto" + "crypto/ecdsa" + "crypto/elliptic" + "crypto/rand" + "crypto/rsa" + "encoding/json" + "expvar" + "fmt" + "io" + "math" + "net" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "slices" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/go-jose/go-jose/v4" + "github.com/golang-jwt/jwt/v5" + "github.com/google/go-cmp/cmp" + "github.com/pierrec/lz4/v4" + "github.com/tailscale/tb/gocache/cachers" +) + +// sha256OfEmpty is the SHA-256 hash of an empty string, used as a well-known +// value in SQLite to store bytes, as it's common. +const sha256OfEmpty = "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855" + +// Package-level ECDSA P-256 keys reused across tests that need OIDC signing +// keys. Generating these is the single most expensive step in the JWT tests, +// so we share them rather than regenerating per test. +var ( + testKey1 = mustGenerateTestKey() + testKey2 = mustGenerateTestKey() + testKey3 = mustGenerateTestKey() +) + +func mustGenerateTestKey() *ecdsa.PrivateKey { + k, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader) + if err != nil { + panic(fmt.Sprintf("generating test ECDSA key: %v", err)) + } + return k +} + +type tester struct { + t testing.TB + srv *Server + hs *httptest.Server + + timeMu sync.Mutex + curTime time.Time +} + +func (t *tester) Logf(format string, args ...any) { + t.t.Logf(format, args...) +} + +func (t *tester) now() time.Time { + t.timeMu.Lock() + defer t.timeMu.Unlock() + return t.curTime +} + +func (t *tester) advanceClock(d time.Duration) { + t.timeMu.Lock() + defer t.timeMu.Unlock() + t.curTime = t.curTime.Add(d) +} + +func (t *tester) mkClient() *cachers.HTTPClient { + clientCacheDir := t.t.TempDir() + return &cachers.HTTPClient{ + BaseURL: t.hs.URL, + Disk: &cachers.DiskCache{ + Dir: clientCacheDir, + Logf: func(format string, args ...any) { + t.Logf("client-disk: "+format, args...) + }, + }, + } +} + +func (st *tester) usageStats() *usageStats { + st.t.Helper() + // Stats are maintained by an async per-shard stats loop in production. Tests + // need a deterministic aggregate, so flush pending access-time bumps and + // rotate every shard before reading. + st.srv.flushAccessTimeBumps() + if err := st.srv.scanAllShards(context.Background()); err != nil { + st.t.Fatalf("scanAllShards: %v", err) + } + stats, err := st.srv.usageStats() + if err != nil { + st.t.Fatalf("usageStats: %v", err) + } + return stats +} + +// cleanOldObjects runs cleanup until it's a no-op, so tests get a fully +// drained eviction in one call. Updates the aggregate (via a forced shard +// rescan) between batches so the cleanupTick's "are we over budget" check +// sees the post-eviction state, not stale pre-eviction shard data. +func (st *tester) cleanOldObjects() countAndSize { + st.t.Helper() + var total countAndSize + ctx := context.Background() + for { + _ = st.usageStats() + res, err := st.srv.cleanupTick(ctx) + if err != nil { + st.t.Fatalf("cleanupTick: %v", err) + } + if res.Count == 0 { + return total + } + total.Count += res.Count + total.Size += res.Size + } +} + +func (st *tester) diskFiles() []string { + st.t.Helper() + var ret []string + + err := filepath.Walk(st.srv.dir, func(path string, fi os.FileInfo, err error) error { + if err != nil { + return err + } + if fi.IsDir() && strings.HasPrefix(fi.Name(), ".") && path != st.srv.dir { + return filepath.SkipDir + } + if !fi.Mode().IsRegular() || strings.HasPrefix(fi.Name(), ".") || strings.HasPrefix(fi.Name(), "gocached") { + return nil + } + ret = append(ret, fi.Name()) + return nil + }) + if err != nil { + st.t.Fatalf("Walk: %v", err) + } + slices.Sort(ret) + return ret +} + +// wantMetric is a helper to check an expvar.Int metric and reset it +// for future tests. +func (st *tester) wantMetric(m *expvar.Int, want int64) { + st.t.Helper() + if got := m.Value(); got != want { + st.t.Errorf("metric = %d, want %d", got, want) + } + m.Set(0) +} + +// close shuts down the tester's HTTP listener and server early so a later +// tester can reopen the same cache directory. The Cleanup-registered closes +// registered by newServerTester then become harmless no-ops. +func (st *tester) close() { + st.t.Helper() + st.hs.Close() + if err := st.srv.Close(); err != nil { + st.t.Fatalf("Close: %v", err) + } +} + +// drain synchronously settles all pending PUTs, standing in for the +// background put-queue pipeline that newServerTester disables. +func (st *tester) drain() { + st.t.Helper() + if err := st.srv.drainPendingPuts(); err != nil { + st.t.Fatalf("drainPendingPuts: %v", err) + } +} + +func (st *tester) wantPut(c *cachers.HTTPClient, actionID, outputID string, val string) { + ctx := context.Background() + st.t.Helper() + clientDiskPath, err := c.Put(ctx, actionID, outputID, int64(len(val)), strings.NewReader(val)) + if err != nil { + st.t.Fatalf("Put: %v", err) + } + if clientDiskPath == "" { + st.t.Fatal("Put returned empty disk path") + } + st.drain() + st.wantMetric(&st.srv.m.Puts, 1) + wrote, err := os.ReadFile(clientDiskPath) + if err != nil { + st.t.Fatalf("ReadFile: %v", err) + } + if string(wrote) != val { + st.t.Errorf("ReadFile got %q, want %q", wrote, val) + } +} + +func (st *tester) wantGet(c *cachers.HTTPClient, actionID, outputID, wantVal string) { + ctx := context.Background() + st.t.Helper() + gotOutputID, diskPath, err := c.Get(ctx, actionID) + if err != nil { + st.t.Fatalf("Get: %v", err) + } + if gotOutputID != outputID { + st.t.Errorf("Get got outputID %q, want %q", gotOutputID, outputID) + } + if diskPath == "" { + st.t.Fatal("Get returned empty disk path") + } + wrote, err := os.ReadFile(diskPath) + if err != nil { + st.t.Fatalf("ReadFile: %v", err) + } + if string(wrote) != wantVal { + st.t.Errorf("ReadFile got %q, want %q", wrote, wantVal) + } +} + +func (st *tester) wantGetMiss(c *cachers.HTTPClient, actionID string) { + ctx := context.Background() + st.t.Helper() + gotOutputID, diskPath, err := c.Get(ctx, actionID) + if err != nil { + st.t.Fatalf("Get: %v", err) + } + if gotOutputID != "" { + st.t.Errorf("Get got outputID %q; want empty", gotOutputID) + } + if diskPath != "" { + st.t.Fatalf("Get returned disk path %q; want empty", diskPath) + } +} + +// lz4Size returns the lz4-compressed size of data. +func lz4Size(t testing.TB, data []byte) int64 { + t.Helper() + var buf bytes.Buffer + w := lz4.NewWriter(&buf) + if err := w.Apply(lz4.SizeOption(uint64(len(data)))); err != nil { + t.Fatal(err) + } + if _, err := w.Write(data); err != nil { + t.Fatal(err) + } + if err := w.Close(); err != nil { + t.Fatal(err) + } + return int64(buf.Len()) +} + +func withClock(clk func() time.Time) ServerOption { + return func(cfg *Server) { + cfg.clock = clk + } +} + +// withoutBackgroundLoops disables every periodic goroutine that start would +// otherwise launch (stats loop, cleanup, checkpoint, DB-size metrics). +// Tests need this because the mocked clock makes scannedAt collide across +// concurrent scans, so a background scan can overwrite a freshly-PUT +// shard with a stale-snapshot empty result. Tests drive scanShard / +// cleanupTick / etc. synchronously instead. +func withoutBackgroundLoops() ServerOption { + return func(cfg *Server) { + cfg.disableBackgroundLoops = true + } +} + +func newServerTester(t testing.TB, extraOpts ...ServerOption) *tester { + st := &tester{ + t: t, + curTime: time.Unix(1234, 0), + } + + opts := []ServerOption{ + WithDir(t.TempDir()), + WithLogf(t.Logf), + WithVerbose(true), + withClock(st.now), + withoutBackgroundLoops(), + // Default to 16 shards instead of production's 256 to keep + // scanAllShards fast in tests (each shard is a separate SQL + // round-trip; on near-empty tables overhead dominates). Tests + // that need a specific shardPrefixLen pass WithShardPrefixLen + // in extraOpts to override. + WithShardPrefixLen(1), + } + srv, err := NewServer(append(opts, extraOpts...)...) + if err != nil { + t.Fatalf("starting gocached: %v", err) + } + st.srv = srv + + // The hot tier index is populated by an async startup scan; wait for it + // so tests observe deterministic hot tier behavior. + if srv.hot != nil { + waitFor(t, "hot index scan", srv.hot.ready.Load) + } + + st.hs = httptest.NewServer(st.srv) + t.Cleanup(func() { st.srv.Close() }) + t.Cleanup(st.hs.Close) + + return st +} + +// signingKey must be a key type the server's signing method accepts; either +// [*ecdsa.PrivateKey] for ES256 and ES384, or [*rsa.PrivateKey] for RS256. +type jwtFunc func(claims jwt.MapClaims, signingKey any) string + +type oidcConfig struct { + alg string + use string + method jwt.SigningMethod + extraKeys []jose.JSONWebKey +} + +type oidcOpt func(*oidcConfig) + +// withJWKAlg sets the published JWK's "alg" member. Empty omits it entirely, +// which is what AWS's STS issuers do on their RSA key. +func withJWKAlg(alg string) oidcOpt { + return func(c *oidcConfig) { c.alg = alg } +} + +// withJWKUse sets the published JWK's "use" member. Empty omits it entirely. +func withJWKUse(use string) oidcOpt { + return func(c *oidcConfig) { c.use = use } +} + +// withSigningMethod changes the algorithm the returned [jwtFunc] signs with, +// and the "alg" it puts in the token header. Callers must pass a key of a type +// the method can sign with. +func withSigningMethod(m jwt.SigningMethod) oidcOpt { + return func(c *oidcConfig) { c.method = m } +} + +// withExtraJWK publishes an additional key alongside the primary one, for +// reproducing issuers that serve a mix of usable and unusable keys. +func withExtraJWK(k jose.JSONWebKey) oidcOpt { + return func(c *oidcConfig) { c.extraKeys = append(c.extraKeys, k) } +} + +// startOIDCServer starts a mock OIDC server that gocached can use for JWT auth. +// The provided publicKey is what JWT signatures will be validated against. +func startOIDCServer(t *testing.T, publicKey crypto.PublicKey, opts ...oidcOpt) (iss string, jwtFunc jwtFunc) { + t.Helper() + cfg := oidcConfig{ + alg: "ES256", + use: "sig", + method: jwt.SigningMethodES256, + } + for _, o := range opts { + o(&cfg) + } + + mux := http.NewServeMux() + srv := httptest.NewServer(mux) + t.Cleanup(srv.Close) + + issuer := fmt.Sprintf("http://%s", srv.Listener.Addr().String()) + + mux.HandleFunc("/.well-known/openid-configuration", func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(map[string]any{ + "issuer": issuer, + "jwks_uri": fmt.Sprintf("%s/jwks", issuer), + }) + }) + mux.HandleFunc("/jwks", func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + keys := []jose.JSONWebKey{ + { + Key: publicKey, + KeyID: "test-key", + Algorithm: cfg.alg, + Use: cfg.use, + }, + } + json.NewEncoder(w).Encode(map[string]any{ + "keys": append(keys, cfg.extraKeys...), + }) + }) + + return issuer, func(claims jwt.MapClaims, signingKey any) string { + t.Helper() + unsignedTk := &jwt.Token{ + Header: map[string]any{ + "typ": "JWT", + "alg": cfg.method.Alg(), + "kid": "test-key", + }, + Claims: claims, + Method: cfg.method, + } + tk, err := unsignedTk.SignedString(signingKey) + if err != nil { + t.Fatalf("error signing token: %v", err) + } + + return tk + } +} + +func TestServer(t *testing.T) { + st := newServerTester(t) + + ctx := context.Background() + + // Make two clients (imagine: two different builder VMs) + c1 := st.mkClient() + c2 := st.mkClient() + + const testActionID = "0001" + const testActionIDMiss = "0002" // this one doesn't exist + const testActionIDBig = "0bbb" // non-inline object + const testActionIDEmpty = "0000" + const testOutputID = "9900" + const testOutputIDBig = "9bbb" + const testOutputIDEmpty = sha256OfEmpty + const testObjectValue = "test data" + testObjectValueBig := strings.Repeat("x", smallObjectSize+1) + + // Populate from the first client. + st.wantPut(c1, testActionID, testOutputID, testObjectValue) + st.wantPut(c1, testActionIDBig, testOutputIDBig, testObjectValueBig) + st.wantPut(c1, testActionIDEmpty, testOutputIDEmpty, "") + + // Read from the second client. + st.wantGet(c2, testActionID, testOutputID, testObjectValue) + st.wantGet(c2, testActionIDBig, testOutputIDBig, testObjectValueBig) + st.wantGet(c2, testActionIDEmpty, testOutputIDEmpty, "") + + // Check metrics + st.wantMetric(&st.srv.m.Gets, 3) + st.wantMetric(&st.srv.m.GetHits, 3) + st.wantMetric(&st.srv.m.GetHitsInline, 1) + + // Do the same get again from the same client. This shouldn't hit the network. + st.wantGet(c2, testActionID, testOutputID, testObjectValue) + st.wantMetric(&st.srv.m.Gets, 0) + + // Cache miss. This should hit the network and fail. + if _, _, err := c2.Get(ctx, testActionIDMiss); err != nil { + t.Fatalf("miss Get: %v", err) + } + st.wantMetric(&st.srv.m.Gets, 1) + st.wantMetric(&st.srv.m.GetHits, 0) + + // Check that access time gets updated. + // Do it from a fresh client without a disk cache. + st.wantMetric(&st.srv.m.GetAccessBumps, 0) + st.advanceClock(relAtimeSeconds * 2 * time.Second) // advance clock by 2 days + c3 := st.mkClient() + st.wantGet(c3, testActionID, testOutputID, testObjectValue) + st.wantMetric(&st.srv.m.GetAccessBumps, 1) + + // Get usage stats. The tester helper flushes pending access-time bumps + // and rotates every shard so the aggregate is deterministic. + stats := st.usageStats() + bigStored := int64(len(testObjectValueBig)) // below lz4CompressThreshold, stored uncompressed + totalSize := int64(9) + bigStored // 9 (inline "test data") + big + 0 (empty) + want := &usageStats{ + MissingBlobRows: 0, + ActionsLE: map[time.Duration]countAndSize{ + 24 * time.Hour: {Count: 1, Size: 9}, + 48 * time.Hour: {Count: 1, Size: 9}, + 96 * time.Hour: {Count: 3, Size: totalSize}, + 168 * time.Hour: {Count: 3, Size: totalSize}, + 336 * time.Hour: {Count: 3, Size: totalSize}, + 720 * time.Hour: {Count: 3, Size: totalSize}, + 2160 * time.Hour: {Count: 3, Size: totalSize}, + math.MaxInt64: {Count: 3, Size: totalSize}, + }, + } + if diff := cmp.Diff(stats, want); diff != "" { + t.Errorf("usageStats mismatch (-got +want):\n%s", diff) + } + + st.advanceClock(relAtimeSeconds * 2 * time.Second) // advance clock by 2 days +} + +// TestEvictionQueryPlan is the lockdown test for the eviction query: it +// runs EXPLAIN QUERY PLAN and asserts SQLite uses idx_actions_access and +// never does a SCAN of the Actions table. INDEXED BY in the query already +// forces this at runtime, but the test catches schema or query drift early +// (and gives a concrete failure message instead of a runtime SQL error). +func TestEvictionQueryPlan(t *testing.T) { + st := newServerTester(t) + row := st.srv.db.QueryRow("EXPLAIN QUERY PLAN "+evictionCandidateQuery, 0, 1) + var id, parent, notused int + var detail string + if err := row.Scan(&id, &parent, ¬used, &detail); err != nil { + t.Fatalf("EXPLAIN QUERY PLAN: %v", err) + } + t.Logf("plan: %s", detail) + if !strings.Contains(detail, "idx_actions_access") { + t.Errorf("plan does not use idx_actions_access: %q", detail) + } + if strings.Contains(detail, "SCAN Actions") || strings.Contains(detail, "SCAN TABLE Actions") { + t.Errorf("plan does a table scan of Actions: %q", detail) + } +} + +// TestGetFromNamespaceQueryPlan is the lockdown test for the GET lookup +// query: the NamespaceID IN (?, ?) + ActionID equality must resolve via the +// (NamespaceID, ActionID) primary key, never a table scan, regardless of the +// global generation in use. +func TestGetFromNamespaceQueryPlan(t *testing.T) { + st := newServerTester(t, WithGlobalGeneration(2)) + rows, err := st.srv.db.Query("EXPLAIN QUERY PLAN "+getFromNamespace, + st.srv.globalNamespaceID, int64(7), "0001", st.srv.globalNamespaceID) + if err != nil { + t.Fatalf("EXPLAIN QUERY PLAN: %v", err) + } + defer rows.Close() + var sawActionsSearch bool + for rows.Next() { + var id, parent, notused int + var detail string + if err := rows.Scan(&id, &parent, ¬used, &detail); err != nil { + t.Fatalf("scan: %v", err) + } + t.Logf("plan: %s", detail) + if strings.Contains(detail, "SCAN a") || strings.Contains(detail, "SCAN TABLE a") { + t.Errorf("plan does a table scan of Actions: %q", detail) + } + if strings.Contains(detail, "SEARCH a ") && strings.Contains(detail, "NamespaceID=? AND ActionID=?") { + sawActionsSearch = true + } + } + if err := rows.Err(); err != nil { + t.Fatal(err) + } + if !sawActionsSearch { + t.Error("plan does not search Actions by (NamespaceID, ActionID)") + } +} + +func TestEvictOldestActions(t *testing.T) { + st := newServerTester(t) + ctx := context.Background() + + // Four unique blobs (different content => different SHA256 => different + // BlobID), inserted 24h apart so AccessTime order matches insert order. + c := st.mkClient() + st.wantPut(c, "0001", "9901", "1") + st.advanceClock(24 * time.Hour) + st.wantPut(c, "0002", "9902", "22") + st.advanceClock(24 * time.Hour) + st.wantPut(c, "0003", "9903", "333") + st.advanceClock(24 * time.Hour) + st.wantPut(c, "0004", "9904", "4444") + + // All four should still be on disk. + if got, want := len(st.diskFiles()), 0; got != want { + // All four are <= smallObjectSize, so they're inline in the DB, + // not on disk. Confirms the test isn't accidentally on the disk path. + t.Errorf("diskFiles before evict: %d, want %d (all should be inline)", got, want) + } + + // Cutoff lets only the two oldest (1 and 22) through. limit=10 is generous. + res, err := st.srv.evictOldestActions(ctx, st.srv.now().Add(-25*time.Hour).Unix(), 10, math.MaxInt64) + if err != nil { + t.Fatalf("evictOldestActions: %v", err) + } + // 1:1 blobs => evicting Actions evicts Blobs too. + if want := (countAndSize{Count: 2, Size: 3}); res != want { + t.Errorf("evict result = %+v, want %+v", res, want) + } + + // Now no cutoff (math.MaxInt64) but limit=1: just the next-oldest ("333"). + res, err = st.srv.evictOldestActions(ctx, math.MaxInt64, 1, math.MaxInt64) + if err != nil { + t.Fatalf("evictOldestActions: %v", err) + } + if want := (countAndSize{Count: 1, Size: 3}); res != want { + t.Errorf("evict result = %+v, want %+v", res, want) + } +} + +// TestActionLRUSharedBlob locks in the Action-LRU semantic on shared blobs: +// when two Actions point at the same Blob, evicting the old Action leaves +// the Blob alive (because the fresh Action still references it). +func TestActionLRUSharedBlob(t *testing.T) { + st := newServerTester(t) + ctx := context.Background() + + c := st.mkClient() + const content = "shared-blob-content" + // Two distinct ActionIDs but identical content => identical SHA256 => + // the same BlobID. Action 0001 is inserted "two days ago", 0002 today. + st.wantPut(c, "0001", "9901", content) + st.advanceClock(48 * time.Hour) + st.wantPut(c, "0002", "9902", content) + + // Evict everything older than 25h. Action-LRU: the old 0001 Action + // must go, but the Blob stays because 0002 still references it. + res, err := st.srv.evictOldestActions(ctx, st.srv.now().Add(-25*time.Hour).Unix(), 10, math.MaxInt64) + if err != nil { + t.Fatalf("evictOldestActions: %v", err) + } + if got, want := res.Count, int64(0); got != want { + t.Errorf("evicted Blob count = %d, want %d (Blob still has a fresher Action)", got, want) + } + var actionCount int + if err := st.srv.db.QueryRow(`SELECT COUNT(*) FROM Actions`).Scan(&actionCount); err != nil { + t.Fatal(err) + } + if actionCount != 1 { + t.Errorf("Actions count after evict = %d, want 1 (0001 evicted, 0002 alive)", actionCount) + } + var blobCount int + if err := st.srv.db.QueryRow(`SELECT COUNT(*) FROM Blobs`).Scan(&blobCount); err != nil { + t.Fatal(err) + } + if blobCount != 1 { + t.Errorf("Blobs count after evict = %d, want 1 (Blob shared with 0002, stays)", blobCount) + } +} + +func TestShardDeltaPUT(t *testing.T) { + st := newServerTester(t) + + // Before any PUT: delta is zero. + if c, b := st.srv.sumShardDeltas(); c != 0 || b != 0 { + t.Errorf("delta before PUT = (%d, %d), want (0, 0)", c, b) + } + + // Two PUTs of unique content. Each is small enough to be stored inline + // so StoredSize == len(content). + c := st.mkClient() + st.wantPut(c, "0001", "9901", "hello") + st.wantPut(c, "0002", "9902", "world!!") + + if got, want := func() int64 { c, _ := st.srv.sumShardDeltas(); return c }(), int64(2); got != want { + t.Errorf("delta count after 2 PUTs = %d, want %d", got, want) + } + if got, want := func() int64 { _, b := st.srv.sumShardDeltas(); return b }(), int64(5+7); got != want { + t.Errorf("delta bytes after 2 PUTs = %d, want %d", got, want) + } + + // Duplicate PUT (same actionID/outputID/content) goes through INSERT OR + // IGNORE: affected=0, so the delta must NOT advance. + st.wantPut(c, "0001", "9901", "hello") + if got, _ := st.srv.sumShardDeltas(); got != 2 { + t.Errorf("delta count after duplicate PUT = %d, want 2 (dup must not double-count)", got) + } +} + +func TestShardDeltaZeroedAfterScan(t *testing.T) { + // The whole point of the preDelta dance in scanAndPersistShard is that + // once a shard is rescanned, anything that was reflected in the new + // persisted snapshot leaves the delta. Verify it: PUT, scanAllShards, + // delta == 0. + st := newServerTester(t) + c := st.mkClient() + st.wantPut(c, "0001", "9901", "hello") + st.wantPut(c, "0002", "9902", "world!!") + + if cnt, b := st.srv.sumShardDeltas(); cnt == 0 && b == 0 { + t.Fatalf("delta is already zero before scan; test isn't exercising the post-scan zeroing") + } + if err := st.srv.scanAllShards(context.Background()); err != nil { + t.Fatal(err) + } + if cnt, b := st.srv.sumShardDeltas(); cnt != 0 || b != 0 { + t.Errorf("delta after scanAllShards = (%d, %d), want (0, 0)", cnt, b) + } +} + +func TestShardDeltaDrivesSizeCleanup(t *testing.T) { + // Dead-reckoning's reason for existing: cleanup must react to PUTs that + // happen between scans. Scenario: scan first so the persisted + // aggregate is at zero, then PUT enough to go over the limit *without* + // a scan, then assert cleanupTick fires. + st := newServerTester(t) + st.srv.maxSize = 8 + ctx := context.Background() + + // Empty cache: aggregate scanned, delta zero, cleanupTick no-op. + if err := st.srv.scanAllShards(ctx); err != nil { + t.Fatal(err) + } + if res, err := st.srv.cleanupTick(ctx); err != nil || res.Count != 0 { + t.Fatalf("cleanupTick on empty: %+v, err=%v; want zero", res, err) + } + + // PUT 10 bytes total. NO scanAllShards — the persisted aggregate still + // says zero. cleanupTick must consult the delta and decide to evict. + c := st.mkClient() + st.wantPut(c, "0001", "9901", "1") + st.advanceClock(time.Second) + st.wantPut(c, "0002", "9902", "22") + st.advanceClock(time.Second) + st.wantPut(c, "0003", "9903", "333") + st.advanceClock(time.Second) + st.wantPut(c, "0004", "9904", "4444") + st.advanceClock(time.Second) + + if us := st.srv.lastUsage.Load(); us.All().Size != 0 { + t.Fatalf("persisted aggregate is non-zero (%v); the test premise (no scan yet) is broken", us.All()) + } + if _, b := st.srv.sumShardDeltas(); b != 10 { + t.Fatalf("delta bytes = %d, want 10 (10 = 1+2+3+4)", b) + } + res, err := st.srv.cleanupTick(ctx) + if err != nil { + t.Fatalf("cleanupTick: %v", err) + } + // maxSize=8, delta-only total=10, so we need at least 2 bytes freed. + // The oldest are "1" (1B) and "22" (2B), summing to 3 ≥ 2 — stop after + // evicting both. Same as TestCleanOldObjectsBySize. + if want := (countAndSize{Count: 2, Size: 3}); res != want { + t.Errorf("cleanupTick result = %+v, want %+v", res, want) + } +} + +func TestCleanOldObjectsByAge(t *testing.T) { + st := newServerTester(t) + st.srv.maxAge = 24 * time.Hour + + // Populate some data. + c1 := st.mkClient() + st.wantPut(c1, "0001", "9901", strings.Repeat("x", smallObjectSize+1)) + st.advanceClock(25 * time.Hour) + st.wantPut(c1, "0002", "9902", strings.Repeat("x", smallObjectSize+2)) + st.wantPut(c1, "0003", "9903", "small") + smallLen := int64(len("small")) + + stored1 := int64(smallObjectSize + 1) // below lz4CompressThreshold, stored uncompressed + stored2 := lz4Size(t, []byte(strings.Repeat("x", smallObjectSize+2))) // at lz4CompressThreshold, lz4-compressed + + st1 := st.usageStats() + if all, want := st1.All(), (countAndSize{Count: 3, Size: stored1 + stored2 + smallLen}); all != want { + t.Errorf("usageStats: %v; want %v", all, want) + } + // First file is uncompressed (no .lz4 suffix), second is lz4-compressed. + if got, want := st.diskFiles(), []string{"333092a3daf718ed8f38a94e302df139edd4e3b5da4239a497995683942cf28c.lz4", "c6d8e9905300876046729949cc95c2385221270d389176f7234fe7ac00c4e430"}; !slices.Equal(got, want) { + t.Errorf("diskFiles: %v; want %v", got, want) + } + + clean1 := st.cleanOldObjects() + if clean1.Count != 1 || clean1.Size != stored1 { + t.Errorf("cleanOldObjects got %v, want {Count: 1, Size: %d}", clean1, stored1) + } + clean2 := st.cleanOldObjects() + if clean2.Count != 0 || clean2.Size != 0 { + t.Errorf("cleanOldObjects got %v, want {Count: 0, Size: 0}", clean2) + } + + st2 := st.usageStats() + if all, want := st2.All(), (countAndSize{Count: 2, Size: stored2 + smallLen}); all != want { + t.Errorf("usageStats after clean: %v; want %v", all, want) + } + if got, want := st.diskFiles(), []string{"333092a3daf718ed8f38a94e302df139edd4e3b5da4239a497995683942cf28c.lz4"}; !slices.Equal(got, want) { + t.Errorf("diskFiles after clean: %v; want %v", got, want) + } +} + +func TestCleanOldObjectsBySize(t *testing.T) { + st := newServerTester(t) + + // Populate some data. + c1 := st.mkClient() + st.wantPut(c1, "0001", "9901", "1") + st.advanceClock(time.Second) + st.wantPut(c1, "0002", "9902", "22") + st.advanceClock(time.Second) + st.wantPut(c1, "0003", "9903", "333") + st.advanceClock(time.Second) + st.wantPut(c1, "0004", "9904", "4444") + st.advanceClock(time.Second) + + st1 := st.usageStats() + if all, want := st1.All(), (countAndSize{Count: 4, Size: 10}); all != want { + t.Errorf("usageStats: %v; want %v", all, want) + } + + clean1 := st.cleanOldObjects() + if clean1.Count != 0 || clean1.Size != 0 { + t.Errorf("cleanOldObjects got %v, want no clean", clean1) + } + + st.srv.maxSize = 8 // the only way get to 8 or under is by deleting "1" and "22" (3 bytes) + + if got, want := st.cleanOldObjects(), (countAndSize{Count: 2, Size: 3}); got != want { + t.Errorf("cleanOldObjects got %v, want %v", got, want) + } + if got, want := st.usageStats().All(), (countAndSize{Count: 2, Size: 7}); got != want { + t.Errorf("usageStats: %v; want %v", got, want) + } +} + +func TestLZ4Storage(t *testing.T) { + st := newServerTester(t) + c := st.mkClient() + + type testCase struct { + name string + size int + wantLZ4 bool // expect .lz4 file on disk + wantDisk bool // expect any disk file (false = inline in DB) + } + tests := []testCase{ + {"empty", 0, false, false}, + {"tiny_1b", 1, false, false}, + {"inline_max", smallObjectSize, false, false}, + {"disk_no_lz4", smallObjectSize + 1, false, true}, // on disk but below lz4CompressThreshold + {"disk_at_lz4_threshold", lz4CompressThreshold, true, true}, // smallest lz4-compressed disk blob + {"disk_2k", 2048, true, true}, + {"disk_64k", 64 << 10, true, true}, + } + + for i, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + actionID := fmt.Sprintf("%04x", i+0x10) + outputID := fmt.Sprintf("%04x", i+0x90) + + data := make([]byte, tt.size) + if len(data) > 0 { + data[0] = 'X' + data[len(data)-1] = 'X' + } + val := string(data) + + // PUT + st.wantPut(c, actionID, outputID, val) + + // GET via client (sends Accept-Encoding: lz4) — verify round-trip. + c2 := st.mkClient() // fresh client, no disk cache + st.wantGet(c2, actionID, outputID, val) + + // Raw HTTP GET with Accept-Encoding: lz4 + req, _ := http.NewRequest("GET", st.hs.URL+"/action/"+actionID, nil) + req.Header.Set("Want-Object", "1") + req.Header.Set("Accept-Encoding", "lz4") + res, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatalf("raw GET: %v", err) + } + body, _ := io.ReadAll(res.Body) + res.Body.Close() + + if tt.wantLZ4 { + if got := res.Header.Get("Content-Encoding"); got != "lz4" { + t.Errorf("Accept lz4: Content-Encoding = %q, want %q", got, "lz4") + } + if got := res.Header.Get("X-Uncompressed-Length"); got != fmt.Sprint(tt.size) { + t.Errorf("Accept lz4: X-Uncompressed-Length = %q, want %q", got, fmt.Sprint(tt.size)) + } + if got := res.Header.Get("Content-Length"); got == fmt.Sprint(tt.size) { + t.Errorf("Accept lz4: Content-Length = %q, should be compressed (smaller)", got) + } + } else { + if got := res.Header.Get("Content-Encoding"); got != "" { + t.Errorf("Accept lz4: Content-Encoding = %q, want empty", got) + } + if got := res.Header.Get("X-Uncompressed-Length"); got != "" { + t.Errorf("Accept lz4: X-Uncompressed-Length = %q, want empty", got) + } + // Uncompressed: body should be the raw data. + if string(body) != val { + t.Errorf("Accept lz4: body length = %d, want %d", len(body), len(val)) + } + } + + // Raw HTTP GET without Accept-Encoding: lz4 — server must decompress. + req2, _ := http.NewRequest("GET", st.hs.URL+"/action/"+actionID, nil) + req2.Header.Set("Want-Object", "1") + // Deliberately no Accept-Encoding. + res2, err := http.DefaultClient.Do(req2) + if err != nil { + t.Fatalf("raw GET (no lz4): %v", err) + } + body2, _ := io.ReadAll(res2.Body) + res2.Body.Close() + + if got := res2.Header.Get("Content-Encoding"); got != "" { + t.Errorf("No Accept lz4: Content-Encoding = %q, want empty", got) + } + if got := res2.Header.Get("Content-Length"); got != fmt.Sprint(tt.size) { + t.Errorf("No Accept lz4: Content-Length = %q, want %q", got, fmt.Sprint(tt.size)) + } + if string(body2) != val { + t.Errorf("No Accept lz4: body length = %d, want %d", len(body2), len(val)) + } + }) + } + + // Verify disk state: expect exactly one plain file (the disk_no_lz4 case) + // and the rest with .lz4 suffix. + var plain, compressed int + for _, f := range st.diskFiles() { + if strings.HasSuffix(f, ".lz4") { + compressed++ + } else { + plain++ + } + } + if plain != 1 { + t.Errorf("disk files: got %d plain (non-lz4), want 1", plain) + } + if compressed != 3 { + t.Errorf("disk files: got %d .lz4, want 3", compressed) + } +} + +func TestShardPrefix(t *testing.T) { + // Verify the %0*x zero-padding contract: the width comes from the * + // argument and pads with leading zeros to that width. A typo (%0x, + // %x, %2x) would silently break shard keys. + cases := []struct { + prefixLen int + idx int + want shardPrefix + }{ + {1, 0, "0"}, + {1, 15, "f"}, + {2, 0, "00"}, + {2, 1, "01"}, + {2, 15, "0f"}, + {2, 16, "10"}, + {2, 255, "ff"}, + {3, 0, "000"}, + {3, 16, "010"}, + {3, 4095, "fff"}, + {4, 0, "0000"}, + {4, 65535, "ffff"}, + } + for _, c := range cases { + srv := &Server{shardPrefixLen: c.prefixLen} + if got := srv.shardPrefix(c.idx); got != c.want { + t.Errorf("shardPrefix(prefixLen=%d, idx=%d) = %q, want %q", + c.prefixLen, c.idx, got, c.want) + } + } +} + +func TestShardRange(t *testing.T) { + // shardRange's last-shard sentinel ("g" chars) is the easiest place for a + // future refactor to silently drop the final shard from the stats loop. + cases := []struct { + prefixLen int + idx int + wantLo, wantHi string + }{ + // First shard: lo = all zeros prefix, hi = next prefix. + {1, 0, "0", "1"}, + {2, 0, "00", "01"}, + // Middle shard. + {2, 15, "0f", "10"}, + {2, 16, "10", "11"}, + // Last shard: hi is the "g"-sentinel that sorts after all hex. + {1, 15, "f", "g"}, + {2, 255, "ff", "gg"}, + {3, 4095, "fff", "ggg"}, + {4, 65535, "ffff", "gggg"}, + } + for _, c := range cases { + srv := &Server{shardPrefixLen: c.prefixLen} + gotLo, gotHi := srv.shardRange(c.idx) + if gotLo != c.wantLo || gotHi != c.wantHi { + t.Errorf("shardRange(prefixLen=%d, idx=%d) = (%q, %q), want (%q, %q)", + c.prefixLen, c.idx, gotLo, gotHi, c.wantLo, c.wantHi) + } + } +} + +func TestLoadShardStats(t *testing.T) { + // Exercise the three branches that decide whether a persisted row gets + // kept or dropped: prefix-length mismatch, cohort mismatch, and a clean + // match. A dropped row stays in the table on disk but doesn't contribute + // to the in-memory aggregate. + // Override the tester default (shardPrefixLen=1) because this test + // asserts the wrong-prefix-length and wrong-cohort branches using + // 2-character prefixes. + st := newServerTester(t, WithShardPrefixLen(2)) + ctx := context.Background() + + if got := st.srv.numShards(); got != 256 { + t.Fatalf("numShards = %d, want 256", got) + } + + // One clean shard scan so we have a row in BlobShardStats whose schema + // matches the current server. + if err := st.srv.scanAndPersistShard(ctx, 0); err != nil { + t.Fatalf("scanAndPersistShard: %v", err) + } + + // Hand-write two more rows that should be rejected on the next load. + mustExec := func(query string, args ...any) { + t.Helper() + if _, err := st.srv.db.ExecContext(ctx, query, args...); err != nil { + t.Fatalf("exec %q: %v", query, err) + } + } + mustExec(`INSERT INTO BlobShardStats (Prefix, ScannedAt, StatsJSON) VALUES (?, ?, ?)`, + "abc", 1, `{"ActionsLE":{}}`) // wrong prefix length + mustExec(`INSERT INTO BlobShardStats (Prefix, ScannedAt, StatsJSON) VALUES (?, ?, ?)`, + "01", 1, `{"ActionsLE":{"60000000000":{"Count":1,"Size":1}}}`) // wrong cohort set + + if err := st.srv.loadShardStats(ctx); err != nil { + t.Fatalf("loadShardStats: %v", err) + } + if got, want := len(st.srv.shardStats), 1; got != want { + t.Errorf("shardStats has %d entries, want %d (only the clean prefix=00 row should survive)", got, want) + } + if _, ok := st.srv.shardStats["00"]; !ok { + t.Errorf("shardStats missing prefix=00") + } +} + +func TestShardScanDurationPercentiles(t *testing.T) { + srv := &Server{ + shardStats: map[shardPrefix]*shardSnapshot{ + "00": {stats: &usageStats{QueryDuration: 10 * time.Millisecond}}, + "01": {stats: &usageStats{QueryDuration: 20 * time.Millisecond}}, + "02": {stats: &usageStats{QueryDuration: 30 * time.Millisecond}}, + "03": {stats: &usageStats{QueryDuration: 40 * time.Millisecond}}, + // Zero-duration shard is excluded so freshly-loaded-but-not-yet- + // observed entries don't skew the low percentile to zero. + "04": {stats: &usageStats{QueryDuration: 0}}, + }, + } + p25, p50, p90, n := srv.shardScanDurationPercentiles() + if n != 4 { + t.Errorf("n = %d, want 4 (zero-duration shard should be excluded)", n) + } + // With sorted = [10,20,30,40] ms and rank-index picks, p25 -> idx 1, + // p50 -> idx 2, p90 -> idx 3. + if p25 != 20*time.Millisecond { + t.Errorf("p25 = %v, want 20ms", p25) + } + if p50 != 30*time.Millisecond { + t.Errorf("p50 = %v, want 30ms", p50) + } + if p90 != 40*time.Millisecond { + t.Errorf("p90 = %v, want 40ms", p90) + } + + empty := &Server{shardStats: map[shardPrefix]*shardSnapshot{}} + if p25, p50, p90, n := empty.shardScanDurationPercentiles(); p25 != 0 || p50 != 0 || p90 != 0 || n != 0 { + t.Errorf("empty server: got (%v, %v, %v, n=%d), want all zero", p25, p50, p90, n) + } +} + +func TestShardFreshness(t *testing.T) { + now := time.Unix(1_000_000, 0) + at := func(secsAgo int64) time.Time { return now.Add(-time.Duration(secsAgo) * time.Second) } + + srv := &Server{ + shardStats: map[shardPrefix]*shardSnapshot{ + "00": {stats: &usageStats{}, scannedAt: at(10)}, // <= 1m + "01": {stats: &usageStats{}, scannedAt: at(200)}, // <= 5m + "02": {stats: &usageStats{}, scannedAt: at(800)}, // <= 15m + "03": {stats: &usageStats{}, scannedAt: at(3000)}, // <= 1h + "04": {stats: &usageStats{}, scannedAt: at(80_000)}, // <= 24h + "05": {stats: &usageStats{}, scannedAt: at(200_000)}, // older than 24h + }, + } + counts, oldest, total := srv.shardFreshness(now) + if total != 6 { + t.Errorf("total = %d, want 6", total) + } + if got, want := oldest, at(200_000); !got.Equal(want) { + t.Errorf("oldest = %v, want %v", got, want) + } + want := map[time.Duration]int{ + 1 * time.Minute: 1, // just the 10s-old shard + 5 * time.Minute: 2, // adds the 200s-old + 15 * time.Minute: 3, // adds the 800s-old + 1 * time.Hour: 4, // adds the 3000s-old + 6 * time.Hour: 4, // unchanged (next is 80000s = ~22h) + 24 * time.Hour: 5, // adds the 22h-old; the 200000s (~2.3d) one is still excluded + } + for d, c := range want { + if got := counts[d]; got != c { + t.Errorf("counts[%v] = %d, want %d", d, got, c) + } + } +} + +func TestShardHealthGauges(t *testing.T) { + st := newServerTester(t) + ctx := context.Background() + + // Before any scan: unscanned should equal numShards, oldest age 0. + body := scrapeMetrics(t, st) + if got, want := promGauge(t, body, "gocached_shard_stats_unscanned"), float64(st.srv.numShards()); got != want { + t.Errorf("unscanned before scan = %v, want %v", got, want) + } + if got := promGauge(t, body, "gocached_shard_stats_oldest_age_seconds"); got != 0 { + t.Errorf("oldest_age_seconds before scan = %v, want 0", got) + } + + // Scan one shard at t=0, then advance the clock and re-scrape. The + // oldest-age gauge should reflect the advance; unscanned should drop by 1. + if err := st.srv.scanAndPersistShard(ctx, 0); err != nil { + t.Fatalf("scanAndPersistShard: %v", err) + } + st.advanceClock(90 * time.Second) + + body = scrapeMetrics(t, st) + if got, want := promGauge(t, body, "gocached_shard_stats_unscanned"), float64(st.srv.numShards()-1); got != want { + t.Errorf("unscanned after 1 scan = %v, want %v", got, want) + } + if got, want := promGauge(t, body, "gocached_shard_stats_oldest_age_seconds"), float64(90); got != want { + t.Errorf("oldest_age_seconds after 90s advance = %v, want %v", got, want) + } +} + +// scrapeMetrics returns the /metrics body served by the test server's +// metrics handler. +func scrapeMetrics(t *testing.T, st *tester) string { + t.Helper() + req := httptest.NewRequest("GET", "/metrics", nil) + rec := httptest.NewRecorder() + st.srv.metricsHandler.ServeHTTP(rec, req) + if rec.Code != http.StatusOK { + t.Fatalf("/metrics status = %d, body=%s", rec.Code, rec.Body.String()) + } + return rec.Body.String() +} + +// promGauge returns the value of the named gauge in a Prometheus text +// exposition body, fatal-ing if the line isn't present. Matches lines like +// " " but not commented HELP/TYPE lines. +func promGauge(t *testing.T, body, name string) float64 { + t.Helper() + for line := range strings.Lines(body) { + line = strings.TrimRight(line, "\n") + if strings.HasPrefix(line, "#") { + continue + } + k, v, ok := strings.Cut(line, " ") + if !ok || k != name { + continue + } + var f float64 + if _, err := fmt.Sscanf(v, "%g", &f); err != nil { + t.Fatalf("parsing %q value %q: %v", name, v, err) + } + return f + } + t.Fatalf("metric %q not found in /metrics body:\n%s", name, body) + return 0 +} + +func TestServeUsage(t *testing.T) { + st := newServerTester(t) + c := st.mkClient() + st.wantPut(c, "0001", "9901", "hello") + // Force a scan pass so the freshness section has data to render. + if err := st.srv.scanAllShards(context.Background()); err != nil { + t.Fatalf("scanAllShards: %v", err) + } + + req := httptest.NewRequest("GET", "/usage", nil) + rec := httptest.NewRecorder() + st.srv.serveUsage(rec, req) + if rec.Code != http.StatusOK { + t.Fatalf("status = %d, want 200", rec.Code) + } + body := rec.Body.String() + // Each new section should render at least its heading. Asserting on the + // headings (rather than exact bytes) lets the formatting tweak without + // breaking the test. + for _, want := range []string{ + "Actions by access-time age", + "Shard scan freshness", + "Per-shard scan duration", + } { + if !strings.Contains(body, want) { + t.Errorf("/usage body missing %q\nbody:\n%s", want, body) + } + } +} + +func TestWALCheckpoint(t *testing.T) { + st := newServerTester(t) + ctx := context.Background() + + walPath := filepath.Join(st.srv.dir, fmt.Sprintf("gocached-v%d.db-wal", schemaVersion)) + + // Generate enough write traffic for the WAL to be a few pages large. + // The exact number isn't important; we just want it big enough that + // "shrank to nearly empty" is a meaningful observation. + for i := range 500 { + if _, err := st.srv.db.ExecContext(ctx, + `INSERT INTO Actions (NamespaceID, ActionID, BlobID, AltOutputID, CreateTime, AccessTime) VALUES (0, ?, 0, '', 0, 0)`, + fmt.Sprintf("%032x", i)); err != nil { + t.Fatalf("insert %d: %v", i, err) + } + } + + before, err := os.Stat(walPath) + if err != nil { + t.Fatalf("stat WAL: %v", err) + } + t.Logf("WAL before checkpoint: %d bytes", before.Size()) + if before.Size() < 4096 { + // 500 row inserts should leave at least one full WAL page even after + // the periodic autocheckpoint reuses space in place. A tiny WAL here + // means the test isn't actually exercising the truncation path. + t.Fatalf("WAL only %d bytes before checkpoint; expected meaningful traffic", before.Size()) + } + + st.srv.updateDBSizeMetrics() + if got := st.srv.m.SQLiteWALBytes.Value(); got != before.Size() { + t.Errorf("sqlite_wal_bytes gauge before checkpoint: got %d, want %d", got, before.Size()) + } + if got := st.srv.m.SQLiteDataBytes.Value(); got <= 0 { + t.Errorf("sqlite_data_bytes gauge: got %d, want > 0", got) + } + + busy, log, ckpt, err := st.srv.checkpointTruncate(ctx) + if err != nil { + t.Fatalf("checkpointTruncate: %v", err) + } + if busy != 0 || log != ckpt { + t.Errorf("checkpoint not fully applied: busy=%d log=%d ckpt=%d", busy, log, ckpt) + } + + after, err := os.Stat(walPath) + if err != nil { + t.Fatalf("stat WAL after checkpoint: %v", err) + } + t.Logf("WAL after checkpoint: %d bytes", after.Size()) + // TRUNCATE checkpoint with no concurrent readers truncates the WAL to 0. + if after.Size() != 0 { + t.Errorf("WAL not truncated: got %d bytes, want 0", after.Size()) + } + + st.srv.updateDBSizeMetrics() + if got := st.srv.m.SQLiteWALBytes.Value(); got != 0 { + t.Errorf("sqlite_wal_bytes gauge after checkpoint: got %d, want 0", got) + } +} + +func TestClientConnReuse(t *testing.T) { + st := newServerTester(t) + + var numDials atomic.Int32 + tr := http.DefaultTransport.(*http.Transport).Clone() + tr.DialContext = func(ctx context.Context, network, addr string) (net.Conn, error) { + num := numDials.Add(1) + t.Logf("DialContext #%d for %s %s", num, network, addr) + var std net.Dialer + return std.DialContext(ctx, network, addr) + } + t.Cleanup(func() { tr.CloseIdleConnections() }) + + c1 := st.mkClient() + c1.HTTPClient = &http.Client{Transport: tr} + const missAction = "0001" + st.wantGetMiss(c1, missAction) + st.wantGetMiss(c1, missAction) + st.wantGetMiss(c1, missAction) + st.wantPut(c1, "0001", "9901", "1") + st.wantGet(c1, "0001", "9901", "1") + if got := numDials.Load(); got != 1 { + t.Errorf("numDials = %d; want 1", got) + } +} + +func TestExchangeToken(t *testing.T) { + privateKey := testKey1 + otherPrivateKey := testKey2 + + for name, tc := range map[string]struct { + mutateClaims func(jwt.MapClaims) + signingKey *ecdsa.PrivateKey + wantStatusCode int + }{ + // Base case: no mutation. + "valid": { + wantStatusCode: http.StatusOK, + }, + // Every other test makes one mutation from the base case that should cause failure. + "missing_sub": { + mutateClaims: func(cl jwt.MapClaims) { + delete(cl, "sub") + }, + wantStatusCode: http.StatusUnauthorized, + }, + "invalid_sub": { + mutateClaims: func(cl jwt.MapClaims) { + cl["sub"] = "user456" + }, + wantStatusCode: http.StatusUnauthorized, + }, + "invalid_iss": { + mutateClaims: func(cl jwt.MapClaims) { + cl["iss"] = "invalid_issuer" + }, + wantStatusCode: http.StatusUnauthorized, + }, + "invalid_aud": { + mutateClaims: func(cl jwt.MapClaims) { + cl["aud"] = "invalid_audience" + }, + wantStatusCode: http.StatusUnauthorized, + }, + "not_yet_valid": { + mutateClaims: func(cl jwt.MapClaims) { + cl["nbf"] = jwt.NewNumericDate(time.Now().Add(10 * time.Minute)) + }, + wantStatusCode: http.StatusUnauthorized, + }, + "expired": { + mutateClaims: func(cl jwt.MapClaims) { + cl["exp"] = jwt.NewNumericDate(time.Now().Add(-time.Minute)) + }, + wantStatusCode: http.StatusUnauthorized, + }, + "invalid_signature": { + signingKey: otherPrivateKey, + wantStatusCode: http.StatusUnauthorized, + }, + } { + t.Run(name, func(t *testing.T) { + issuer, createJWT := startOIDCServer(t, privateKey.Public()) + st := newServerTester(t, + WithJWTAuth(issuer), + WithNamespaceMapping(func(ctx context.Context, claims map[string]any) (*NamespaceGrant, error) { + if claims["sub"] != "user123" { + return nil, fmt.Errorf("sub = %v, want user123", claims["sub"]) + } + return &NamespaceGrant{WriteNamespace: new(GlobalNamespace)}, nil + }), + ) + + // Generate JWT. + tokenClaims := jwt.MapClaims{ + "sub": "user123", + "num": 42, + "iss": issuer, + "aud": gocachedAudience, + "nbf": jwt.NewNumericDate(time.Now().Add(-time.Minute)), + "exp": jwt.NewNumericDate(time.Now().Add(time.Hour)), + } + if tc.mutateClaims != nil { + tc.mutateClaims(tokenClaims) + } + signingKey := privateKey + if tc.signingKey != nil { + signingKey = tc.signingKey + } + body, err := json.Marshal(map[string]any{ + "jwt": createJWT(tokenClaims, signingKey), + }) + if err != nil { + t.Fatalf("error marshaling request body: %v", err) + } + + // Exchange JWT for access token. + req, err := http.NewRequest("POST", st.hs.URL+"/auth/exchange-token", bytes.NewReader(body)) + if err != nil { + t.Fatalf("error creating request: %v", err) + } + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatalf("error making request: %v", err) + } + defer resp.Body.Close() + if resp.StatusCode != tc.wantStatusCode { + t.Fatalf("unexpected status code: want %d, got %d", tc.wantStatusCode, resp.StatusCode) + } + body, err = io.ReadAll(resp.Body) + if err != nil { + t.Fatalf("error reading response body: %v", err) + } + + if tc.wantStatusCode != http.StatusOK { + if string(body) != "unauthorized\n" { + t.Fatalf("unexpected error body: %s", string(body)) + } + + // No access token to do further checks with; test finished. + return + } + + // Check returned access token. + var d struct { + AccessToken string `json:"access_token"` + } + if err := json.Unmarshal(body, &d); err != nil { + t.Fatalf("error decoding response body: %v", err) + } + if d.AccessToken == "" { + t.Fatalf("expected access_token in response, got %s", string(body)) + } + + cl := st.mkClient() + if _, _, err := cl.Get(t.Context(), "abc123"); err == nil { + t.Fatalf("Get without access token succeeded unexpectedly") + } + + cl.AccessToken = d.AccessToken + st.wantGetMiss(cl, "abc123") + + st.wantPut(cl, "abc123", "def456", "data789") + st.wantGet(cl, "abc123", "def456", "data789") + + // Check session stats. + reqStats, err := http.NewRequest("GET", st.hs.URL+"/session/stats", nil) + if err != nil { + t.Fatalf("error creating stats request: %v", err) + } + reqStats.Header.Set("Authorization", "Bearer "+d.AccessToken) + respStats, err := http.DefaultClient.Do(reqStats) + if err != nil { + t.Fatalf("error making stats request: %v", err) + } + defer respStats.Body.Close() + if respStats.StatusCode != http.StatusOK { + t.Fatalf("unexpected stats status code: want %d, got %d", http.StatusOK, respStats.StatusCode) + } + bodyStats, err := io.ReadAll(respStats.Body) + if err != nil { + t.Fatalf("error reading stats response body: %v", err) + } + var stats stats + if err := json.Unmarshal(bodyStats, &stats); err != nil { + t.Fatalf("error decoding stats response body: %v", err) + } + t.Logf("stats: %v", stats) + if stats.Gets == 0 { + t.Errorf("expected non-zero gets in session stats") + } + if stats.Puts == 0 { + t.Errorf("expected non-zero puts in session stats") + } + }) + } +} + +// Test an issuer that looks like AWS's STS issuers, which omit the optional +// "alg" member (RFC 7517 4.4) on their RSA signing key and also publish an +// ES384 EC key. +func TestExchangeTokenAWSShapedJWKS(t *testing.T) { + rsaKey, err := rsa.GenerateKey(rand.Reader, 2048) + if err != nil { + t.Fatalf("generating RSA key: %v", err) + } + ecKey, err := ecdsa.GenerateKey(elliptic.P384(), rand.Reader) + if err != nil { + t.Fatalf("generating P-384 key: %v", err) + } + + for name, tc := range map[string]struct { + publicKey crypto.PublicKey + signingKey any + opts []oidcOpt + }{ + // The key AWS actually signs identity tokens with. + "rsa_without_alg": { + publicKey: rsaKey.Public(), + signingKey: rsaKey, + opts: []oidcOpt{ + withJWKAlg(""), + withSigningMethod(jwt.SigningMethodRS256), + withExtraJWK(jose.JSONWebKey{ + Key: ecKey.Public(), + KeyID: "EC384_0", + Algorithm: "ES384", + Use: "sig", + }), + }, + }, + "es384": { + publicKey: ecKey.Public(), + signingKey: ecKey, + opts: []oidcOpt{ + withJWKAlg("ES384"), + withSigningMethod(jwt.SigningMethodES384), + }, + }, + // "use" is optional too, so its absence mustn't disqualify a key. + "without_use": { + publicKey: testKey1.Public(), + signingKey: testKey1, + opts: []oidcOpt{ + withJWKUse(""), + }, + }, + } { + t.Run(name, func(t *testing.T) { + issuer, createJWT := startOIDCServer(t, tc.publicKey, tc.opts...) + st := newServerTester(t, + WithJWTAuth(issuer), + WithNamespaceMapping(func(ctx context.Context, claims map[string]any) (*NamespaceGrant, error) { + return &NamespaceGrant{WriteNamespace: new(GlobalNamespace)}, nil + }), + ) + + resp, accessToken := exchangeToken(t, st.hs.URL, createJWT(baseClaims(issuer, "user123"), tc.signingKey)) + if resp.StatusCode != http.StatusOK { + t.Fatalf("got status %d, want %d", resp.StatusCode, http.StatusOK) + } + if accessToken == "" { + t.Error("expected an access token") + } + }) + } +} + +func TestJWKSNoUsableSigningKeys(t *testing.T) { + for name, opts := range map[string][]oidcOpt{ + "unsupported_alg": {withJWKAlg("PS512")}, + "encryption_key": {withJWKUse("enc")}, + } { + t.Run(name, func(t *testing.T) { + issuer, _ := startOIDCServer(t, testKey1.Public(), opts...) + _, err := NewServer( + WithDir(t.TempDir()), + WithLogf(t.Logf), + WithJWTAuth(issuer), + WithNamespaceMapping(func(ctx context.Context, claims map[string]any) (*NamespaceGrant, error) { + return &NamespaceGrant{WriteNamespace: new(GlobalNamespace)}, nil + }), + ) + if err == nil { + t.Fatal("NewServer succeeded with an issuer that has no usable signing keys, want error") + } + if !strings.Contains(err.Error(), "no usable signing keys") { + t.Errorf("error = %v, want it to mention no usable signing keys", err) + } + }) + } +} + +func TestExchangeTokenNamespaceValidation(t *testing.T) { + for name, tc := range map[string]struct { + namespace string + wantStatusCode int + }{ + "simple": {"user123", http.StatusOK}, + "structured_sub": {"repo:octo-org/octo-repo:environment:prod", http.StatusOK}, + "auth0_sub": {"auth0|507f1f77bcf86cd799439020", http.StatusOK}, + "email": {"alice+ci@example.com", http.StatusOK}, + "dependabot": {"dependabot[bot]", http.StatusOK}, + "space": {"has space", http.StatusUnauthorized}, + "non_ascii": {"héllo", http.StatusUnauthorized}, + "control_char": {"bad\tns", http.StatusUnauthorized}, + "html_meta": {"