Commit a0ec3bd

Antonio Mika  ·  2026-10-09 09:22:34 -0400 EDT
parent 5947599
perf(pgs): speed up uploads and fix quota accounting
6 files changed,  +182, -153
+42, -5
......@@ -33,10 +33,11 @@ const rsyncFileMax = 24 << 20
3333 // rsyncEnv is a pgs SSH server backed by one storage implementation, and a
3434 // real rsync client set up to talk to it.
3535 type rsyncEnv struct {
36- t *testing.T
37- st storage.StorageServe
38- bucket storage.Bucket
39- rsh string
36+ t *testing.T
37+ st storage.StorageServe
38+ bucket storage.Bucket
39+ rsh string
40+ feature *db.FeatureFlag
4041 }
4142
4243 func newRsyncEnv(t *testing.T, st storage.StorageServe) *rsyncEnv {
......@@ -94,7 +95,7 @@ func newRsyncEnv(t *testing.T, st storage.StorageServe) *rsyncEnv {
9495 "ssh -p %s -i %s -o IdentitiesOnly=yes -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null -o LogLevel=ERROR",
9596 port, keyFile,
9697 )
97- return &rsyncEnv{t: t, st: st, bucket: bucket, rsh: rsh}
98+ return &rsyncEnv{t: t, st: st, bucket: bucket, rsh: rsh, feature: dbpool.Feature}
9899 }
99100
100101 // discardPubsub drops cache purges, which pgs sends from goroutines that can
......@@ -395,3 +396,39 @@ func TestRsyncStorage(t *testing.T) {
395396 })
396397 }
397398 }
399+
400+// TestRsyncStorageQuota checks that shrinking and deleting files frees quota
401+// for files later in the same run.
402+func TestRsyncStorageQuota(t *testing.T) {
403+ for name, newStorage := range rsyncStorageBackends() {
404+ t.Run(name, func(t *testing.T) {
405+ env := newRsyncEnv(t, newStorage(t))
406+ env.feature.Data.StorageMax = 3 << 20
407+ src := tree{t, t.TempDir()}
408+ remote := "localhost:/site"
409+ t0 := time.Unix(1_600_000_000, 0)
410+
411+ src.write("a.bin", randomBytes(2<<20), t0)
412+ env.rsync("-rt", src.dir+"/", remote)
413+
414+ // a.bin is sent before b.bin, which only fits once a.bin's
415+ // shrink is counted.
416+ src.write("a.bin", randomBytes(512<<10), t0)
417+ src.write("b.bin", randomBytes(2<<20), t0)
418+ if s := env.rsync("-rt", src.dir+"/", remote); s.transferred != 2 {
419+ t.Fatalf("transferred %d files, want 2\n%s", s.transferred, s.output)
420+ }
421+
422+ if err := os.Remove(filepath.Join(src.dir, "a.bin")); err != nil {
423+ t.Fatal(err)
424+ }
425+ src.write("c.bin", randomBytes(768<<10), t0)
426+ if s := env.rsync("-rt", "--delete-before", src.dir+"/", remote); s.deleted != 1 || s.transferred != 1 {
427+ t.Fatalf("deleted %d and transferred %d files, want 1 and 1\n%s", s.deleted, s.transferred, s.output)
428+ }
429+ if _, _, ok := env.stored("/site/c.bin"); !ok {
430+ t.Fatal("c.bin was not stored")
431+ }
432+ })
433+ }
434+}
+82, -125
......@@ -23,82 +23,41 @@ import (
2323 ignore "github.com/sabhiram/go-gitignore"
2424 )
2525
26-type ctxBucketKey struct{}
27-type ctxStorageSizeKey struct{}
28-type ctxProjectKey struct{}
29-type ctxFeatureFlagKey struct{}
30-type ctxDenylistKey struct{}
31-
32-type DenyList struct {
33- Denylist string
26+type ctxUploadStateKey struct{}
27+
28+// uploadState is what a session's uploads share. It is set on the session
29+// once because every SetValue nests the session context one level deeper,
30+// and SFTP can run several writes at a time.
31+type uploadState struct {
32+ mu sync.Mutex
33+ bucket storage.Bucket
34+ featureFlag *db.FeatureFlag
35+ storageSize int64
36+ project *db.Project
37+ denylist *ignore.GitIgnore
3438 }
3539
36-func getDenylist(s *pssh.SSHServerConnSession) *DenyList {
37- v := s.Context().Value(ctxDenylistKey{})
38- if v == nil {
39- return nil
40- }
41- denylist := s.Context().Value(ctxDenylistKey{}).(*DenyList)
42- return denylist
43-}
44-
45-func setDenylist(s *pssh.SSHServerConnSession, denylist string) {
46- s.SetValue(ctxDenylistKey{}, &DenyList{Denylist: denylist})
47-}
48-
49-func getProject(s *pssh.SSHServerConnSession) *db.Project {
50- v := s.Context().Value(ctxProjectKey{})
51- if v == nil {
52- return nil
53- }
54- project := s.Context().Value(ctxProjectKey{}).(*db.Project)
55- return project
56-}
57-
58-func setProject(s *pssh.SSHServerConnSession, project *db.Project) {
59- s.SetValue(ctxProjectKey{}, project)
60-}
61-
62-func getFeatureFlag(s *pssh.SSHServerConnSession) *db.FeatureFlag {
63- v := s.Context().Value(ctxFeatureFlagKey{})
64- if v == nil {
65- return nil
66- }
67- ff := s.Context().Value(ctxFeatureFlagKey{}).(*db.FeatureFlag)
68- return ff
69-}
70-
71-func setFeatureFlag(s *pssh.SSHServerConnSession, ff *db.FeatureFlag) {
72- s.SetValue(ctxFeatureFlagKey{}, ff)
73-}
74-
75-func getBucket(s *pssh.SSHServerConnSession) (storage.Bucket, error) {
76- bucket := s.Context().Value(ctxBucketKey{}).(storage.Bucket)
77- if bucket.Name == "" {
78- return bucket, fmt.Errorf("bucket not set on `ssh.Context()` for connection")
40+func getUploadState(s *pssh.SSHServerConnSession) (*uploadState, error) {
41+ state, ok := s.Context().Value(ctxUploadStateKey{}).(*uploadState)
42+ if !ok {
43+ return nil, fmt.Errorf("upload state not set on `ssh.Context()` for connection")
7944 }
80- return bucket, nil
45+ return state, nil
8146 }
8247
83-func getStorageSize(s *pssh.SSHServerConnSession) uint64 {
84- return s.Context().Value(ctxStorageSizeKey{}).(uint64)
48+func (u *uploadState) getStorageSize() int64 {
49+ u.mu.Lock()
50+ defer u.mu.Unlock()
51+ return u.storageSize
8552 }
8653
87-func incrementStorageSize(s *pssh.SSHServerConnSession, fileSize int64) uint64 {
88- curSize := getStorageSize(s)
89- var nextStorageSize uint64
90- if fileSize < 0 {
91- nextStorageSize = curSize - uint64(fileSize)
92- } else {
93- nextStorageSize = curSize + uint64(fileSize)
94- }
95- s.SetValue(ctxStorageSizeKey{}, nextStorageSize)
96- return nextStorageSize
97-}
98-
99-func shouldIgnoreFile(fp, ignoreStr string) bool {
100- object := ignore.CompileIgnoreLines(strings.Split(ignoreStr, "\n")...)
101- return object.MatchesPath(fp)
54+// addStorageSize adds delta bytes to the bucket's size and returns the new
55+// size.
56+func (u *uploadState) addStorageSize(delta int64) int64 {
57+ u.mu.Lock()
58+ defer u.mu.Unlock()
59+ u.storageSize = max(u.storageSize+delta, 0)
60+ return u.storageSize
10261 }
10362
10463 type FileData struct {
......@@ -106,7 +65,7 @@ type FileData struct {
10665 User *db.User
10766 Bucket storage.Bucket
10867 Project *db.Project
109- DenyList string
68+ DenyList *ignore.GitIgnore
11069 }
11170
11271 type UploadAssetHandler struct {
......@@ -222,20 +181,22 @@ func (h *UploadAssetHandler) Validate(s *pssh.SSHServerConnSession) error {
222181 if err != nil {
223182 return err
224183 }
225- setFeatureFlag(s, ff)
226184
227185 assetBucket := shared.GetAssetBucketName(user.ID)
228186 bucket, err := h.Cfg.Storage.UpsertBucket(assetBucket)
229187 if err != nil {
230188 return err
231189 }
232- s.SetValue(ctxBucketKey{}, bucket)
233190
234191 totalStorageSize, err := h.Cfg.Storage.GetBucketQuota(bucket)
235192 if err != nil {
236193 return err
237194 }
238- s.SetValue(ctxStorageSizeKey{}, totalStorageSize)
195+ s.SetValue(ctxUploadStateKey{}, &uploadState{
196+ bucket: bucket,
197+ featureFlag: ff,
198+ storageSize: int64(totalStorageSize),
199+ })
239200
240201 logger.Info(
241202 "bucket size",
......@@ -252,6 +213,33 @@ func (h *UploadAssetHandler) Validate(s *pssh.SSHServerConnSession) error {
252213 return nil
253214 }
254215
216+// sessionProject returns the project being uploaded to and its compiled
217+// _pgs_ignore, finding or creating the project when it changes. A session
218+// can upload to many projects because SFTP connections are kept alive.
219+func (h *UploadAssetHandler) sessionProject(state *uploadState, user *db.User, projectName string, logger *slog.Logger) (*db.Project, *ignore.GitIgnore, error) {
220+ state.mu.Lock()
221+ defer state.mu.Unlock()
222+ if state.project != nil && state.project.Name == projectName {
223+ return state.project, state.denylist, nil
224+ }
225+
226+ project, err := h.Cfg.DB.UpsertProject(user.ID, projectName, projectName)
227+ if err != nil {
228+ logger.Error("upsert project", "err", err.Error())
229+ return nil, nil, err
230+ }
231+
232+ dlist, err := h.findDenylist(state.bucket, project, logger)
233+ if err != nil {
234+ logger.Info("failed to get denylist, setting default (.*)", "err", err.Error())
235+ dlist = ".*"
236+ }
237+
238+ state.project = project
239+ state.denylist = ignore.CompileIgnoreLines(strings.Split(dlist, "\n")...)
240+ return state.project, state.denylist, nil
241+}
242+
255243 func (h *UploadAssetHandler) findDenylist(bucket storage.Bucket, project *db.Project, logger *slog.Logger) (string, error) {
256244 fp, _, err := h.Cfg.Storage.GetObject(bucket, filepath.Join(project.ProjectDir, "_pgs_ignore"))
257245 if err != nil {
......@@ -299,27 +287,19 @@ func (h *UploadAssetHandler) Write(s *pssh.SSHServerConnSession, entry *sendutil
299287 "size", entry.Size,
300288 )
301289
302- bucket, err := getBucket(s)
290+ state, err := getUploadState(s)
303291 if err != nil {
304- logger.Error("could not find bucket in ctx", "err", err.Error())
292+ logger.Error("could not find upload state in ctx", "err", err.Error())
305293 return "", err
306294 }
295+ bucket := state.bucket
307296
308- project := getProject(s)
309297 projectName := shared.GetProjectName(entry)
310298 logger = logger.With("project", projectName)
311299
312- // find, create, or update project if we haven't already done it
313- // we need to also check if the project stored in ctx is the same project
314- // being uploaded since users can keep an ssh connection alive via sftp
315- // and created many projects in a single session
316- if project == nil || project.Name != projectName {
317- project, err = h.Cfg.DB.UpsertProject(user.ID, projectName, projectName)
318- if err != nil {
319- logger.Error("upsert project", "err", err.Error())
320- return "", err
321- }
322- setProject(s, project)
300+ project, denylist, err := h.sessionProject(state, user, projectName, logger)
301+ if err != nil {
302+ return "", err
323303 }
324304
325305 if project.Blocked != "" {
......@@ -343,33 +323,16 @@ func (h *UploadAssetHandler) Write(s *pssh.SSHServerConnSession, entry *sendutil
343323 // calculate the filsize difference between the same file already
344324 // stored and the updated file being uploaded
345325 assetFilename := shared.GetAssetFileName(entry)
346- obj, info, _ := h.Cfg.Storage.GetObject(bucket, assetFilename)
347326 var curFileSize int64
348- if info != nil {
327+ if info, err := h.Cfg.Storage.StatObject(bucket, assetFilename); err == nil {
349328 curFileSize = info.Size
350329 }
351- if obj != nil {
352- defer func() {
353- _ = obj.Close()
354- }()
355- }
356-
357- denylist := getDenylist(s)
358- if denylist == nil {
359- dlist, err := h.findDenylist(bucket, project, logger)
360- if err != nil {
361- logger.Info("failed to get denylist, setting default (.*)", "err", err.Error())
362- dlist = ".*"
363- }
364- setDenylist(s, dlist)
365- denylist = &DenyList{Denylist: dlist}
366- }
367330
368331 data := &FileData{
369332 FileEntry: entry,
370333 User: user,
371334 Bucket: bucket,
372- DenyList: denylist.Denylist,
335+ DenyList: denylist,
373336 Project: project,
374337 }
375338
......@@ -378,18 +341,15 @@ func (h *UploadAssetHandler) Write(s *pssh.SSHServerConnSession, entry *sendutil
378341 return "", err
379342 }
380343
381- featureFlag := getFeatureFlag(s)
382- if featureFlag == nil {
383- return "", fmt.Errorf("pico+ feature flag ctx not set")
384- }
344+ featureFlag := state.featureFlag
385345 // SFTP does not report file size so the more performant way to
386346 // check filesize constraints is to try and upload the file to s3
387347 // with a specialized reader that raises an error if the filesize limit
388348 // has been reached
389349 storageMax := featureFlag.Data.StorageMax
390350 fileMax := featureFlag.Data.FileMax
391- curStorageSize := getStorageSize(s)
392- remaining := int64(storageMax) - int64(curStorageSize)
351+ curStorageSize := state.getStorageSize()
352+ remaining := int64(storageMax) - curStorageSize
393353 sizeRemaining := min(remaining+curFileSize, fileMax)
394354 if sizeRemaining <= 0 {
395355 _, _ = fmt.Fprintln(s.Stderr(), "storage quota reached")
......@@ -428,8 +388,7 @@ func (h *UploadAssetHandler) Write(s *pssh.SSHServerConnSession, entry *sendutil
428388 return "", cerr
429389 }
430390
431- deltaFileSize := curFileSize - fsize
432- nextStorageSize := incrementStorageSize(s, deltaFileSize)
391+ nextStorageSize := state.addStorageSize(fsize - curFileSize)
433392
434393 url := h.Cfg.AssetURL(
435394 user.Name,
......@@ -477,11 +436,12 @@ func (h *UploadAssetHandler) Delete(s *pssh.SSHServerConnSession, entry *senduti
477436 "file", assetFilepath,
478437 )
479438
480- bucket, err := getBucket(s)
439+ state, err := getUploadState(s)
481440 if err != nil {
482- logger.Error("could not find bucket in ctx", "err", err.Error())
441+ logger.Error("could not find upload state in ctx", "err", err.Error())
483442 return err
484443 }
444+ bucket := state.bucket
485445
486446 projectName := shared.GetProjectName(entry)
487447 logger = logger.With("project", projectName)
......@@ -490,14 +450,7 @@ func (h *UploadAssetHandler) Delete(s *pssh.SSHServerConnSession, entry *senduti
490450
491451 // Check if this path represents a directory (has a . _pico_keep_dir marker)
492452 keepDirPath := filepath.Join(assetFilepath, "._pico_keep_dir")
493- keepDirReader, _, keepDirErr := h.Cfg.Storage.GetObject(bucket, keepDirPath)
494- if keepDirReader != nil {
495- defer func() {
496- _ = keepDirReader.Close()
497- }()
498- }
499-
500- if keepDirErr == nil {
453+ if _, err := h.Cfg.Storage.StatObject(bucket, keepDirPath); err == nil {
501454 // This is a directory being deleted. We must delete all nested
502455 // . _pico_keep_dir files first, otherwise os.Remove() on the
503456 // directory will fail with "directory not empty".
......@@ -563,7 +516,11 @@ func (h *UploadAssetHandler) Delete(s *pssh.SSHServerConnSession, entry *senduti
563516 return err
564517 }
565518 }
519+ info, statErr := h.Cfg.Storage.StatObject(bucket, assetFilepath)
566520 err = h.Cfg.Storage.DeleteObject(bucket, assetFilepath)
521+ if err == nil && statErr == nil {
522+ state.addStorageSize(-info.Size)
523+ }
567524
568525 surrogate := getSurrogateKey(user.Name, projectName)
569526 h.Cfg.CacheClearingQueue <- surrogate
......@@ -589,7 +546,7 @@ func (h *UploadAssetHandler) validateAsset(data *FileData) (bool, error) {
589546 }
590547
591548 fpath := strings.Replace(data.Filepath, "/"+projectName, "", 1)
592- if shouldIgnoreFile(fpath, data.DenyList) {
549+ if data.DenyList.MatchesPath(fpath) {
593550 err := fmt.Errorf(
594551 "ERROR: (%s) file rejected, https://pico.sh/pgs#-pgs-ignore",
595552 data.Filepath,
+20, -0
......@@ -3,8 +3,10 @@ package storage_test
33 import (
44 "bytes"
55 "crypto/rand"
6+ "errors"
67 "fmt"
78 "io"
9+ "io/fs"
810 "log/slog"
911 "slices"
1012 "strings"
......@@ -106,6 +108,21 @@ func TestStorageConformance(t *testing.T) {
106108 }
107109 })
108110
111+ t.Run("stat", func(t *testing.T) {
112+ info, err := st.StatObject(b, "/proj/css/site.css")
113+ if err != nil {
114+ t.Fatal(err)
115+ }
116+ if info.Size != 6 || !info.LastModified.Equal(t2) {
117+ t.Fatalf("size=%d mtime=%v", info.Size, info.LastModified)
118+ }
119+ for _, missing := range []string{"/proj/missing.txt", "/proj/css", "/proj/css/"} {
120+ if _, err := st.StatObject(b, missing); !errors.Is(err, fs.ErrNotExist) {
121+ t.Errorf("StatObject(%q) returned %v, want fs.ErrNotExist", missing, err)
122+ }
123+ }
124+ })
125+
109126 t.Run("list", func(t *testing.T) {
110127 now := list(t, st, b, "/proj/now.txt", false)[0]
111128 tests := []struct {
......@@ -214,6 +231,9 @@ func TestStorageConformance(t *testing.T) {
214231 if _, _, err := st.GetObject(b, "/proj/css/deep/x.css"); err == nil {
215232 t.Fatal("deleted object still readable")
216233 }
234+ if _, err := st.StatObject(b, "/proj/css/deep/x.css"); !errors.Is(err, fs.ErrNotExist) {
235+ t.Fatalf("StatObject after delete returned %v", err)
236+ }
217237 got := list(t, st, b, "/proj/css/", true)
218238 want := []string{entry("site.css", false, 15, t1)}
219239 if !slices.Equal(got, want) {
+20, -23
......@@ -1,8 +1,6 @@
11 package storage
22
33 import (
4- "crypto/md5"
5- "encoding/hex"
64 "fmt"
75 "io"
86 "io/fs"
......@@ -120,31 +118,30 @@ func (s *StorageFS) GetObject(bucket Bucket, fpath string) (utils.ReadAndReaderA
120118 _ = dat.Close()
121119 return nil, objInfo, err
122120 }
121+ fillObjectInfo(objInfo, info)
122+ return dat, objInfo, nil
123+}
124+
125+// fillObjectInfo sets an object's size, mtime and ETag from its file. The
126+// ETag is built from size and mtime, as web servers do, so serving a file
127+// doesn't mean reading it twice.
128+func fillObjectInfo(objInfo *ObjectInfo, info os.FileInfo) {
123129 objInfo.Size = info.Size()
124130 objInfo.LastModified = info.ModTime()
131+ objInfo.ETag = fmt.Sprintf("%x-%x", info.ModTime().UnixNano(), info.Size())
132+}
125133
126- etag := ""
127- // only generate etag if file is less than 10MB
128- if info.Size() <= int64(10*MB) {
129- // calculate etag
130- h := md5.New()
131- if _, err := io.Copy(h, dat); err != nil {
132- _ = dat.Close()
133- return nil, objInfo, err
134- }
135- md5Sum := h.Sum(nil)
136- etag = hex.EncodeToString(md5Sum)
137-
138- // reset os.File reader
139- _, err = dat.Seek(0, io.SeekStart)
140- if err != nil {
141- _ = dat.Close()
142- return nil, objInfo, err
143- }
134+func (s *StorageFS) StatObject(bucket Bucket, fpath string) (*ObjectInfo, error) {
135+ info, err := os.Stat(filepath.Join(bucket.Path, fpath))
136+ if err != nil {
137+ return nil, err
144138 }
145-
146- objInfo.ETag = etag
147- return dat, objInfo, nil
139+ if info.IsDir() {
140+ return nil, fmt.Errorf("%s is a directory: %w", fpath, fs.ErrNotExist)
141+ }
142+ objInfo := &ObjectInfo{ContentType: mime.GetMimeType(fpath)}
143+ fillObjectInfo(objInfo, info)
144+ return objInfo, nil
148145 }
149146
150147 func (s *StorageFS) PutObject(bucket Bucket, fpath string, contents io.Reader, info *ObjectInfo) (string, int64, error) {
+16, -0
......@@ -4,6 +4,7 @@ import (
44 "bytes"
55 "fmt"
66 "io"
7+ "io/fs"
78 "os"
89 "path/filepath"
910 "strings"
......@@ -119,6 +120,21 @@ func (s *StorageMemory) GetObject(bucket Bucket, fpath string) (utils.ReadAndRea
119120 return &seekableReader{bytes.NewReader([]byte(dat))}, objInfo, nil
120121 }
121122
123+func (s *StorageMemory) StatObject(bucket Bucket, fpath string) (*ObjectInfo, error) {
124+ s.mu.RLock()
125+ defer s.mu.RUnlock()
126+
127+ fpath = memoryKey(fpath)
128+ dat, ok := s.storage[bucket.Path][fpath]
129+ if !ok {
130+ return nil, fmt.Errorf("object does not exist: %s: %w", fpath, fs.ErrNotExist)
131+ }
132+ return &ObjectInfo{
133+ Size: int64(len(dat)),
134+ LastModified: s.mtimes[bucket.Path][fpath],
135+ }, nil
136+}
137+
122138 func (s *StorageMemory) PutObject(bucket Bucket, fpath string, contents io.Reader, info *ObjectInfo) (string, int64, error) {
123139 s.mu.Lock()
124140 defer s.mu.Unlock()
+2, -0
......@@ -31,6 +31,8 @@ type BucketStorage interface {
3131
3232 type ObjectStorage interface {
3333 GetObject(bucket Bucket, fpath string) (utils.ReadAndReaderAtCloser, *ObjectInfo, error)
34+ // StatObject returns an object's info without opening it.
35+ StatObject(bucket Bucket, fpath string) (*ObjectInfo, error)
3436 PutObject(bucket Bucket, fpath string, contents io.Reader, info *ObjectInfo) (string, int64, error)
3537 DeleteObject(bucket Bucket, fpath string) error
3638 ListObjects(bucket Bucket, dir string, recursive bool) ([]os.FileInfo, error)