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 | |
| 33 | 33 | // rsyncEnv is a pgs SSH server backed by one storage implementation, and a | |
| 34 | 34 | // real rsync client set up to talk to it. | |
| 35 | 35 | 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 | |
| 40 | 41 | } | |
| 41 | 42 | ||
| 42 | 43 | func newRsyncEnv(t *testing.T, st storage.StorageServe) *rsyncEnv { |
| ... | ... | @@ -94,7 +95,7 @@ func newRsyncEnv(t *testing.T, st storage.StorageServe) *rsyncEnv { | |
| 94 | 95 | "ssh -p %s -i %s -o IdentitiesOnly=yes -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null -o LogLevel=ERROR", | |
| 95 | 96 | port, keyFile, | |
| 96 | 97 | ) | |
| 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} | |
| 98 | 99 | } | |
| 99 | 100 | ||
| 100 | 101 | // discardPubsub drops cache purges, which pgs sends from goroutines that can |
| ... | ... | @@ -395,3 +396,39 @@ func TestRsyncStorage(t *testing.T) { | |
| 395 | 396 | }) | |
| 396 | 397 | } | |
| 397 | 398 | } | |
| 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 ( | |
| 23 | 23 | ignore "github.com/sabhiram/go-gitignore" | |
| 24 | 24 | ) | |
| 25 | 25 | ||
| 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 | |
| 34 | 38 | } | |
| 35 | 39 | ||
| 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") | |
| 79 | 44 | } | |
| 80 | - | return bucket, nil | |
| 45 | + | return state, nil | |
| 81 | 46 | } | |
| 82 | 47 | ||
| 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 | |
| 85 | 52 | } | |
| 86 | 53 | ||
| 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 | |
| 102 | 61 | } | |
| 103 | 62 | ||
| 104 | 63 | type FileData struct { |
| ... | ... | @@ -106,7 +65,7 @@ type FileData struct { | |
| 106 | 65 | User *db.User | |
| 107 | 66 | Bucket storage.Bucket | |
| 108 | 67 | Project *db.Project | |
| 109 | - | DenyList string | |
| 68 | + | DenyList *ignore.GitIgnore | |
| 110 | 69 | } | |
| 111 | 70 | ||
| 112 | 71 | type UploadAssetHandler struct { |
| ... | ... | @@ -222,20 +181,22 @@ func (h *UploadAssetHandler) Validate(s *pssh.SSHServerConnSession) error { | |
| 222 | 181 | if err != nil { | |
| 223 | 182 | return err | |
| 224 | 183 | } | |
| 225 | - | setFeatureFlag(s, ff) | |
| 226 | 184 | ||
| 227 | 185 | assetBucket := shared.GetAssetBucketName(user.ID) | |
| 228 | 186 | bucket, err := h.Cfg.Storage.UpsertBucket(assetBucket) | |
| 229 | 187 | if err != nil { | |
| 230 | 188 | return err | |
| 231 | 189 | } | |
| 232 | - | s.SetValue(ctxBucketKey{}, bucket) | |
| 233 | 190 | ||
| 234 | 191 | totalStorageSize, err := h.Cfg.Storage.GetBucketQuota(bucket) | |
| 235 | 192 | if err != nil { | |
| 236 | 193 | return err | |
| 237 | 194 | } | |
| 238 | - | s.SetValue(ctxStorageSizeKey{}, totalStorageSize) | |
| 195 | + | s.SetValue(ctxUploadStateKey{}, &uploadState{ | |
| 196 | + | bucket: bucket, | |
| 197 | + | featureFlag: ff, | |
| 198 | + | storageSize: int64(totalStorageSize), | |
| 199 | + | }) | |
| 239 | 200 | ||
| 240 | 201 | logger.Info( | |
| 241 | 202 | "bucket size", |
| ... | ... | @@ -252,6 +213,33 @@ func (h *UploadAssetHandler) Validate(s *pssh.SSHServerConnSession) error { | |
| 252 | 213 | return nil | |
| 253 | 214 | } | |
| 254 | 215 | ||
| 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 | + | ||
| 255 | 243 | func (h *UploadAssetHandler) findDenylist(bucket storage.Bucket, project *db.Project, logger *slog.Logger) (string, error) { | |
| 256 | 244 | fp, _, err := h.Cfg.Storage.GetObject(bucket, filepath.Join(project.ProjectDir, "_pgs_ignore")) | |
| 257 | 245 | if err != nil { |
| ... | ... | @@ -299,27 +287,19 @@ func (h *UploadAssetHandler) Write(s *pssh.SSHServerConnSession, entry *sendutil | |
| 299 | 287 | "size", entry.Size, | |
| 300 | 288 | ) | |
| 301 | 289 | ||
| 302 | - | bucket, err := getBucket(s) | |
| 290 | + | state, err := getUploadState(s) | |
| 303 | 291 | 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()) | |
| 305 | 293 | return "", err | |
| 306 | 294 | } | |
| 295 | + | bucket := state.bucket | |
| 307 | 296 | ||
| 308 | - | project := getProject(s) | |
| 309 | 297 | projectName := shared.GetProjectName(entry) | |
| 310 | 298 | logger = logger.With("project", projectName) | |
| 311 | 299 | ||
| 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 | |
| 323 | 303 | } | |
| 324 | 304 | ||
| 325 | 305 | if project.Blocked != "" { |
| ... | ... | @@ -343,33 +323,16 @@ func (h *UploadAssetHandler) Write(s *pssh.SSHServerConnSession, entry *sendutil | |
| 343 | 323 | // calculate the filsize difference between the same file already | |
| 344 | 324 | // stored and the updated file being uploaded | |
| 345 | 325 | assetFilename := shared.GetAssetFileName(entry) | |
| 346 | - | obj, info, _ := h.Cfg.Storage.GetObject(bucket, assetFilename) | |
| 347 | 326 | var curFileSize int64 | |
| 348 | - | if info != nil { | |
| 327 | + | if info, err := h.Cfg.Storage.StatObject(bucket, assetFilename); err == nil { | |
| 349 | 328 | curFileSize = info.Size | |
| 350 | 329 | } | |
| 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 | - | } | |
| 367 | 330 | ||
| 368 | 331 | data := &FileData{ | |
| 369 | 332 | FileEntry: entry, | |
| 370 | 333 | User: user, | |
| 371 | 334 | Bucket: bucket, | |
| 372 | - | DenyList: denylist.Denylist, | |
| 335 | + | DenyList: denylist, | |
| 373 | 336 | Project: project, | |
| 374 | 337 | } | |
| 375 | 338 |
| ... | ... | @@ -378,18 +341,15 @@ func (h *UploadAssetHandler) Write(s *pssh.SSHServerConnSession, entry *sendutil | |
| 378 | 341 | return "", err | |
| 379 | 342 | } | |
| 380 | 343 | ||
| 381 | - | featureFlag := getFeatureFlag(s) | |
| 382 | - | if featureFlag == nil { | |
| 383 | - | return "", fmt.Errorf("pico+ feature flag ctx not set") | |
| 384 | - | } | |
| 344 | + | featureFlag := state.featureFlag | |
| 385 | 345 | // SFTP does not report file size so the more performant way to | |
| 386 | 346 | // check filesize constraints is to try and upload the file to s3 | |
| 387 | 347 | // with a specialized reader that raises an error if the filesize limit | |
| 388 | 348 | // has been reached | |
| 389 | 349 | storageMax := featureFlag.Data.StorageMax | |
| 390 | 350 | fileMax := featureFlag.Data.FileMax | |
| 391 | - | curStorageSize := getStorageSize(s) | |
| 392 | - | remaining := int64(storageMax) - int64(curStorageSize) | |
| 351 | + | curStorageSize := state.getStorageSize() | |
| 352 | + | remaining := int64(storageMax) - curStorageSize | |
| 393 | 353 | sizeRemaining := min(remaining+curFileSize, fileMax) | |
| 394 | 354 | if sizeRemaining <= 0 { | |
| 395 | 355 | _, _ = fmt.Fprintln(s.Stderr(), "storage quota reached") |
| ... | ... | @@ -428,8 +388,7 @@ func (h *UploadAssetHandler) Write(s *pssh.SSHServerConnSession, entry *sendutil | |
| 428 | 388 | return "", cerr | |
| 429 | 389 | } | |
| 430 | 390 | ||
| 431 | - | deltaFileSize := curFileSize - fsize | |
| 432 | - | nextStorageSize := incrementStorageSize(s, deltaFileSize) | |
| 391 | + | nextStorageSize := state.addStorageSize(fsize - curFileSize) | |
| 433 | 392 | ||
| 434 | 393 | url := h.Cfg.AssetURL( | |
| 435 | 394 | user.Name, |
| ... | ... | @@ -477,11 +436,12 @@ func (h *UploadAssetHandler) Delete(s *pssh.SSHServerConnSession, entry *senduti | |
| 477 | 436 | "file", assetFilepath, | |
| 478 | 437 | ) | |
| 479 | 438 | ||
| 480 | - | bucket, err := getBucket(s) | |
| 439 | + | state, err := getUploadState(s) | |
| 481 | 440 | 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()) | |
| 483 | 442 | return err | |
| 484 | 443 | } | |
| 444 | + | bucket := state.bucket | |
| 485 | 445 | ||
| 486 | 446 | projectName := shared.GetProjectName(entry) | |
| 487 | 447 | logger = logger.With("project", projectName) |
| ... | ... | @@ -490,14 +450,7 @@ func (h *UploadAssetHandler) Delete(s *pssh.SSHServerConnSession, entry *senduti | |
| 490 | 450 | ||
| 491 | 451 | // Check if this path represents a directory (has a . _pico_keep_dir marker) | |
| 492 | 452 | 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 { | |
| 501 | 454 | // This is a directory being deleted. We must delete all nested | |
| 502 | 455 | // . _pico_keep_dir files first, otherwise os.Remove() on the | |
| 503 | 456 | // directory will fail with "directory not empty". |
| ... | ... | @@ -563,7 +516,11 @@ func (h *UploadAssetHandler) Delete(s *pssh.SSHServerConnSession, entry *senduti | |
| 563 | 516 | return err | |
| 564 | 517 | } | |
| 565 | 518 | } | |
| 519 | + | info, statErr := h.Cfg.Storage.StatObject(bucket, assetFilepath) | |
| 566 | 520 | err = h.Cfg.Storage.DeleteObject(bucket, assetFilepath) | |
| 521 | + | if err == nil && statErr == nil { | |
| 522 | + | state.addStorageSize(-info.Size) | |
| 523 | + | } | |
| 567 | 524 | ||
| 568 | 525 | surrogate := getSurrogateKey(user.Name, projectName) | |
| 569 | 526 | h.Cfg.CacheClearingQueue <- surrogate |
| ... | ... | @@ -589,7 +546,7 @@ func (h *UploadAssetHandler) validateAsset(data *FileData) (bool, error) { | |
| 589 | 546 | } | |
| 590 | 547 | ||
| 591 | 548 | fpath := strings.Replace(data.Filepath, "/"+projectName, "", 1) | |
| 592 | - | if shouldIgnoreFile(fpath, data.DenyList) { | |
| 549 | + | if data.DenyList.MatchesPath(fpath) { | |
| 593 | 550 | err := fmt.Errorf( | |
| 594 | 551 | "ERROR: (%s) file rejected, https://pico.sh/pgs#-pgs-ignore", | |
| 595 | 552 | data.Filepath, |
+20,
-0
| ... | ... | @@ -3,8 +3,10 @@ package storage_test | |
| 3 | 3 | import ( | |
| 4 | 4 | "bytes" | |
| 5 | 5 | "crypto/rand" | |
| 6 | + | "errors" | |
| 6 | 7 | "fmt" | |
| 7 | 8 | "io" | |
| 9 | + | "io/fs" | |
| 8 | 10 | "log/slog" | |
| 9 | 11 | "slices" | |
| 10 | 12 | "strings" |
| ... | ... | @@ -106,6 +108,21 @@ func TestStorageConformance(t *testing.T) { | |
| 106 | 108 | } | |
| 107 | 109 | }) | |
| 108 | 110 | ||
| 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 | + | ||
| 109 | 126 | t.Run("list", func(t *testing.T) { | |
| 110 | 127 | now := list(t, st, b, "/proj/now.txt", false)[0] | |
| 111 | 128 | tests := []struct { |
| ... | ... | @@ -214,6 +231,9 @@ func TestStorageConformance(t *testing.T) { | |
| 214 | 231 | if _, _, err := st.GetObject(b, "/proj/css/deep/x.css"); err == nil { | |
| 215 | 232 | t.Fatal("deleted object still readable") | |
| 216 | 233 | } | |
| 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 | + | } | |
| 217 | 237 | got := list(t, st, b, "/proj/css/", true) | |
| 218 | 238 | want := []string{entry("site.css", false, 15, t1)} | |
| 219 | 239 | if !slices.Equal(got, want) { |
+20,
-23
| ... | ... | @@ -1,8 +1,6 @@ | |
| 1 | 1 | package storage | |
| 2 | 2 | ||
| 3 | 3 | import ( | |
| 4 | - | "crypto/md5" | |
| 5 | - | "encoding/hex" | |
| 6 | 4 | "fmt" | |
| 7 | 5 | "io" | |
| 8 | 6 | "io/fs" |
| ... | ... | @@ -120,31 +118,30 @@ func (s *StorageFS) GetObject(bucket Bucket, fpath string) (utils.ReadAndReaderA | |
| 120 | 118 | _ = dat.Close() | |
| 121 | 119 | return nil, objInfo, err | |
| 122 | 120 | } | |
| 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) { | |
| 123 | 129 | objInfo.Size = info.Size() | |
| 124 | 130 | objInfo.LastModified = info.ModTime() | |
| 131 | + | objInfo.ETag = fmt.Sprintf("%x-%x", info.ModTime().UnixNano(), info.Size()) | |
| 132 | + | } | |
| 125 | 133 | ||
| 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 | |
| 144 | 138 | } | |
| 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 | |
| 148 | 145 | } | |
| 149 | 146 | ||
| 150 | 147 | func (s *StorageFS) PutObject(bucket Bucket, fpath string, contents io.Reader, info *ObjectInfo) (string, int64, error) { |
+16,
-0
| ... | ... | @@ -4,6 +4,7 @@ import ( | |
| 4 | 4 | "bytes" | |
| 5 | 5 | "fmt" | |
| 6 | 6 | "io" | |
| 7 | + | "io/fs" | |
| 7 | 8 | "os" | |
| 8 | 9 | "path/filepath" | |
| 9 | 10 | "strings" |
| ... | ... | @@ -119,6 +120,21 @@ func (s *StorageMemory) GetObject(bucket Bucket, fpath string) (utils.ReadAndRea | |
| 119 | 120 | return &seekableReader{bytes.NewReader([]byte(dat))}, objInfo, nil | |
| 120 | 121 | } | |
| 121 | 122 | ||
| 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 | + | ||
| 122 | 138 | func (s *StorageMemory) PutObject(bucket Bucket, fpath string, contents io.Reader, info *ObjectInfo) (string, int64, error) { | |
| 123 | 139 | s.mu.Lock() | |
| 124 | 140 | defer s.mu.Unlock() |
+2,
-0
| ... | ... | @@ -31,6 +31,8 @@ type BucketStorage interface { | |
| 31 | 31 | ||
| 32 | 32 | type ObjectStorage interface { | |
| 33 | 33 | 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) | |
| 34 | 36 | PutObject(bucket Bucket, fpath string, contents io.Reader, info *ObjectInfo) (string, int64, error) | |
| 35 | 37 | DeleteObject(bucket Bucket, fpath string) error | |
| 36 | 38 | ListObjects(bucket Bucket, dir string, recursive bool) ([]os.FileInfo, error) |