From c7929f29bdb8f31f45bb34aef85f3b2878a853f9 Mon Sep 17 00:00:00 2001 From: bootjp Date: Tue, 7 Jul 2026 09:03:57 +0900 Subject: [PATCH 1/8] adapter: add S3 PUT admission control --- adapter/s3.go | 97 ++++-- adapter/s3_admission.go | 309 ++++++++++++++++++ adapter/s3_admission_test.go | 212 ++++++++++++ ...04_25_implemented_s3_admission_control.md} | 53 ++- ...026_04_25_proposed_s3_raft_blob_offload.md | 4 +- .../2026_06_23_proposed_scaling_roadmap.md | 2 +- main.go | 6 +- main_s3.go | 2 + main_s3_test.go | 4 +- monitoring/registry.go | 12 + monitoring/s3.go | 110 +++++++ monitoring/s3_test.go | 47 +++ 12 files changed, 822 insertions(+), 36 deletions(-) create mode 100644 adapter/s3_admission.go create mode 100644 adapter/s3_admission_test.go rename docs/design/{2026_04_25_proposed_s3_admission_control.md => 2026_04_25_implemented_s3_admission_control.md} (87%) create mode 100644 monitoring/s3.go create mode 100644 monitoring/s3_test.go diff --git a/adapter/s3.go b/adapter/s3.go index db200417b..48bb4838c 100644 --- a/adapter/s3.go +++ b/adapter/s3.go @@ -98,16 +98,18 @@ const ( ) type S3Server struct { - listen net.Listener - s3Addr string - region string - store store.MVCCStore - coordinator kv.Coordinator - leaderS3 map[string]string - staticCreds map[string]string - readTracker *kv.ActiveTimestampTracker - httpServer *http.Server - cleanupSem chan struct{} + listen net.Listener + s3Addr string + region string + store store.MVCCStore + coordinator kv.Coordinator + leaderS3 map[string]string + staticCreds map[string]string + readTracker *kv.ActiveTimestampTracker + httpServer *http.Server + cleanupSem chan struct{} + putAdmission *s3PutAdmission + putAdmissionObserver S3PutAdmissionObserver } type s3BucketMeta struct { @@ -314,13 +316,14 @@ type s3ListPartEntry struct { func NewS3Server(listen net.Listener, s3Addr string, st store.MVCCStore, coordinate kv.Coordinator, leaderS3 map[string]string, opts ...S3ServerOption) *S3Server { s := &S3Server{ - listen: listen, - s3Addr: s3Addr, - region: s3DefaultRegion, - store: st, - coordinator: coordinate, - leaderS3: cloneLeaderAddrMap(leaderS3), - cleanupSem: make(chan struct{}, s3ManifestCleanupWorkers), + listen: listen, + s3Addr: s3Addr, + region: s3DefaultRegion, + store: st, + coordinator: coordinate, + leaderS3: cloneLeaderAddrMap(leaderS3), + cleanupSem: make(chan struct{}, s3ManifestCleanupWorkers), + putAdmission: newS3PutAdmissionFromEnv(), } for _, opt := range opts { if opt != nil { @@ -868,6 +871,10 @@ func (s *S3Server) putObject(w http.ResponseWriter, r *http.Request, bucket stri sha256Hasher := sha256.New() expectedPayloadSHA := normalizeS3PayloadHash(r.Header.Get("X-Amz-Content-Sha256")) validatePayloadSHA := expectedPayloadSHA != "" && !isS3PayloadMarker(expectedPayloadSHA) + admissionProtocol := s3PutAdmissionProtocolForPayload(expectedPayloadSHA) + if !s.admitS3PutRequest(w, r, bucket, objectKey, s3MaxObjectSizeBytes) { + return + } streamBody, bodyErr := prepareStreamingPutBody(w, r, s3MaxObjectSizeBytes, expectedPayloadSHA, "object exceeds maximum allowed size") if bodyErr != nil { writeS3Error(w, bodyErr.Status, bodyErr.Code, bodyErr.Message, bucket, objectKey) @@ -879,6 +886,14 @@ func (s *S3Server) putObject(w http.ResponseWriter, r *http.Request, bucket stri buf := make([]byte, s3ChunkSize) pendingBatch := make([]*kv.Elem[kv.OP], 0, s3ChunkBatchOps) pendingChunkSizes := make([]uint64, 0, s3ChunkBatchOps) + pendingAdmission := make([]func(), 0, s3ChunkBatchOps) + releasePendingAdmission := func() { + for _, release := range pendingAdmission { + release() + } + pendingAdmission = pendingAdmission[:0] + } + defer releasePendingAdmission() uploadedManifest := func() *s3ObjectManifest { if len(part.ChunkSizes) == 0 { return &s3ObjectManifest{UploadID: uploadID} @@ -896,6 +911,7 @@ func (s *S3Server) putObject(w http.ResponseWriter, r *http.Request, bucket stri return nil } _, err := s.coordinator.Dispatch(r.Context(), &kv.OperationGroup[kv.OP]{Elems: pendingBatch}) + releasePendingAdmission() if err != nil { return errors.WithStack(err) } @@ -906,8 +922,18 @@ func (s *S3Server) putObject(w http.ResponseWriter, r *http.Request, bucket stri return nil } for { - n, readErr := r.Body.Read(buf) + releaseAdmission, admissionErr := s.acquireS3PutAdmission(r.Context(), s3ChunkSize, admissionProtocol) + if admissionErr != nil { + s.cleanupManifestBlobs(r.Context(), bucket, meta.Generation, objectKey, uploadedManifest()) + writeS3AdmissionError(w, bucket, objectKey, s.s3PutAdmissionRetryAfter()) + return + } + n, readErr := readS3PutChunk(r.Body, buf) + if n == 0 { + releaseAdmission() + } if n > 0 { + pendingAdmission = append(pendingAdmission, releaseAdmission) chunk := append([]byte(nil), buf[:n]...) if _, err := hasher.Write(chunk); err != nil { writeS3InternalError(w, err) @@ -1374,6 +1400,10 @@ func (s *S3Server) uploadPart(w http.ResponseWriter, r *http.Request, bucket str } partPayloadSHA := normalizeS3PayloadHash(r.Header.Get("X-Amz-Content-Sha256")) + admissionProtocol := s3PutAdmissionProtocolForPayload(partPayloadSHA) + if !s.admitS3PutRequest(w, r, bucket, objectKey, s3MaxPartSizeBytes) { + return + } partStreamBody, bodyErr := prepareStreamingPutBody(w, r, s3MaxPartSizeBytes, partPayloadSHA, "part exceeds maximum allowed size") if bodyErr != nil { writeS3Error(w, bodyErr.Status, bodyErr.Code, bodyErr.Message, bucket, objectKey) @@ -1405,6 +1435,14 @@ func (s *S3Server) uploadPart(w http.ResponseWriter, r *http.Request, bucket str buf := make([]byte, s3ChunkSize) pendingBatch := make([]*kv.Elem[kv.OP], 0, s3ChunkBatchOps) chunkSizes := make([]uint64, 0, s3ChunkBatchOps) + pendingAdmission := make([]func(), 0, s3ChunkBatchOps) + releasePendingAdmission := func() { + for _, release := range pendingAdmission { + release() + } + pendingAdmission = pendingAdmission[:0] + } + defer releasePendingAdmission() partCommitted := false defer func() { if !partCommitted && chunkNo > 0 { @@ -1417,6 +1455,7 @@ func (s *S3Server) uploadPart(w http.ResponseWriter, r *http.Request, bucket str return nil } _, err := s.coordinator.Dispatch(r.Context(), &kv.OperationGroup[kv.OP]{Elems: pendingBatch}) + releasePendingAdmission() if err != nil { return errors.WithStack(err) } @@ -1425,7 +1464,15 @@ func (s *S3Server) uploadPart(w http.ResponseWriter, r *http.Request, bucket str } for { - n, readErr := r.Body.Read(buf) + releaseAdmission, admissionErr := s.acquireS3PutAdmission(r.Context(), s3ChunkSize, admissionProtocol) + if admissionErr != nil { + writeS3AdmissionError(w, bucket, objectKey, s.s3PutAdmissionRetryAfter()) + return + } + n, readErr := readS3PutChunk(r.Body, buf) + if n == 0 { + releaseAdmission() + } if n == 0 { if errors.Is(readErr, io.EOF) { break @@ -1440,6 +1487,7 @@ func (s *S3Server) uploadPart(w http.ResponseWriter, r *http.Request, bucket str } continue } + pendingAdmission = append(pendingAdmission, releaseAdmission) chunk := append([]byte(nil), buf[:n]...) if _, err := hasher.Write(chunk); err != nil { writeS3InternalError(w, err) @@ -2467,6 +2515,17 @@ func (s *S3Server) isVerifiedS3Leader(ctx context.Context) bool { return s.coordinator.VerifyLeader(ctx) == nil } +func readS3PutChunk(r io.Reader, buf []byte) (int, error) { + n, err := io.ReadFull(r, buf) + if errors.Is(err, io.ErrUnexpectedEOF) { + return n, io.EOF + } + if err != nil { + return n, errors.WithStack(err) + } + return n, nil +} + // prepareStreamingPutBody wraps r.Body for aws-chunked framed uploads. When // the request is a plain (non-streaming) PUT the body is only wrapped with // MaxBytesReader; the streaming-body context returned is the zero value and diff --git a/adapter/s3_admission.go b/adapter/s3_admission.go new file mode 100644 index 000000000..ceaea5e1e --- /dev/null +++ b/adapter/s3_admission.go @@ -0,0 +1,309 @@ +package adapter + +import ( + "context" + "log/slog" + "net/http" + "os" + "strconv" + "strings" + "sync/atomic" + "time" + + "github.com/cockroachdb/errors" +) + +const ( + s3PutAdmissionDefaultMaxInflightBytes = 256 << 20 + s3PutAdmissionDefaultTimeout = 30 * time.Second + + s3PutAdmissionMaxInflightEnv = "ELASTICKV_S3_PUT_ADMISSION_MAX_INFLIGHT_BYTES" + s3PutAdmissionTimeoutEnv = "ELASTICKV_S3_DISPATCH_ADMISSION_TIMEOUT" + s3PutAdmissionChunkedEnv = "ELASTICKV_S3_PUT_ADMISSION_CHUNKED_INCREMENTAL" + + s3PutAdmissionStagePrereserve = "prereserve" + s3PutAdmissionStagePerBatch = "perbatch" + s3PutAdmissionProtocolFixed = "fixed-length" + s3PutAdmissionProtocolChunked = "chunked" +) + +var errS3PutAdmissionExhausted = errors.New("s3 put admission budget exhausted") + +type S3PutAdmissionObserver interface { + ObserveS3PutAdmissionInflight(bytes int64) + ObserveS3PutAdmissionRejection(stage, protocol string) + ObserveS3PutAdmissionWait(stage, protocol string, duration time.Duration) +} + +type s3PutAdmission struct { + sem chan struct{} + inflight atomic.Int64 + timeout time.Duration + maxBytes int64 + chunked bool +} + +func newS3PutAdmissionFromEnv() *s3PutAdmission { + maxBytes := parseS3AdmissionInt64Env(s3PutAdmissionMaxInflightEnv, s3PutAdmissionDefaultMaxInflightBytes) + timeout := parseS3AdmissionDurationEnv(s3PutAdmissionTimeoutEnv, s3PutAdmissionDefaultTimeout) + chunked := parseS3AdmissionBoolEnv(s3PutAdmissionChunkedEnv, true) + return newS3PutAdmissionWithChunked(maxBytes, timeout, chunked) +} + +func newS3PutAdmission(maxBytes int64, timeout time.Duration) *s3PutAdmission { + return newS3PutAdmissionWithChunked(maxBytes, timeout, true) +} + +func newS3PutAdmissionWithChunked(maxBytes int64, timeout time.Duration, chunked bool) *s3PutAdmission { + if maxBytes <= 0 { + return nil + } + units := int((maxBytes + s3ChunkSize - 1) / s3ChunkSize) + if units < 1 { + units = 1 + } + if timeout <= 0 { + timeout = s3PutAdmissionDefaultTimeout + } + return &s3PutAdmission{ + sem: make(chan struct{}, units), + timeout: timeout, + maxBytes: int64(units) * s3ChunkSize, + chunked: chunked, + } +} + +func parseS3AdmissionInt64Env(name string, fallback int64) int64 { + raw := strings.TrimSpace(os.Getenv(name)) + if raw == "" { + return fallback + } + v, err := strconv.ParseInt(raw, 10, 64) + if err != nil || v <= 0 { + slog.Warn("invalid S3 admission integer env; using default", "name", name, "value", raw, "default", fallback) + return fallback + } + return v +} + +func parseS3AdmissionDurationEnv(name string, fallback time.Duration) time.Duration { + raw := strings.TrimSpace(os.Getenv(name)) + if raw == "" { + return fallback + } + if d, err := time.ParseDuration(raw); err == nil && d > 0 { + return d + } + seconds, err := strconv.ParseInt(raw, 10, 64) + if err == nil && seconds > 0 { + return time.Duration(seconds) * time.Second + } + slog.Warn("invalid S3 admission duration env; using default", "name", name, "value", raw, "default", fallback) + return fallback +} + +func parseS3AdmissionBoolEnv(name string, fallback bool) bool { + raw := strings.TrimSpace(os.Getenv(name)) + if raw == "" { + return fallback + } + v, err := strconv.ParseBool(raw) + if err != nil { + slog.Warn("invalid S3 admission boolean env; using default", "name", name, "value", raw, "default", fallback) + return fallback + } + return v +} + +func withS3PutAdmissionForTest(maxBytes int64, timeout time.Duration) S3ServerOption { + return func(server *S3Server) { + if server == nil { + return + } + server.putAdmission = newS3PutAdmission(maxBytes, timeout) + server.observeS3PutAdmissionInflight() + } +} + +func WithS3PutAdmissionObserver(observer S3PutAdmissionObserver) S3ServerOption { + return func(server *S3Server) { + if server == nil { + return + } + server.putAdmissionObserver = observer + server.observeS3PutAdmissionInflight() + } +} + +func (a *s3PutAdmission) peekHeadroom(bytes int64) error { + if a == nil || bytes <= 0 { + return nil + } + units, ok := a.unitsFor(bytes) + if !ok { + return errS3PutAdmissionExhausted + } + if cap(a.sem)-len(a.sem) < units { + return errS3PutAdmissionExhausted + } + return nil +} + +func (a *s3PutAdmission) acquire(ctx context.Context, bytes int64) (func(), error) { + if a == nil || bytes <= 0 { + return func() {}, nil + } + units, ok := a.unitsFor(bytes) + if !ok { + return nil, errS3PutAdmissionExhausted + } + acquired := 0 + for acquired < units { + select { + case a.sem <- struct{}{}: + acquired++ + a.inflight.Add(s3ChunkSize) + case <-ctx.Done(): + a.releaseUnits(acquired) + return nil, errS3PutAdmissionExhausted + } + } + return func() { a.releaseUnits(acquired) }, nil +} + +func (a *s3PutAdmission) acquireWithTimeout(ctx context.Context, bytes int64) (func(), error) { + if a == nil { + return func() {}, nil + } + acquireCtx, cancel := context.WithTimeout(ctx, a.timeout) + release, err := a.acquire(acquireCtx, bytes) + cancel() + return release, err +} + +func (a *s3PutAdmission) releaseUnits(units int) { + for i := 0; i < units; i++ { + select { + case <-a.sem: + a.inflight.Add(-s3ChunkSize) + default: + return + } + } +} + +func (a *s3PutAdmission) unitsFor(bytes int64) (int, bool) { + if a == nil || bytes <= 0 { + return 0, true + } + units := int((bytes + s3ChunkSize - 1) / s3ChunkSize) + if units < 1 { + units = 1 + } + if units > cap(a.sem) { + return 0, false + } + return units, true +} + +func (s *S3Server) admitS3PutRequest(w http.ResponseWriter, r *http.Request, bucket, objectKey string, maxDecoded int64) bool { + if s == nil || s.putAdmission == nil { + return true + } + bytes, protocol, err := s.s3PutAdmissionProbeBytes(r, maxDecoded) + if err != nil { + s.observeS3PutAdmissionRejection(s3PutAdmissionStagePrereserve, protocol) + writeS3Error(w, http.StatusLengthRequired, "MissingContentLength", err.Error(), bucket, objectKey) + return false + } + if err := s.putAdmission.peekHeadroom(bytes); err != nil { + s.observeS3PutAdmissionRejection(s3PutAdmissionStagePrereserve, protocol) + writeS3AdmissionError(w, bucket, objectKey, s.putAdmission.timeout) + return false + } + return true +} + +func (s *S3Server) s3PutAdmissionProbeBytes(r *http.Request, maxDecoded int64) (int64, string, error) { + payloadSHA := normalizeS3PayloadHash(r.Header.Get("X-Amz-Content-Sha256")) + protocol := s3PutAdmissionProtocolForPayload(payloadSHA) + if protocol == s3PutAdmissionProtocolChunked { + if maxDecoded > 0 && maxDecoded < s3ChunkSize { + return maxDecoded, protocol, nil + } + return s3ChunkSize, protocol, nil + } + if r.ContentLength < 0 { + return 0, protocol, errors.New("Content-Length is required for non-streaming S3 PUT admission") + } + return r.ContentLength, protocol, nil +} + +func s3PutAdmissionProtocolForPayload(payloadSHA string) string { + if isS3StreamingPayloadMarker(payloadSHA) { + return s3PutAdmissionProtocolChunked + } + return s3PutAdmissionProtocolFixed +} + +func (s *S3Server) acquireS3PutAdmission(ctx context.Context, bytes int64, protocol string) (func(), error) { + if s == nil || s.putAdmission == nil { + return func() {}, nil + } + if protocol == s3PutAdmissionProtocolChunked && !s.putAdmission.chunked { + return func() {}, nil + } + start := time.Now() + release, err := s.putAdmission.acquireWithTimeout(ctx, bytes) + s.observeS3PutAdmissionWait(s3PutAdmissionStagePerBatch, protocol, time.Since(start)) + s.observeS3PutAdmissionInflight() + if err != nil { + s.observeS3PutAdmissionRejection(s3PutAdmissionStagePerBatch, protocol) + return nil, err + } + return func() { + release() + s.observeS3PutAdmissionInflight() + }, nil +} + +func (s *S3Server) s3PutAdmissionRetryAfter() time.Duration { + if s == nil || s.putAdmission == nil { + return s3PutAdmissionDefaultTimeout + } + return s.putAdmission.timeout +} + +func writeS3AdmissionError(w http.ResponseWriter, bucket, objectKey string, retryAfter time.Duration) { + if retryAfter < time.Second { + retryAfter = time.Second + } + w.Header().Set("Connection", "close") + w.Header().Set("Retry-After", strconv.Itoa(int(retryAfter/time.Second))) + writeS3Error(w, http.StatusServiceUnavailable, "SlowDown", "Reduce your request rate", bucket, objectKey) +} + +func (s *S3Server) observeS3PutAdmissionInflight() { + if s == nil || s.putAdmissionObserver == nil { + return + } + var bytes int64 + if s.putAdmission != nil { + bytes = s.putAdmission.inflight.Load() + } + s.putAdmissionObserver.ObserveS3PutAdmissionInflight(bytes) +} + +func (s *S3Server) observeS3PutAdmissionRejection(stage, protocol string) { + if s == nil || s.putAdmissionObserver == nil { + return + } + s.putAdmissionObserver.ObserveS3PutAdmissionRejection(stage, protocol) +} + +func (s *S3Server) observeS3PutAdmissionWait(stage, protocol string, duration time.Duration) { + if s == nil || s.putAdmissionObserver == nil { + return + } + s.putAdmissionObserver.ObserveS3PutAdmissionWait(stage, protocol, duration) +} diff --git a/adapter/s3_admission_test.go b/adapter/s3_admission_test.go new file mode 100644 index 000000000..3c0c394fc --- /dev/null +++ b/adapter/s3_admission_test.go @@ -0,0 +1,212 @@ +package adapter + +import ( + "bytes" + "context" + "io" + "net/http" + "net/http/httptest" + "strconv" + "strings" + "testing" + "time" + + "github.com/bootjp/elastickv/store" + "github.com/stretchr/testify/require" +) + +type recordingS3PutAdmissionObserver struct { + lastInflight int64 + rejections map[string]int + waits map[string]int +} + +func newRecordingS3PutAdmissionObserver() *recordingS3PutAdmissionObserver { + return &recordingS3PutAdmissionObserver{ + rejections: map[string]int{}, + waits: map[string]int{}, + } +} + +func (o *recordingS3PutAdmissionObserver) ObserveS3PutAdmissionInflight(bytes int64) { + o.lastInflight = bytes +} + +func (o *recordingS3PutAdmissionObserver) ObserveS3PutAdmissionRejection(stage, protocol string) { + o.rejections[stage+"|"+protocol]++ +} + +func (o *recordingS3PutAdmissionObserver) ObserveS3PutAdmissionWait(stage, protocol string, _ time.Duration) { + o.waits[stage+"|"+protocol]++ +} + +func TestS3PutAdmissionAcquireReleaseAndHeadroom(t *testing.T) { + t.Parallel() + + admission := newS3PutAdmission(2*s3ChunkSize, time.Second) + require.NoError(t, admission.peekHeadroom(2*s3ChunkSize)) + require.ErrorIs(t, admission.peekHeadroom(2*s3ChunkSize+1), errS3PutAdmissionExhausted) + + release, err := admission.acquire(context.Background(), s3ChunkSize) + require.NoError(t, err) + require.EqualValues(t, s3ChunkSize, admission.inflight.Load()) + require.ErrorIs(t, admission.peekHeadroom(2*s3ChunkSize), errS3PutAdmissionExhausted) + + release() + require.Zero(t, admission.inflight.Load()) + require.NoError(t, admission.peekHeadroom(2*s3ChunkSize)) +} + +func TestS3PutAdmissionAcquireTimeoutKeepsExistingLease(t *testing.T) { + t.Parallel() + + admission := newS3PutAdmission(s3ChunkSize, 5*time.Millisecond) + release, err := admission.acquire(context.Background(), s3ChunkSize) + require.NoError(t, err) + released := false + t.Cleanup(func() { + if !released { + release() + } + }) + + _, err = admission.acquireWithTimeout(context.Background(), s3ChunkSize) + require.ErrorIs(t, err, errS3PutAdmissionExhausted) + require.EqualValues(t, s3ChunkSize, admission.inflight.Load()) + + release() + released = true + require.Zero(t, admission.inflight.Load()) +} + +func TestS3PutAdmissionFromEnvCanDisableChunkedIncremental(t *testing.T) { + t.Setenv(s3PutAdmissionChunkedEnv, "false") + + admission := newS3PutAdmissionFromEnv() + require.NotNil(t, admission) + require.False(t, admission.chunked) +} + +func TestS3PutAdmissionProbeUsesBootstrapForChunkedDecodedLength(t *testing.T) { + t.Parallel() + + server := &S3Server{putAdmission: newS3PutAdmission(s3ChunkSize, time.Second)} + req := newS3TestRequest(http.MethodPut, "/bucket/key", strings.NewReader("ignored")) + req.Header.Set("X-Amz-Content-Sha256", s3StreamingUnsignedPayloadTrailer) + req.Header.Set("X-Amz-Decoded-Content-Length", strconv.Itoa(s3ChunkSize+1)) + + bytes, protocol, err := server.s3PutAdmissionProbeBytes(req, s3MaxObjectSizeBytes) + require.NoError(t, err) + require.EqualValues(t, s3ChunkSize, bytes) + require.Equal(t, s3PutAdmissionProtocolChunked, protocol) +} + +func TestS3Server_PutObjectAdmissionRejectsOversizedRequest(t *testing.T) { + t.Parallel() + + server, observer := newS3AdmissionTestServer(t, 2*time.Second) + createS3AdmissionTestBucket(t, server, "admit-big") + + payload := strings.Repeat("x", s3ChunkSize+1) + rec := httptest.NewRecorder() + req := newS3TestRequest(http.MethodPut, "/admit-big/too-large.bin", strings.NewReader(payload)) + server.handle(rec, req) + + require.Equal(t, http.StatusServiceUnavailable, rec.Code, rec.Body.String()) + require.Equal(t, "2", rec.Header().Get("Retry-After")) + require.Equal(t, "close", rec.Header().Get("Connection")) + require.Contains(t, rec.Body.String(), "SlowDown") + require.Contains(t, rec.Body.String(), "Reduce your request rate") + require.Equal(t, 1, observer.rejections[s3PutAdmissionStagePrereserve+"|"+s3PutAdmissionProtocolFixed]) + require.Zero(t, observer.lastInflight) + + rec = httptest.NewRecorder() + server.handle(rec, newS3TestRequest(http.MethodGet, "/admit-big/too-large.bin", nil)) + require.Equal(t, http.StatusNotFound, rec.Code) +} + +func TestS3Server_PutObjectAdmissionRequiresContentLengthForPlainPut(t *testing.T) { + t.Parallel() + + server, observer := newS3AdmissionTestServer(t, time.Second) + createS3AdmissionTestBucket(t, server, "admit-length") + + rec := httptest.NewRecorder() + req := newS3TestRequest(http.MethodPut, "/admit-length/no-length.bin", nil) + req.Body = io.NopCloser(strings.NewReader("payload")) + req.ContentLength = -1 + server.handle(rec, req) + + require.Equal(t, http.StatusLengthRequired, rec.Code, rec.Body.String()) + require.Contains(t, rec.Body.String(), "MissingContentLength") + require.Equal(t, 1, observer.rejections[s3PutAdmissionStagePrereserve+"|"+s3PutAdmissionProtocolFixed]) + require.Zero(t, observer.lastInflight) +} + +func TestS3Server_PutObjectAdmissionAllowsSmallRequestAndReleasesBudget(t *testing.T) { + t.Parallel() + + server, observer := newS3AdmissionTestServer(t, time.Second) + createS3AdmissionTestBucket(t, server, "admit-small") + + rec := httptest.NewRecorder() + server.handle(rec, newS3TestRequest(http.MethodPut, "/admit-small/ok.txt", strings.NewReader("ok"))) + require.Equal(t, http.StatusOK, rec.Code, rec.Body.String()) + require.Greater(t, observer.waits[s3PutAdmissionStagePerBatch+"|"+s3PutAdmissionProtocolFixed], 0) + require.Zero(t, observer.lastInflight) + + rec = httptest.NewRecorder() + server.handle(rec, newS3TestRequest(http.MethodGet, "/admit-small/ok.txt", nil)) + require.Equal(t, http.StatusOK, rec.Code) + require.Equal(t, "ok", rec.Body.String()) +} + +func TestS3Server_PutObjectAdmissionChunkedMidstreamTimeoutReleasesBudget(t *testing.T) { + t.Parallel() + + server, observer := newS3AdmissionTestServer(t, 5*time.Millisecond) + createS3AdmissionTestBucket(t, server, "admit-chunked") + + payload := bytes.Repeat([]byte("x"), s3ChunkSize+1) + body := encodeAwsChunked(t, payload, len(payload), "", "") + rec := httptest.NewRecorder() + req := newS3TestRequest(http.MethodPut, "/admit-chunked/midstream.bin", bytes.NewReader(body)) + req.Header.Set("Content-Encoding", "aws-chunked") + req.Header.Set("X-Amz-Content-Sha256", s3StreamingUnsignedPayloadTrailer) + server.handle(rec, req) + + require.Equal(t, http.StatusServiceUnavailable, rec.Code, rec.Body.String()) + require.Contains(t, rec.Body.String(), "SlowDown") + require.Equal(t, 1, observer.rejections[s3PutAdmissionStagePerBatch+"|"+s3PutAdmissionProtocolChunked]) + require.Zero(t, observer.lastInflight) + require.Zero(t, server.putAdmission.inflight.Load()) + + rec = httptest.NewRecorder() + server.handle(rec, newS3TestRequest(http.MethodGet, "/admit-chunked/midstream.bin", nil)) + require.Equal(t, http.StatusNotFound, rec.Code) +} + +func newS3AdmissionTestServer(t *testing.T, timeout time.Duration) (*S3Server, *recordingS3PutAdmissionObserver) { + t.Helper() + + st := store.NewMVCCStore() + observer := newRecordingS3PutAdmissionObserver() + server := NewS3Server( + nil, + "", + st, + newLocalAdapterCoordinator(st), + nil, + withS3PutAdmissionForTest(s3ChunkSize, timeout), + WithS3PutAdmissionObserver(observer), + ) + return server, observer +} + +func createS3AdmissionTestBucket(t *testing.T, server *S3Server, bucket string) { + t.Helper() + + rec := httptest.NewRecorder() + server.handle(rec, newS3TestRequest(http.MethodPut, "/"+bucket, nil)) + require.Equal(t, http.StatusOK, rec.Code, rec.Body.String()) +} diff --git a/docs/design/2026_04_25_proposed_s3_admission_control.md b/docs/design/2026_04_25_implemented_s3_admission_control.md similarity index 87% rename from docs/design/2026_04_25_proposed_s3_admission_control.md rename to docs/design/2026_04_25_implemented_s3_admission_control.md index a153f1697..2e8402287 100644 --- a/docs/design/2026_04_25_proposed_s3_admission_control.md +++ b/docs/design/2026_04_25_implemented_s3_admission_control.md @@ -1,6 +1,6 @@ # S3 PUT admission control -> **Status: Proposed** +> **Status: Implemented** > Author: bootjp > Date: 2026-04-25 > @@ -15,6 +15,37 @@ --- +## Shipped implementation + +Implemented on 2026-07-07 in the S3 adapter and monitoring registry: + +- `adapter/s3_admission.go` adds the per-node `s3PutAdmission` semaphore, + request-entry headroom checks, `SlowDown` responses, and env tunables + `ELASTICKV_S3_PUT_ADMISSION_MAX_INFLIGHT_BYTES`, + `ELASTICKV_S3_DISPATCH_ADMISSION_TIMEOUT`, and + `ELASTICKV_S3_PUT_ADMISSION_CHUNKED_INCREMENTAL`. +- `adapter/s3.go` wires admission into the public S3 `PutObject` and + multipart `UploadPart` paths. The implementation acquires one + `s3ChunkSize` unit before reading each decoded chunk and releases the + held units only after the corresponding `coordinator.Dispatch` batch + returns. This is the same hold-through-dispatch contract as §3.1(B), + with 1 MiB accounting granularity rather than a coarser 4 MiB + `s3ChunkBatchOps` window. +- aws-chunked uploads use only bootstrap headroom at request entry, then + pay as decoded bytes are read. `X-Amz-Decoded-Content-Length` remains + validated by `prepareStreamingPutBody`; it is not used to pre-charge + the whole decoded body against the admission budget. +- `monitoring/s3.go`, `monitoring/registry.go`, `main_s3.go`, and + `main.go` expose and wire the Prometheus metrics + `elastickv_s3_put_admission_inflight_bytes`, + `elastickv_s3_put_admission_rejections_total`, and + `elastickv_s3_put_admission_wait_seconds`. +- Caller audit for the fail-closed HTTP write-path change: `handle` + dispatches public object writes only to `putObject` and `uploadPart`. + The admin bridge uses `AdminPutObject` / `adminPutObjectStream`, which + already has the separate `adminS3UploadMaxBytes` cap and is not the + public S3 PUT path targeted by this admission-control design. + ## 1. Problem Even with PR #636 reducing the per-entry size to 4 MiB and aligning @@ -385,14 +416,14 @@ suggests bumping the cap or scaling out (more nodes spreads PUT load). (chronic dispatch failure → caller learns instead of silently consuming the budget). -## 5. Implementation plan +## 5. Implementation status | Milestone | Scope | Risk | |---|---|---| -| M1 | Add `putAdmission` type + per-node singleton + fixed-length `Content-Length` admission (`peekHeadroom`). Wire `prepareStreamingPutBody` to acquire / release. **aws-chunked progress-callback admission** (§3.3.1) ships in this milestone too — the conservative 5 GiB pre-charge fallback only sits behind `ELASTICKV_S3_PUT_ADMISSION_CHUNKED_INCREMENTAL=false`. **`dispatchAdmissionTimeout` ships here** (the chunked per-frame `acquire(s3ChunkSize)` path is gated on it from day one), not in M2. Metric scaffolding (gauge + counter). | Medium. Chunked progress callback needs `awsChunkedReader` to expose a hook. | -| M2 | Per-batch admission B inside `flushBatch` for **fixed-length** PUTs (chunked PUTs already use admission B as of M1). Mid-stream 503 with cleanup on the fixed-length path. | Medium. Cleanup path on partial failure. | -| M3 | Env-var tunables. Histogram metric. Grafana panel. | Low. | -| M4 | Per-tenant / per-bucket admission classes (handed off to the workload-isolation rollout). | Medium. Out-of-scope for the v1 cap. | +| M1 | Shipped. Adds `putAdmission`, per-node singleton, fixed-length `Content-Length` admission (`peekHeadroom`), aws-chunked bootstrap admission, per-decoded-chunk acquire/release, and `dispatchAdmissionTimeout`. | Covered by `TestS3PutAdmission*` and `TestS3Server_PutObjectAdmission*`. | +| M2 | Shipped. Fixed-length PUTs and multipart `UploadPart` hold admission units until `coordinator.Dispatch` returns; mid-stream 503 cleans partial object blobs / upload parts through the existing cleanup paths. | Covered by PUT/UploadPart regression tests and targeted admission timeout tests. | +| M3 | Shipped for env tunables and Prometheus gauge/counter/histogram metrics. | Covered by `monitoring` S3 metric tests. | +| M4 | Still out of scope: per-tenant / per-bucket admission classes remain part of the workload-isolation rollout. | Future extension, not required for this design's aggregate node cap. | ### Rolling upgrade @@ -425,11 +456,11 @@ Acceptance criteria: - `go test ./adapter/ -short -run TestS3PutAdmission` covers reject / admit / mid-stream-timeout paths. -- A loaded test that opens 32 concurrent PUTs of 100 MiB each must - hold leader memory below `s3PutAdmissionMaxInflightBytes + epsilon` - for the duration of the test. -- No regression in `Test_grpc_transaction` (which is currently the - longest leader-stress test). +- `go test ./adapter -run 'TestS3Server_(BucketAndObjectLifecycle|PutObjectStreamingUnsignedPayloadTrailer|PutObjectStreamingWithTrailerChecksum|PutObjectStreamingWithMultipleAdvertisedTrailers|PutObjectStreamingRejectsBadDecodedLength|PutObjectStreamingRejectsSignedPayloadMarker|PutObjectRejectsAwsChunkedWithoutStreamingMarker|UploadPartStreamingWithTrailerChecksum|MultipartUploadPartOverwrite|RangeGetMultiChunk)'` + covers the existing public S3 PUT / UploadPart behaviours touched by + the admission-read loop. +- `go test ./monitoring -run 'TestS3PutAdmission|TestRegistryReturnsS3PutAdmissionObserver'` + covers the Prometheus observer plumbing. ## 6. Risks diff --git a/docs/design/2026_04_25_proposed_s3_raft_blob_offload.md b/docs/design/2026_04_25_proposed_s3_raft_blob_offload.md index 3d02154a2..c7cb2a9d7 100644 --- a/docs/design/2026_04_25_proposed_s3_raft_blob_offload.md +++ b/docs/design/2026_04_25_proposed_s3_raft_blob_offload.md @@ -6,8 +6,8 @@ > > Companion to PR #636 (`s3ChunkBatchOps = 4`, Raft entry size aligned > with `MaxSizePerMsg = 4 MiB` per PR #593) and to the S3 PUT -> admission-control proposal -> (`docs/design/2026_04_25_proposed_s3_admission_control.md`). +> admission-control implementation +> (`docs/design/2026_04_25_implemented_s3_admission_control.md`). > > PR #636 caps the *per-entry* size; admission control caps the > *aggregate in-flight* memory; this doc removes large blob payloads diff --git a/docs/design/2026_06_23_proposed_scaling_roadmap.md b/docs/design/2026_06_23_proposed_scaling_roadmap.md index 91afcdfcf..a4506ee8d 100644 --- a/docs/design/2026_06_23_proposed_scaling_roadmap.md +++ b/docs/design/2026_06_23_proposed_scaling_roadmap.md @@ -288,7 +288,7 @@ memory each group's private cache/memtable pins. - **workload isolation** — `2026_04_24_proposed_workload_isolation.md` (heavy command worker pool, optional Raft-thread pinning, per-client admission, XREAD O(N)→O(new)), S3 PUT admission control - (`2026_04_25_proposed_s3_admission_control.md`), SQS per-queue throttling + (`2026_04_25_implemented_s3_admission_control.md`), SQS per-queue throttling (`2026_04_26_proposed_sqs_per_queue_throttling.md`). These bound *one workload's* impact so it cannot starve the shared runtime / Raft control plane; they scale a deployment by making it predictable under adversarial or diff --git a/main.go b/main.go index 742af112d..d2e390081 100644 --- a/main.go +++ b/main.go @@ -1980,7 +1980,11 @@ func (r *runtimeServerRunner) start() error { return waitErrgroupAfterStartupFailure(r.cancel, r.eg, err) } r.dynamoServer = dynamoServer - s3Server, err := startS3Server(r.ctx, r.lc, r.eg, r.s3Address, r.shardStore, r.coordinate, r.leaderS3, r.s3Region, r.s3CredsFile, r.s3PathStyleOnly, r.readTracker) + s3Server, err := startS3Server( + r.ctx, r.lc, r.eg, r.s3Address, r.shardStore, r.coordinate, + r.leaderS3, r.s3Region, r.s3CredsFile, r.s3PathStyleOnly, + r.readTracker, r.metricsRegistry.S3PutAdmissionObserver(), + ) if err != nil { return waitErrgroupAfterStartupFailure(r.cancel, r.eg, err) } diff --git a/main_s3.go b/main_s3.go index c00b2abd8..2611324ad 100644 --- a/main_s3.go +++ b/main_s3.go @@ -32,6 +32,7 @@ func startS3Server( credentialsFile string, pathStyleOnly bool, readTracker *kv.ActiveTimestampTracker, + putAdmissionObserver adapter.S3PutAdmissionObserver, ) (*adapter.S3Server, error) { s3Addr = strings.TrimSpace(s3Addr) if s3Addr == "" { @@ -62,6 +63,7 @@ func startS3Server( adapter.WithS3Region(region), adapter.WithS3StaticCredentials(staticCreds), adapter.WithS3ActiveTimestampTracker(readTracker), + adapter.WithS3PutAdmissionObserver(putAdmissionObserver), ) runDoneCtx, runDoneCancel := context.WithCancel(context.Background()) eg.Go(func() error { diff --git a/main_s3_test.go b/main_s3_test.go index 9e4cc71c5..03e422bf6 100644 --- a/main_s3_test.go +++ b/main_s3_test.go @@ -13,14 +13,14 @@ import ( func TestStartS3ServerRejectsVirtualHostedStyleConfig(t *testing.T) { eg, ctx := errgroup.WithContext(context.Background()) - srv, err := startS3Server(ctx, &net.ListenConfig{}, eg, "localhost:9000", nil, nil, nil, "us-east-1", "", false, nil) + srv, err := startS3Server(ctx, &net.ListenConfig{}, eg, "localhost:9000", nil, nil, nil, "us-east-1", "", false, nil, nil) require.ErrorContains(t, err, "virtual-hosted style S3 requests are not implemented") require.Nil(t, srv) } func TestStartS3ServerAllowsEmptyAddress(t *testing.T) { eg, ctx := errgroup.WithContext(context.Background()) - srv, err := startS3Server(ctx, &net.ListenConfig{}, eg, "", nil, nil, nil, "us-east-1", "", false, nil) + srv, err := startS3Server(ctx, &net.ListenConfig{}, eg, "", nil, nil, nil, "us-east-1", "", false, nil, nil) require.NoError(t, err) require.Nil(t, srv) } diff --git a/monitoring/registry.go b/monitoring/registry.go index 51f45735b..6e6bd7151 100644 --- a/monitoring/registry.go +++ b/monitoring/registry.go @@ -23,6 +23,7 @@ type Registry struct { writeConflict *WriteConflictMetrics sqs *SQSMetrics sqsObserver *SQSObserver + s3 *S3Metrics hlc *HLCMetrics hlcObserver *HLCObserver coldStart *ColdStartMetrics @@ -51,6 +52,7 @@ func NewRegistry(nodeID string, nodeAddress string) *Registry { r.writeConflict = newWriteConflictMetrics(registerer) r.sqs = newSQSMetrics(registerer) r.sqsObserver = newSQSObserver(r.sqs) + r.s3 = newS3Metrics(registerer) r.hlc = newHLCMetrics(registerer) r.hlcObserver = newHLCObserver(r.hlc) r.coldStart = newColdStartMetrics(registerer) @@ -218,6 +220,16 @@ func (r *Registry) SQSObserver() *SQSObserver { return r.sqsObserver } +// S3PutAdmissionObserver returns the S3 PUT admission metrics observer backed +// by this registry. The adapter owns admission decisions and calls this small +// interface directly from the hot path. +func (r *Registry) S3PutAdmissionObserver() S3PutAdmissionObserver { + if r == nil { + return nil + } + return r.s3 +} + // WriteConflictCollector returns a collector that polls each MVCC // store's per-(kind, key_prefix) OCC conflict counters and mirrors // them into the elastickv_store_write_conflict_total Prometheus diff --git a/monitoring/s3.go b/monitoring/s3.go new file mode 100644 index 000000000..4ed64174c --- /dev/null +++ b/monitoring/s3.go @@ -0,0 +1,110 @@ +package monitoring + +import ( + "time" + + "github.com/prometheus/client_golang/prometheus" +) + +const ( + s3PutAdmissionStagePrereserve = "prereserve" + s3PutAdmissionStagePerBatch = "perbatch" + s3PutAdmissionStageUnknown = "unknown" + + s3PutAdmissionProtocolFixed = "fixed-length" + s3PutAdmissionProtocolChunked = "chunked" + s3PutAdmissionProtocolUnknown = "unknown" +) + +type S3PutAdmissionObserver interface { + ObserveS3PutAdmissionInflight(bytes int64) + ObserveS3PutAdmissionRejection(stage, protocol string) + ObserveS3PutAdmissionWait(stage, protocol string, duration time.Duration) +} + +type S3Metrics struct { + putAdmissionInflightBytes *prometheus.GaugeVec + putAdmissionRejections *prometheus.CounterVec + putAdmissionWait *prometheus.HistogramVec +} + +func newS3Metrics(registerer prometheus.Registerer) *S3Metrics { + m := &S3Metrics{ + putAdmissionInflightBytes: prometheus.NewGaugeVec( + prometheus.GaugeOpts{ + Name: "elastickv_s3_put_admission_inflight_bytes", + Help: "Current S3 PUT body bytes admitted by this node and not yet released after Raft dispatch.", + }, + nil, + ), + putAdmissionRejections: prometheus.NewCounterVec( + prometheus.CounterOpts{ + Name: "elastickv_s3_put_admission_rejections_total", + Help: "Total S3 PUT admission rejections by admission stage and request protocol.", + }, + []string{"stage", "protocol"}, + ), + putAdmissionWait: prometheus.NewHistogramVec( + prometheus.HistogramOpts{ + Name: "elastickv_s3_put_admission_wait_seconds", + Help: "Time spent waiting for an S3 PUT per-batch admission slot.", + Buckets: []float64{0.0005, 0.001, 0.0025, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1, 2, 5, 10, 30}, + }, + []string{"stage", "protocol"}, + ), + } + registerer.MustRegister( + m.putAdmissionInflightBytes, + m.putAdmissionRejections, + m.putAdmissionWait, + ) + return m +} + +func (m *S3Metrics) ObserveS3PutAdmissionInflight(bytes int64) { + if m == nil { + return + } + m.putAdmissionInflightBytes.WithLabelValues().Set(float64(max(int64(0), bytes))) +} + +func (m *S3Metrics) ObserveS3PutAdmissionRejection(stage, protocol string) { + if m == nil { + return + } + m.putAdmissionRejections.WithLabelValues( + normalizeS3PutAdmissionStage(stage), + normalizeS3PutAdmissionProtocol(protocol), + ).Inc() +} + +func (m *S3Metrics) ObserveS3PutAdmissionWait(stage, protocol string, duration time.Duration) { + if m == nil { + return + } + if duration < 0 { + duration = 0 + } + m.putAdmissionWait.WithLabelValues( + normalizeS3PutAdmissionStage(stage), + normalizeS3PutAdmissionProtocol(protocol), + ).Observe(duration.Seconds()) +} + +func normalizeS3PutAdmissionStage(stage string) string { + switch stage { + case s3PutAdmissionStagePrereserve, s3PutAdmissionStagePerBatch: + return stage + default: + return s3PutAdmissionStageUnknown + } +} + +func normalizeS3PutAdmissionProtocol(protocol string) string { + switch protocol { + case s3PutAdmissionProtocolFixed, s3PutAdmissionProtocolChunked: + return protocol + default: + return s3PutAdmissionProtocolUnknown + } +} diff --git a/monitoring/s3_test.go b/monitoring/s3_test.go new file mode 100644 index 000000000..03ffb3fc0 --- /dev/null +++ b/monitoring/s3_test.go @@ -0,0 +1,47 @@ +package monitoring + +import ( + "strings" + "testing" + "time" + + "github.com/prometheus/client_golang/prometheus" + "github.com/prometheus/client_golang/prometheus/testutil" + "github.com/stretchr/testify/require" +) + +func TestS3PutAdmissionMetricsObserve(t *testing.T) { + t.Parallel() + + reg := prometheus.NewRegistry() + metrics := newS3Metrics(reg) + + metrics.ObserveS3PutAdmissionInflight(1024) + metrics.ObserveS3PutAdmissionRejection(s3PutAdmissionStagePrereserve, s3PutAdmissionProtocolFixed) + metrics.ObserveS3PutAdmissionRejection(s3PutAdmissionStagePerBatch, s3PutAdmissionProtocolChunked) + metrics.ObserveS3PutAdmissionWait(s3PutAdmissionStagePerBatch, s3PutAdmissionProtocolChunked, 12*time.Millisecond) + + err := testutil.GatherAndCompare( + reg, + strings.NewReader(` +# HELP elastickv_s3_put_admission_inflight_bytes Current S3 PUT body bytes admitted by this node and not yet released after Raft dispatch. +# TYPE elastickv_s3_put_admission_inflight_bytes gauge +elastickv_s3_put_admission_inflight_bytes 1024 +# HELP elastickv_s3_put_admission_rejections_total Total S3 PUT admission rejections by admission stage and request protocol. +# TYPE elastickv_s3_put_admission_rejections_total counter +elastickv_s3_put_admission_rejections_total{protocol="chunked",stage="perbatch"} 1 +elastickv_s3_put_admission_rejections_total{protocol="fixed-length",stage="prereserve"} 1 +`), + "elastickv_s3_put_admission_inflight_bytes", + "elastickv_s3_put_admission_rejections_total", + ) + require.NoError(t, err) + require.Equal(t, 1, testutil.CollectAndCount(metrics.putAdmissionWait)) +} + +func TestRegistryReturnsS3PutAdmissionObserver(t *testing.T) { + t.Parallel() + + registry := NewRegistry("n1", "127.0.0.1:0") + require.NotNil(t, registry.S3PutAdmissionObserver()) +} From 08f8f3c4b26e02edf32ec8c74b8c1a8eb685d003 Mon Sep 17 00:00:00 2001 From: bootjp Date: Tue, 7 Jul 2026 21:31:20 +0900 Subject: [PATCH 2/8] Harden S3 put admission accounting --- adapter/s3_admission.go | 19 +++++++------------ adapter/s3_admission_test.go | 11 +++++++++++ monitoring/registry.go | 2 +- monitoring/s3.go | 7 +++---- monitoring/s3_test.go | 7 +++++++ 5 files changed, 29 insertions(+), 17 deletions(-) diff --git a/adapter/s3_admission.go b/adapter/s3_admission.go index ceaea5e1e..76984b6c7 100644 --- a/adapter/s3_admission.go +++ b/adapter/s3_admission.go @@ -154,21 +154,16 @@ func (a *s3PutAdmission) acquire(ctx context.Context, bytes int64) (func(), erro return func() {}, nil } units, ok := a.unitsFor(bytes) - if !ok { + if !ok || units > 1 { return nil, errS3PutAdmissionExhausted } - acquired := 0 - for acquired < units { - select { - case a.sem <- struct{}{}: - acquired++ - a.inflight.Add(s3ChunkSize) - case <-ctx.Done(): - a.releaseUnits(acquired) - return nil, errS3PutAdmissionExhausted - } + select { + case a.sem <- struct{}{}: + a.inflight.Add(s3ChunkSize) + return func() { a.releaseUnits(1) }, nil + case <-ctx.Done(): + return nil, errS3PutAdmissionExhausted } - return func() { a.releaseUnits(acquired) }, nil } func (a *s3PutAdmission) acquireWithTimeout(ctx context.Context, bytes int64) (func(), error) { diff --git a/adapter/s3_admission_test.go b/adapter/s3_admission_test.go index 3c0c394fc..c82df11d3 100644 --- a/adapter/s3_admission_test.go +++ b/adapter/s3_admission_test.go @@ -57,6 +57,17 @@ func TestS3PutAdmissionAcquireReleaseAndHeadroom(t *testing.T) { require.NoError(t, admission.peekHeadroom(2*s3ChunkSize)) } +func TestS3PutAdmissionAcquireRejectsMultiUnitRequest(t *testing.T) { + t.Parallel() + + admission := newS3PutAdmission(2*s3ChunkSize, time.Second) + release, err := admission.acquire(context.Background(), 2*s3ChunkSize) + require.ErrorIs(t, err, errS3PutAdmissionExhausted) + require.Nil(t, release) + require.Zero(t, admission.inflight.Load()) + require.NoError(t, admission.peekHeadroom(2*s3ChunkSize)) +} + func TestS3PutAdmissionAcquireTimeoutKeepsExistingLease(t *testing.T) { t.Parallel() diff --git a/monitoring/registry.go b/monitoring/registry.go index 6e6bd7151..ac934ecb7 100644 --- a/monitoring/registry.go +++ b/monitoring/registry.go @@ -224,7 +224,7 @@ func (r *Registry) SQSObserver() *SQSObserver { // by this registry. The adapter owns admission decisions and calls this small // interface directly from the hot path. func (r *Registry) S3PutAdmissionObserver() S3PutAdmissionObserver { - if r == nil { + if r == nil || r.s3 == nil { return nil } return r.s3 diff --git a/monitoring/s3.go b/monitoring/s3.go index 4ed64174c..363f24634 100644 --- a/monitoring/s3.go +++ b/monitoring/s3.go @@ -23,19 +23,18 @@ type S3PutAdmissionObserver interface { } type S3Metrics struct { - putAdmissionInflightBytes *prometheus.GaugeVec + putAdmissionInflightBytes prometheus.Gauge putAdmissionRejections *prometheus.CounterVec putAdmissionWait *prometheus.HistogramVec } func newS3Metrics(registerer prometheus.Registerer) *S3Metrics { m := &S3Metrics{ - putAdmissionInflightBytes: prometheus.NewGaugeVec( + putAdmissionInflightBytes: prometheus.NewGauge( prometheus.GaugeOpts{ Name: "elastickv_s3_put_admission_inflight_bytes", Help: "Current S3 PUT body bytes admitted by this node and not yet released after Raft dispatch.", }, - nil, ), putAdmissionRejections: prometheus.NewCounterVec( prometheus.CounterOpts{ @@ -65,7 +64,7 @@ func (m *S3Metrics) ObserveS3PutAdmissionInflight(bytes int64) { if m == nil { return } - m.putAdmissionInflightBytes.WithLabelValues().Set(float64(max(int64(0), bytes))) + m.putAdmissionInflightBytes.Set(float64(max(int64(0), bytes))) } func (m *S3Metrics) ObserveS3PutAdmissionRejection(stage, protocol string) { diff --git a/monitoring/s3_test.go b/monitoring/s3_test.go index 03ffb3fc0..8029a07b1 100644 --- a/monitoring/s3_test.go +++ b/monitoring/s3_test.go @@ -45,3 +45,10 @@ func TestRegistryReturnsS3PutAdmissionObserver(t *testing.T) { registry := NewRegistry("n1", "127.0.0.1:0") require.NotNil(t, registry.S3PutAdmissionObserver()) } + +func TestRegistryS3PutAdmissionObserverNilWhenMetricsMissing(t *testing.T) { + t.Parallel() + + registry := &Registry{} + require.Nil(t, registry.S3PutAdmissionObserver()) +} From a73046d2feaa10270f617aab9e4d073a9f07a6c8 Mon Sep 17 00:00:00 2001 From: bootjp Date: Wed, 8 Jul 2026 07:34:07 +0900 Subject: [PATCH 3/8] Fix S3 admission body handling --- adapter/s3.go | 57 +++++++++++++++--- adapter/s3_admission.go | 39 ++++++++---- adapter/s3_admission_test.go | 59 +++++++++++++++++++ adapter/s3_test.go | 27 +++++++++ ..._04_25_implemented_s3_admission_control.md | 2 +- 5 files changed, 164 insertions(+), 20 deletions(-) diff --git a/adapter/s3.go b/adapter/s3.go index 48bb4838c..67766a012 100644 --- a/adapter/s3.go +++ b/adapter/s3.go @@ -872,7 +872,7 @@ func (s *S3Server) putObject(w http.ResponseWriter, r *http.Request, bucket stri expectedPayloadSHA := normalizeS3PayloadHash(r.Header.Get("X-Amz-Content-Sha256")) validatePayloadSHA := expectedPayloadSHA != "" && !isS3PayloadMarker(expectedPayloadSHA) admissionProtocol := s3PutAdmissionProtocolForPayload(expectedPayloadSHA) - if !s.admitS3PutRequest(w, r, bucket, objectKey, s3MaxObjectSizeBytes) { + if !s.admitS3PutRequest(w, r, bucket, objectKey, s3MaxObjectSizeBytes, "object exceeds maximum allowed size") { return } streamBody, bodyErr := prepareStreamingPutBody(w, r, s3MaxObjectSizeBytes, expectedPayloadSHA, "object exceeds maximum allowed size") @@ -894,6 +894,9 @@ func (s *S3Server) putObject(w http.ResponseWriter, r *http.Request, bucket stri pendingAdmission = pendingAdmission[:0] } defer releasePendingAdmission() + nextReadBuffer := func() ([]byte, bool) { + return nextS3PutReadBuffer(buf, admissionProtocol, r.ContentLength, sizeBytes) + } uploadedManifest := func() *s3ObjectManifest { if len(part.ChunkSizes) == 0 { return &s3ObjectManifest{UploadID: uploadID} @@ -922,19 +925,23 @@ func (s *S3Server) putObject(w http.ResponseWriter, r *http.Request, bucket stri return nil } for { - releaseAdmission, admissionErr := s.acquireS3PutAdmission(r.Context(), s3ChunkSize, admissionProtocol) + readBuf, exactLength := nextReadBuffer() + if len(readBuf) == 0 { + break + } + releaseAdmission, admissionErr := s.acquireS3PutAdmission(r.Context(), int64(len(readBuf)), admissionProtocol) if admissionErr != nil { s.cleanupManifestBlobs(r.Context(), bucket, meta.Generation, objectKey, uploadedManifest()) writeS3AdmissionError(w, bucket, objectKey, s.s3PutAdmissionRetryAfter()) return } - n, readErr := readS3PutChunk(r.Body, buf) + n, readErr := readS3PutChunk(r.Body, readBuf, !exactLength) if n == 0 { releaseAdmission() } if n > 0 { pendingAdmission = append(pendingAdmission, releaseAdmission) - chunk := append([]byte(nil), buf[:n]...) + chunk := append([]byte(nil), readBuf[:n]...) if _, err := hasher.Write(chunk); err != nil { writeS3InternalError(w, err) return @@ -1401,7 +1408,7 @@ func (s *S3Server) uploadPart(w http.ResponseWriter, r *http.Request, bucket str partPayloadSHA := normalizeS3PayloadHash(r.Header.Get("X-Amz-Content-Sha256")) admissionProtocol := s3PutAdmissionProtocolForPayload(partPayloadSHA) - if !s.admitS3PutRequest(w, r, bucket, objectKey, s3MaxPartSizeBytes) { + if !s.admitS3PutRequest(w, r, bucket, objectKey, s3MaxPartSizeBytes, "part exceeds maximum allowed size") { return } partStreamBody, bodyErr := prepareStreamingPutBody(w, r, s3MaxPartSizeBytes, partPayloadSHA, "part exceeds maximum allowed size") @@ -1443,6 +1450,9 @@ func (s *S3Server) uploadPart(w http.ResponseWriter, r *http.Request, bucket str pendingAdmission = pendingAdmission[:0] } defer releasePendingAdmission() + nextReadBuffer := func() ([]byte, bool) { + return nextS3PutReadBuffer(buf, admissionProtocol, r.ContentLength, sizeBytes) + } partCommitted := false defer func() { if !partCommitted && chunkNo > 0 { @@ -1464,12 +1474,16 @@ func (s *S3Server) uploadPart(w http.ResponseWriter, r *http.Request, bucket str } for { - releaseAdmission, admissionErr := s.acquireS3PutAdmission(r.Context(), s3ChunkSize, admissionProtocol) + readBuf, exactLength := nextReadBuffer() + if len(readBuf) == 0 { + break + } + releaseAdmission, admissionErr := s.acquireS3PutAdmission(r.Context(), int64(len(readBuf)), admissionProtocol) if admissionErr != nil { writeS3AdmissionError(w, bucket, objectKey, s.s3PutAdmissionRetryAfter()) return } - n, readErr := readS3PutChunk(r.Body, buf) + n, readErr := readS3PutChunk(r.Body, readBuf, !exactLength) if n == 0 { releaseAdmission() } @@ -1488,7 +1502,7 @@ func (s *S3Server) uploadPart(w http.ResponseWriter, r *http.Request, bucket str continue } pendingAdmission = append(pendingAdmission, releaseAdmission) - chunk := append([]byte(nil), buf[:n]...) + chunk := append([]byte(nil), readBuf[:n]...) if _, err := hasher.Write(chunk); err != nil { writeS3InternalError(w, err) return @@ -2515,11 +2529,33 @@ func (s *S3Server) isVerifiedS3Leader(ctx context.Context) bool { return s.coordinator.VerifyLeader(ctx) == nil } -func readS3PutChunk(r io.Reader, buf []byte) (int, error) { +var errS3PutIncompleteBody = errors.New("request body shorter than Content-Length") + +func nextS3PutReadBuffer(buf []byte, protocol string, contentLength int64, readBytes int64) ([]byte, bool) { + if protocol != s3PutAdmissionProtocolFixed || contentLength < 0 { + return buf, false + } + remaining := contentLength - readBytes + if remaining <= 0 { + return nil, true + } + if remaining < int64(len(buf)) { + return buf[:int(remaining)], true + } + return buf, true +} + +func readS3PutChunk(r io.Reader, buf []byte, allowFinalPartial bool) (int, error) { n, err := io.ReadFull(r, buf) if errors.Is(err, io.ErrUnexpectedEOF) { + if !allowFinalPartial { + return n, errS3PutIncompleteBody + } return n, io.EOF } + if errors.Is(err, io.EOF) && !allowFinalPartial && len(buf) > 0 { + return n, errS3PutIncompleteBody + } if err != nil { return n, errors.WithStack(err) } @@ -2665,6 +2701,9 @@ func classifyS3BodyReadErr(err error, tooLargeMessage string) (*s3PutBodyError, if errors.As(err, &chunkedErr) { return &s3PutBodyError{Status: http.StatusBadRequest, Code: "InvalidRequest", Message: chunkedErr.Error()}, true } + if errors.Is(err, errS3PutIncompleteBody) { + return &s3PutBodyError{Status: http.StatusBadRequest, Code: "IncompleteBody", Message: errS3PutIncompleteBody.Error()}, true + } return nil, false } diff --git a/adapter/s3_admission.go b/adapter/s3_admission.go index 76984b6c7..cfc90996e 100644 --- a/adapter/s3_admission.go +++ b/adapter/s3_admission.go @@ -16,6 +16,7 @@ import ( const ( s3PutAdmissionDefaultMaxInflightBytes = 256 << 20 s3PutAdmissionDefaultTimeout = 30 * time.Second + s3PutAdmissionDisableMaxInflightBytes = int64(1<<63 - 1) s3PutAdmissionMaxInflightEnv = "ELASTICKV_S3_PUT_ADMISSION_MAX_INFLIGHT_BYTES" s3PutAdmissionTimeoutEnv = "ELASTICKV_S3_DISPATCH_ADMISSION_TIMEOUT" @@ -27,7 +28,10 @@ const ( s3PutAdmissionProtocolChunked = "chunked" ) -var errS3PutAdmissionExhausted = errors.New("s3 put admission budget exhausted") +var ( + errS3PutAdmissionExhausted = errors.New("s3 put admission budget exhausted") + errS3PutAdmissionEntityTooLarge = errors.New("s3 put body exceeds maximum allowed size") +) type S3PutAdmissionObserver interface { ObserveS3PutAdmissionInflight(bytes int64) @@ -55,20 +59,26 @@ func newS3PutAdmission(maxBytes int64, timeout time.Duration) *s3PutAdmission { } func newS3PutAdmissionWithChunked(maxBytes int64, timeout time.Duration, chunked bool) *s3PutAdmission { - if maxBytes <= 0 { + if maxBytes <= 0 || maxBytes == s3PutAdmissionDisableMaxInflightBytes { return nil } - units := int((maxBytes + s3ChunkSize - 1) / s3ChunkSize) - if units < 1 { - units = 1 + units64 := 1 + (maxBytes-1)/s3ChunkSize + maxInt := int64(int(^uint(0) >> 1)) + if units64 > maxInt { + return nil } + units := int(units64) if timeout <= 0 { timeout = s3PutAdmissionDefaultTimeout } + roundedMaxBytes := units64 * s3ChunkSize + if roundedMaxBytes < 0 || roundedMaxBytes < maxBytes { + roundedMaxBytes = s3PutAdmissionDisableMaxInflightBytes + } return &s3PutAdmission{ sem: make(chan struct{}, units), timeout: timeout, - maxBytes: int64(units) * s3ChunkSize, + maxBytes: roundedMaxBytes, chunked: chunked, } } @@ -191,23 +201,29 @@ func (a *s3PutAdmission) unitsFor(bytes int64) (int, bool) { if a == nil || bytes <= 0 { return 0, true } - units := int((bytes + s3ChunkSize - 1) / s3ChunkSize) - if units < 1 { - units = 1 + units64 := 1 + (bytes-1)/s3ChunkSize + maxInt := int64(int(^uint(0) >> 1)) + if units64 > maxInt { + return 0, false } + units := int(units64) if units > cap(a.sem) { return 0, false } return units, true } -func (s *S3Server) admitS3PutRequest(w http.ResponseWriter, r *http.Request, bucket, objectKey string, maxDecoded int64) bool { +func (s *S3Server) admitS3PutRequest(w http.ResponseWriter, r *http.Request, bucket, objectKey string, maxDecoded int64, tooLargeMessage string) bool { if s == nil || s.putAdmission == nil { return true } bytes, protocol, err := s.s3PutAdmissionProbeBytes(r, maxDecoded) if err != nil { s.observeS3PutAdmissionRejection(s3PutAdmissionStagePrereserve, protocol) + if errors.Is(err, errS3PutAdmissionEntityTooLarge) { + writeS3Error(w, http.StatusRequestEntityTooLarge, "EntityTooLarge", tooLargeMessage, bucket, objectKey) + return false + } writeS3Error(w, http.StatusLengthRequired, "MissingContentLength", err.Error(), bucket, objectKey) return false } @@ -231,6 +247,9 @@ func (s *S3Server) s3PutAdmissionProbeBytes(r *http.Request, maxDecoded int64) ( if r.ContentLength < 0 { return 0, protocol, errors.New("Content-Length is required for non-streaming S3 PUT admission") } + if maxDecoded > 0 && r.ContentLength > maxDecoded { + return 0, protocol, errS3PutAdmissionEntityTooLarge + } return r.ContentLength, protocol, nil } diff --git a/adapter/s3_admission_test.go b/adapter/s3_admission_test.go index c82df11d3..73af192fd 100644 --- a/adapter/s3_admission_test.go +++ b/adapter/s3_admission_test.go @@ -98,6 +98,12 @@ func TestS3PutAdmissionFromEnvCanDisableChunkedIncremental(t *testing.T) { require.False(t, admission.chunked) } +func TestS3PutAdmissionFromEnvCanDisableCapWithMaxInt(t *testing.T) { + t.Setenv(s3PutAdmissionMaxInflightEnv, "9223372036854775807") + + require.Nil(t, newS3PutAdmissionFromEnv()) +} + func TestS3PutAdmissionProbeUsesBootstrapForChunkedDecodedLength(t *testing.T) { t.Parallel() @@ -136,6 +142,24 @@ func TestS3Server_PutObjectAdmissionRejectsOversizedRequest(t *testing.T) { require.Equal(t, http.StatusNotFound, rec.Code) } +func TestS3Server_PutObjectAdmissionRejectsFixedLengthOverS3LimitAsEntityTooLarge(t *testing.T) { + t.Parallel() + + server, observer := newS3AdmissionTestServer(t, time.Second) + createS3AdmissionTestBucket(t, server, "admit-s3-limit") + + rec := httptest.NewRecorder() + req := newS3TestRequest(http.MethodPut, "/admit-s3-limit/too-large.bin", http.NoBody) + req.ContentLength = s3MaxObjectSizeBytes + 1 + server.handle(rec, req) + + require.Equal(t, http.StatusRequestEntityTooLarge, rec.Code, rec.Body.String()) + require.Contains(t, rec.Body.String(), "EntityTooLarge") + require.Contains(t, rec.Body.String(), "object exceeds maximum allowed size") + require.Equal(t, 1, observer.rejections[s3PutAdmissionStagePrereserve+"|"+s3PutAdmissionProtocolFixed]) + require.Zero(t, observer.lastInflight) +} + func TestS3Server_PutObjectAdmissionRequiresContentLengthForPlainPut(t *testing.T) { t.Parallel() @@ -172,6 +196,41 @@ func TestS3Server_PutObjectAdmissionAllowsSmallRequestAndReleasesBudget(t *testi require.Equal(t, "ok", rec.Body.String()) } +func TestS3Server_PutObjectAdmissionExactFixedLengthChunkDoesNotChargeEOF(t *testing.T) { + t.Parallel() + + server, observer := newS3AdmissionTestServer(t, time.Second) + createS3AdmissionTestBucket(t, server, "admit-exact") + + payload := bytes.Repeat([]byte("x"), s3ChunkSize) + rec := httptest.NewRecorder() + server.handle(rec, newS3TestRequest(http.MethodPut, "/admit-exact/exact.bin", bytes.NewReader(payload))) + + require.Equal(t, http.StatusOK, rec.Code, rec.Body.String()) + require.Equal(t, 1, observer.waits[s3PutAdmissionStagePerBatch+"|"+s3PutAdmissionProtocolFixed]) + require.Zero(t, observer.lastInflight) +} + +func TestS3Server_PutObjectRejectsTruncatedFixedLengthBody(t *testing.T) { + t.Parallel() + + server, observer := newS3AdmissionTestServer(t, time.Second) + createS3AdmissionTestBucket(t, server, "admit-truncated") + + rec := httptest.NewRecorder() + req := newS3TestRequest(http.MethodPut, "/admit-truncated/short.bin", strings.NewReader("short")) + req.ContentLength = int64(len("short") + 1) + server.handle(rec, req) + + require.Equal(t, http.StatusBadRequest, rec.Code, rec.Body.String()) + require.Contains(t, rec.Body.String(), "IncompleteBody") + require.Zero(t, observer.lastInflight) + + rec = httptest.NewRecorder() + server.handle(rec, newS3TestRequest(http.MethodGet, "/admit-truncated/short.bin", nil)) + require.Equal(t, http.StatusNotFound, rec.Code) +} + func TestS3Server_PutObjectAdmissionChunkedMidstreamTimeoutReleasesBudget(t *testing.T) { t.Parallel() diff --git a/adapter/s3_test.go b/adapter/s3_test.go index cfc595a9a..7f03c2114 100644 --- a/adapter/s3_test.go +++ b/adapter/s3_test.go @@ -993,6 +993,33 @@ func TestS3Server_MultipartUploadHappyPath(t *testing.T) { require.Contains(t, rec.Body.String(), "large-file.bin") } +func TestS3Server_UploadPartRejectsTruncatedFixedLengthBody(t *testing.T) { + t.Parallel() + + st := store.NewMVCCStore() + server := NewS3Server(nil, "", st, newLocalAdapterCoordinator(st), nil) + + rec := httptest.NewRecorder() + server.handle(rec, newS3TestRequest(http.MethodPut, "/bucket-mp-truncated", nil)) + require.Equal(t, http.StatusOK, rec.Code) + + rec = httptest.NewRecorder() + server.handle(rec, newS3TestRequest(http.MethodPost, "/bucket-mp-truncated/obj?uploads=", nil)) + require.Equal(t, http.StatusOK, rec.Code) + var initResult s3InitiateMultipartUploadResult + require.NoError(t, xml.Unmarshal(rec.Body.Bytes(), &initResult)) + + rec = httptest.NewRecorder() + req := newS3TestRequest(http.MethodPut, + fmt.Sprintf("/bucket-mp-truncated/obj?uploadId=%s&partNumber=1", initResult.UploadId), + strings.NewReader("short")) + req.ContentLength = int64(len("short") + 1) + server.handle(rec, req) + + require.Equal(t, http.StatusBadRequest, rec.Code, rec.Body.String()) + require.Contains(t, rec.Body.String(), "IncompleteBody") +} + func TestS3Server_AbortMultipartUpload(t *testing.T) { t.Parallel() diff --git a/docs/design/2026_04_25_implemented_s3_admission_control.md b/docs/design/2026_04_25_implemented_s3_admission_control.md index 2e8402287..434ab9890 100644 --- a/docs/design/2026_04_25_implemented_s3_admission_control.md +++ b/docs/design/2026_04_25_implemented_s3_admission_control.md @@ -437,7 +437,7 @@ generous enough that even single-node M1 traffic falls below the threshold under typical load, so the rollout signature is "503 SlowDown rate goes from 0 to negligible" rather than a step function. Operators can pin -`ELASTICKV_S3_PUT_ADMISSION_MAX_INFLIGHT_BYTES=$((1<<63))` to +`ELASTICKV_S3_PUT_ADMISSION_MAX_INFLIGHT_BYTES=9223372036854775807` to disable the cap on M1 nodes during the burn-in window if desired. The aws-chunked progress-callback path is the only behaviour From 0a6ff2b822483ba8441bae7fb8a551efad2f4357 Mon Sep 17 00:00:00 2001 From: bootjp Date: Wed, 8 Jul 2026 10:05:25 +0900 Subject: [PATCH 4/8] s3: stream PUT admission batches --- adapter/s3.go | 59 ++++++++++++++++++--------- adapter/s3_admission.go | 13 ++++++ adapter/s3_admission_test.go | 78 +++++++++++++++++++++++++++++------- 3 files changed, 116 insertions(+), 34 deletions(-) diff --git a/adapter/s3.go b/adapter/s3.go index 67766a012..0aaedf22f 100644 --- a/adapter/s3.go +++ b/adapter/s3.go @@ -925,21 +925,25 @@ func (s *S3Server) putObject(w http.ResponseWriter, r *http.Request, bucket stri return nil } for { + if s.shouldFlushS3PutBatchBeforeRead(admissionProtocol, len(pendingAdmission)) { + if err := flushBatch(); err != nil { + s.cleanupManifestBlobs(r.Context(), bucket, meta.Generation, objectKey, uploadedManifest()) + writeS3InternalError(w, err) + return + } + } readBuf, exactLength := nextReadBuffer() if len(readBuf) == 0 { break } - releaseAdmission, admissionErr := s.acquireS3PutAdmission(r.Context(), int64(len(readBuf)), admissionProtocol) - if admissionErr != nil { - s.cleanupManifestBlobs(r.Context(), bucket, meta.Generation, objectKey, uploadedManifest()) - writeS3AdmissionError(w, bucket, objectKey, s.s3PutAdmissionRetryAfter()) - return - } - n, readErr := readS3PutChunk(r.Body, readBuf, !exactLength) - if n == 0 { - releaseAdmission() - } + n, readErr := s.readS3PutChunkWithAdmissionDeadline(w, r.Body, readBuf, !exactLength) if n > 0 { + releaseAdmission, admissionErr := s.acquireS3PutAdmission(r.Context(), int64(n), admissionProtocol) + if admissionErr != nil { + s.cleanupManifestBlobs(r.Context(), bucket, meta.Generation, objectKey, uploadedManifest()) + writeS3AdmissionError(w, bucket, objectKey, s.s3PutAdmissionRetryAfter()) + return + } pendingAdmission = append(pendingAdmission, releaseAdmission) chunk := append([]byte(nil), readBuf[:n]...) if _, err := hasher.Write(chunk); err != nil { @@ -1474,19 +1478,17 @@ func (s *S3Server) uploadPart(w http.ResponseWriter, r *http.Request, bucket str } for { + if s.shouldFlushS3PutBatchBeforeRead(admissionProtocol, len(pendingAdmission)) { + if err := flushBatch(); err != nil { + writeS3InternalError(w, err) + return + } + } readBuf, exactLength := nextReadBuffer() if len(readBuf) == 0 { break } - releaseAdmission, admissionErr := s.acquireS3PutAdmission(r.Context(), int64(len(readBuf)), admissionProtocol) - if admissionErr != nil { - writeS3AdmissionError(w, bucket, objectKey, s.s3PutAdmissionRetryAfter()) - return - } - n, readErr := readS3PutChunk(r.Body, readBuf, !exactLength) - if n == 0 { - releaseAdmission() - } + n, readErr := s.readS3PutChunkWithAdmissionDeadline(w, r.Body, readBuf, !exactLength) if n == 0 { if errors.Is(readErr, io.EOF) { break @@ -1501,6 +1503,11 @@ func (s *S3Server) uploadPart(w http.ResponseWriter, r *http.Request, bucket str } continue } + releaseAdmission, admissionErr := s.acquireS3PutAdmission(r.Context(), int64(n), admissionProtocol) + if admissionErr != nil { + writeS3AdmissionError(w, bucket, objectKey, s.s3PutAdmissionRetryAfter()) + return + } pendingAdmission = append(pendingAdmission, releaseAdmission) chunk := append([]byte(nil), readBuf[:n]...) if _, err := hasher.Write(chunk); err != nil { @@ -2562,6 +2569,20 @@ func readS3PutChunk(r io.Reader, buf []byte, allowFinalPartial bool) (int, error return n, nil } +func (s *S3Server) readS3PutChunkWithAdmissionDeadline(w http.ResponseWriter, r io.Reader, buf []byte, allowFinalPartial bool) (int, error) { + if s == nil || s.putAdmission == nil || s.putAdmission.timeout <= 0 { + return readS3PutChunk(r, buf, allowFinalPartial) + } + controller := http.NewResponseController(w) + if err := controller.SetReadDeadline(time.Now().Add(s.putAdmission.timeout)); err != nil { + return readS3PutChunk(r, buf, allowFinalPartial) + } + defer func() { + _ = controller.SetReadDeadline(time.Time{}) + }() + return readS3PutChunk(r, buf, allowFinalPartial) +} + // prepareStreamingPutBody wraps r.Body for aws-chunked framed uploads. When // the request is a plain (non-streaming) PUT the body is only wrapped with // MaxBytesReader; the streaming-body context returned is the zero value and diff --git a/adapter/s3_admission.go b/adapter/s3_admission.go index cfc90996e..053acd8a1 100644 --- a/adapter/s3_admission.go +++ b/adapter/s3_admission.go @@ -250,6 +250,9 @@ func (s *S3Server) s3PutAdmissionProbeBytes(r *http.Request, maxDecoded int64) ( if maxDecoded > 0 && r.ContentLength > maxDecoded { return 0, protocol, errS3PutAdmissionEntityTooLarge } + if r.ContentLength > s3ChunkSize { + return s3ChunkSize, protocol, nil + } return r.ContentLength, protocol, nil } @@ -281,6 +284,16 @@ func (s *S3Server) acquireS3PutAdmission(ctx context.Context, bytes int64, proto }, nil } +func (s *S3Server) shouldFlushS3PutBatchBeforeRead(protocol string, pendingAdmission int) bool { + if pendingAdmission == 0 || s == nil || s.putAdmission == nil { + return false + } + if protocol == s3PutAdmissionProtocolChunked { + return true + } + return pendingAdmission >= cap(s.putAdmission.sem) +} + func (s *S3Server) s3PutAdmissionRetryAfter() time.Duration { if s == nil || s.putAdmission == nil { return s3PutAdmissionDefaultTimeout diff --git a/adapter/s3_admission_test.go b/adapter/s3_admission_test.go index 73af192fd..3435edaac 100644 --- a/adapter/s3_admission_test.go +++ b/adapter/s3_admission_test.go @@ -3,6 +3,8 @@ package adapter import ( "bytes" "context" + "encoding/xml" + "fmt" "io" "net/http" "net/http/httptest" @@ -118,7 +120,20 @@ func TestS3PutAdmissionProbeUsesBootstrapForChunkedDecodedLength(t *testing.T) { require.Equal(t, s3PutAdmissionProtocolChunked, protocol) } -func TestS3Server_PutObjectAdmissionRejectsOversizedRequest(t *testing.T) { +func TestS3PutAdmissionProbeUsesChunkHeadroomForFixedLength(t *testing.T) { + t.Parallel() + + server := &S3Server{putAdmission: newS3PutAdmission(s3ChunkSize, time.Second)} + req := newS3TestRequest(http.MethodPut, "/bucket/key", http.NoBody) + req.ContentLength = 2*s3ChunkSize + 1 + + bytes, protocol, err := server.s3PutAdmissionProbeBytes(req, s3MaxObjectSizeBytes) + require.NoError(t, err) + require.EqualValues(t, s3ChunkSize, bytes) + require.Equal(t, s3PutAdmissionProtocolFixed, protocol) +} + +func TestS3Server_PutObjectAdmissionAllowsFixedLengthLargerThanBudget(t *testing.T) { t.Parallel() server, observer := newS3AdmissionTestServer(t, 2*time.Second) @@ -126,20 +141,19 @@ func TestS3Server_PutObjectAdmissionRejectsOversizedRequest(t *testing.T) { payload := strings.Repeat("x", s3ChunkSize+1) rec := httptest.NewRecorder() - req := newS3TestRequest(http.MethodPut, "/admit-big/too-large.bin", strings.NewReader(payload)) + req := newS3TestRequest(http.MethodPut, "/admit-big/larger-than-budget.bin", strings.NewReader(payload)) server.handle(rec, req) - require.Equal(t, http.StatusServiceUnavailable, rec.Code, rec.Body.String()) - require.Equal(t, "2", rec.Header().Get("Retry-After")) - require.Equal(t, "close", rec.Header().Get("Connection")) - require.Contains(t, rec.Body.String(), "SlowDown") - require.Contains(t, rec.Body.String(), "Reduce your request rate") - require.Equal(t, 1, observer.rejections[s3PutAdmissionStagePrereserve+"|"+s3PutAdmissionProtocolFixed]) + require.Equal(t, http.StatusOK, rec.Code, rec.Body.String()) + require.Greater(t, observer.waits[s3PutAdmissionStagePerBatch+"|"+s3PutAdmissionProtocolFixed], 1) + require.Zero(t, observer.rejections[s3PutAdmissionStagePrereserve+"|"+s3PutAdmissionProtocolFixed]) + require.Zero(t, observer.rejections[s3PutAdmissionStagePerBatch+"|"+s3PutAdmissionProtocolFixed]) require.Zero(t, observer.lastInflight) rec = httptest.NewRecorder() - server.handle(rec, newS3TestRequest(http.MethodGet, "/admit-big/too-large.bin", nil)) - require.Equal(t, http.StatusNotFound, rec.Code) + server.handle(rec, newS3TestRequest(http.MethodGet, "/admit-big/larger-than-budget.bin", nil)) + require.Equal(t, http.StatusOK, rec.Code) + require.Equal(t, payload, rec.Body.String()) } func TestS3Server_PutObjectAdmissionRejectsFixedLengthOverS3LimitAsEntityTooLarge(t *testing.T) { @@ -231,7 +245,7 @@ func TestS3Server_PutObjectRejectsTruncatedFixedLengthBody(t *testing.T) { require.Equal(t, http.StatusNotFound, rec.Code) } -func TestS3Server_PutObjectAdmissionChunkedMidstreamTimeoutReleasesBudget(t *testing.T) { +func TestS3Server_PutObjectAdmissionChunkedStreamsPastSingleChunkBudget(t *testing.T) { t.Parallel() server, observer := newS3AdmissionTestServer(t, 5*time.Millisecond) @@ -245,15 +259,49 @@ func TestS3Server_PutObjectAdmissionChunkedMidstreamTimeoutReleasesBudget(t *tes req.Header.Set("X-Amz-Content-Sha256", s3StreamingUnsignedPayloadTrailer) server.handle(rec, req) - require.Equal(t, http.StatusServiceUnavailable, rec.Code, rec.Body.String()) - require.Contains(t, rec.Body.String(), "SlowDown") - require.Equal(t, 1, observer.rejections[s3PutAdmissionStagePerBatch+"|"+s3PutAdmissionProtocolChunked]) + require.Equal(t, http.StatusOK, rec.Code, rec.Body.String()) + require.Greater(t, observer.waits[s3PutAdmissionStagePerBatch+"|"+s3PutAdmissionProtocolChunked], 1) + require.Zero(t, observer.rejections[s3PutAdmissionStagePerBatch+"|"+s3PutAdmissionProtocolChunked]) require.Zero(t, observer.lastInflight) require.Zero(t, server.putAdmission.inflight.Load()) rec = httptest.NewRecorder() server.handle(rec, newS3TestRequest(http.MethodGet, "/admit-chunked/midstream.bin", nil)) - require.Equal(t, http.StatusNotFound, rec.Code) + require.Equal(t, http.StatusOK, rec.Code) + require.Equal(t, string(payload), rec.Body.String()) +} + +func TestS3Server_UploadPartAdmissionAllowsFixedLengthLargerThanBudget(t *testing.T) { + t.Parallel() + + server, observer := newS3AdmissionTestServer(t, 2*time.Second) + createS3AdmissionTestBucket(t, server, "admit-part") + + rec := httptest.NewRecorder() + req := newS3TestRequest(http.MethodPost, "/admit-part/object.bin?uploads=", nil) + server.handle(rec, req) + require.Equal(t, http.StatusOK, rec.Code, rec.Body.String()) + + var initResult s3InitiateMultipartUploadResult + require.NoError(t, xml.Unmarshal(rec.Body.Bytes(), &initResult)) + require.NotEmpty(t, initResult.UploadId) + + payload := strings.Repeat("p", s3ChunkSize+1) + rec = httptest.NewRecorder() + req = newS3TestRequest( + http.MethodPut, + fmt.Sprintf("/admit-part/object.bin?uploadId=%s&partNumber=1", initResult.UploadId), + strings.NewReader(payload), + ) + server.handle(rec, req) + + require.Equal(t, http.StatusOK, rec.Code, rec.Body.String()) + require.Equal(t, quoteS3ETag(md5Hex(payload)), rec.Header().Get("ETag")) + require.Greater(t, observer.waits[s3PutAdmissionStagePerBatch+"|"+s3PutAdmissionProtocolFixed], 1) + require.Zero(t, observer.rejections[s3PutAdmissionStagePrereserve+"|"+s3PutAdmissionProtocolFixed]) + require.Zero(t, observer.rejections[s3PutAdmissionStagePerBatch+"|"+s3PutAdmissionProtocolFixed]) + require.Zero(t, observer.lastInflight) + require.Zero(t, server.putAdmission.inflight.Load()) } func newS3AdmissionTestServer(t *testing.T, timeout time.Duration) (*S3Server, *recordingS3PutAdmissionObserver) { From de31938b4cbacd0091c601d1edaa26ee4962a274 Mon Sep 17 00:00:00 2001 From: bootjp Date: Wed, 8 Jul 2026 10:44:14 +0900 Subject: [PATCH 5/8] s3: preserve streaming body errors --- adapter/s3.go | 46 ++++++++++------------ adapter/s3_admission.go | 37 ++++++++++++++++-- adapter/s3_admission_test.go | 37 ++++++++++++++++++ adapter/s3_test.go | 75 ++++++++++++++++++++++++++++++++++++ 4 files changed, 166 insertions(+), 29 deletions(-) diff --git a/adapter/s3.go b/adapter/s3.go index 0aaedf22f..ebbe4f6b9 100644 --- a/adapter/s3.go +++ b/adapter/s3.go @@ -936,7 +936,7 @@ func (s *S3Server) putObject(w http.ResponseWriter, r *http.Request, bucket stri if len(readBuf) == 0 { break } - n, readErr := s.readS3PutChunkWithAdmissionDeadline(w, r.Body, readBuf, !exactLength) + n, readErr := readS3PutChunk(r.Body, readBuf, !exactLength) if n > 0 { releaseAdmission, admissionErr := s.acquireS3PutAdmission(r.Context(), int64(n), admissionProtocol) if admissionErr != nil { @@ -1488,7 +1488,7 @@ func (s *S3Server) uploadPart(w http.ResponseWriter, r *http.Request, bucket str if len(readBuf) == 0 { break } - n, readErr := s.readS3PutChunkWithAdmissionDeadline(w, r.Body, readBuf, !exactLength) + n, readErr := readS3PutChunk(r.Body, readBuf, !exactLength) if n == 0 { if errors.Is(readErr, io.EOF) { break @@ -2553,34 +2553,30 @@ func nextS3PutReadBuffer(buf []byte, protocol string, contentLength int64, readB } func readS3PutChunk(r io.Reader, buf []byte, allowFinalPartial bool) (int, error) { - n, err := io.ReadFull(r, buf) - if errors.Is(err, io.ErrUnexpectedEOF) { - if !allowFinalPartial { - return n, errS3PutIncompleteBody + total := 0 + for total < len(buf) { + n, err := r.Read(buf[total:]) + if n > 0 { + total += n + } + if errors.Is(err, io.EOF) { + return s3PutChunkEOFResult(total, len(buf), allowFinalPartial) + } + if err != nil { + return total, errors.WithStack(err) + } + if n == 0 { + return total, io.ErrNoProgress } - return n, io.EOF - } - if errors.Is(err, io.EOF) && !allowFinalPartial && len(buf) > 0 { - return n, errS3PutIncompleteBody - } - if err != nil { - return n, errors.WithStack(err) } - return n, nil + return total, nil } -func (s *S3Server) readS3PutChunkWithAdmissionDeadline(w http.ResponseWriter, r io.Reader, buf []byte, allowFinalPartial bool) (int, error) { - if s == nil || s.putAdmission == nil || s.putAdmission.timeout <= 0 { - return readS3PutChunk(r, buf, allowFinalPartial) +func s3PutChunkEOFResult(total, bufLen int, allowFinalPartial bool) (int, error) { + if !allowFinalPartial && (total > 0 || bufLen > 0) { + return total, errS3PutIncompleteBody } - controller := http.NewResponseController(w) - if err := controller.SetReadDeadline(time.Now().Add(s.putAdmission.timeout)); err != nil { - return readS3PutChunk(r, buf, allowFinalPartial) - } - defer func() { - _ = controller.SetReadDeadline(time.Time{}) - }() - return readS3PutChunk(r, buf, allowFinalPartial) + return total, io.EOF } // prepareStreamingPutBody wraps r.Body for aws-chunked framed uploads. When diff --git a/adapter/s3_admission.go b/adapter/s3_admission.go index 053acd8a1..21933c835 100644 --- a/adapter/s3_admission.go +++ b/adapter/s3_admission.go @@ -31,6 +31,7 @@ const ( var ( errS3PutAdmissionExhausted = errors.New("s3 put admission budget exhausted") errS3PutAdmissionEntityTooLarge = errors.New("s3 put body exceeds maximum allowed size") + errS3PutAdmissionInvalidDecoded = errors.New("invalid X-Amz-Decoded-Content-Length") ) type S3PutAdmissionObserver interface { @@ -224,6 +225,10 @@ func (s *S3Server) admitS3PutRequest(w http.ResponseWriter, r *http.Request, buc writeS3Error(w, http.StatusRequestEntityTooLarge, "EntityTooLarge", tooLargeMessage, bucket, objectKey) return false } + if errors.Is(err, errS3PutAdmissionInvalidDecoded) { + writeS3Error(w, http.StatusBadRequest, "InvalidRequest", err.Error(), bucket, objectKey) + return false + } writeS3Error(w, http.StatusLengthRequired, "MissingContentLength", err.Error(), bucket, objectKey) return false } @@ -239,10 +244,8 @@ func (s *S3Server) s3PutAdmissionProbeBytes(r *http.Request, maxDecoded int64) ( payloadSHA := normalizeS3PayloadHash(r.Header.Get("X-Amz-Content-Sha256")) protocol := s3PutAdmissionProtocolForPayload(payloadSHA) if protocol == s3PutAdmissionProtocolChunked { - if maxDecoded > 0 && maxDecoded < s3ChunkSize { - return maxDecoded, protocol, nil - } - return s3ChunkSize, protocol, nil + bytes, err := s3ChunkedPutAdmissionProbeBytes(r, maxDecoded) + return bytes, protocol, err } if r.ContentLength < 0 { return 0, protocol, errors.New("Content-Length is required for non-streaming S3 PUT admission") @@ -256,6 +259,32 @@ func (s *S3Server) s3PutAdmissionProbeBytes(r *http.Request, maxDecoded int64) ( return r.ContentLength, protocol, nil } +func s3ChunkedPutAdmissionProbeBytes(r *http.Request, maxDecoded int64) (int64, error) { + declared, ok, err := s3DecodedContentLength(r) + if err != nil { + return 0, err + } + if ok && maxDecoded > 0 && declared > maxDecoded { + return 0, errS3PutAdmissionEntityTooLarge + } + if maxDecoded > 0 && maxDecoded < s3ChunkSize { + return maxDecoded, nil + } + return s3ChunkSize, nil +} + +func s3DecodedContentLength(r *http.Request) (int64, bool, error) { + declaredRaw := strings.TrimSpace(r.Header.Get("X-Amz-Decoded-Content-Length")) + if declaredRaw == "" { + return 0, false, nil + } + declared, err := strconv.ParseInt(declaredRaw, 10, 64) + if err != nil || declared < 0 { + return 0, true, errS3PutAdmissionInvalidDecoded + } + return declared, true, nil +} + func s3PutAdmissionProtocolForPayload(payloadSHA string) string { if isS3StreamingPayloadMarker(payloadSHA) { return s3PutAdmissionProtocolChunked diff --git a/adapter/s3_admission_test.go b/adapter/s3_admission_test.go index 3435edaac..5d85a122e 100644 --- a/adapter/s3_admission_test.go +++ b/adapter/s3_admission_test.go @@ -120,6 +120,20 @@ func TestS3PutAdmissionProbeUsesBootstrapForChunkedDecodedLength(t *testing.T) { require.Equal(t, s3PutAdmissionProtocolChunked, protocol) } +func TestS3PutAdmissionProbeRejectsChunkedDecodedLengthOverS3Limit(t *testing.T) { + t.Parallel() + + server := &S3Server{putAdmission: newS3PutAdmission(s3ChunkSize, time.Second)} + req := newS3TestRequest(http.MethodPut, "/bucket/key", strings.NewReader("ignored")) + req.Header.Set("X-Amz-Content-Sha256", s3StreamingUnsignedPayloadTrailer) + req.Header.Set("X-Amz-Decoded-Content-Length", strconv.FormatInt(s3MaxObjectSizeBytes+1, 10)) + + bytes, protocol, err := server.s3PutAdmissionProbeBytes(req, s3MaxObjectSizeBytes) + require.ErrorIs(t, err, errS3PutAdmissionEntityTooLarge) + require.Zero(t, bytes) + require.Equal(t, s3PutAdmissionProtocolChunked, protocol) +} + func TestS3PutAdmissionProbeUsesChunkHeadroomForFixedLength(t *testing.T) { t.Parallel() @@ -174,6 +188,29 @@ func TestS3Server_PutObjectAdmissionRejectsFixedLengthOverS3LimitAsEntityTooLarg require.Zero(t, observer.lastInflight) } +func TestS3Server_PutObjectAdmissionRejectsChunkedOverS3LimitBeforeSlowDown(t *testing.T) { + t.Parallel() + + server, observer := newS3AdmissionTestServer(t, time.Second) + createS3AdmissionTestBucket(t, server, "admit-chunked-limit") + + release, err := server.putAdmission.acquire(context.Background(), s3ChunkSize) + require.NoError(t, err) + defer release() + + rec := httptest.NewRecorder() + req := newS3TestRequest(http.MethodPut, "/admit-chunked-limit/too-large.bin", http.NoBody) + req.Header.Set("Content-Encoding", "aws-chunked") + req.Header.Set("X-Amz-Content-Sha256", s3StreamingUnsignedPayloadTrailer) + req.Header.Set("X-Amz-Decoded-Content-Length", strconv.FormatInt(s3MaxObjectSizeBytes+1, 10)) + server.handle(rec, req) + + require.Equal(t, http.StatusRequestEntityTooLarge, rec.Code, rec.Body.String()) + require.Contains(t, rec.Body.String(), "EntityTooLarge") + require.Equal(t, 1, observer.rejections[s3PutAdmissionStagePrereserve+"|"+s3PutAdmissionProtocolChunked]) + require.Zero(t, observer.rejections[s3PutAdmissionStagePerBatch+"|"+s3PutAdmissionProtocolChunked]) +} + func TestS3Server_PutObjectAdmissionRequiresContentLengthForPlainPut(t *testing.T) { t.Parallel() diff --git a/adapter/s3_test.go b/adapter/s3_test.go index 7f03c2114..7287a4b3c 100644 --- a/adapter/s3_test.go +++ b/adapter/s3_test.go @@ -382,6 +382,53 @@ func TestS3Server_PutObjectStreamingRejectsBadDecodedLength(t *testing.T) { require.Contains(t, rec.Body.String(), "InvalidRequest") } +func TestReadS3PutChunkPreservesFullBufferReaderError(t *testing.T) { + t.Parallel() + + buf := make([]byte, 4) + reader := &s3FullBufferErrorReader{err: newAwsChunkedError(io.ErrUnexpectedEOF)} + + n, err := readS3PutChunk(reader, buf, true) + + require.Equal(t, len(buf), n) + var chunkedErr *awsChunkedError + require.ErrorAs(t, err, &chunkedErr) +} + +type s3FullBufferErrorReader struct { + err error +} + +func (r *s3FullBufferErrorReader) Read(p []byte) (int, error) { + for i := range p { + p[i] = 'x' + } + return len(p), r.err +} + +func TestS3Server_PutObjectStreamingRejectsTruncatedChunk(t *testing.T) { + t.Parallel() + + st := store.NewMVCCStore() + server := NewS3Server(nil, "", st, newLocalAdapterCoordinator(st), nil) + rec := httptest.NewRecorder() + server.handle(rec, newS3TestRequest(http.MethodPut, "/bkt-truncated", nil)) + require.Equal(t, http.StatusOK, rec.Code) + + rec = httptest.NewRecorder() + req := newS3TestRequest(http.MethodPut, "/bkt-truncated/obj.bin", strings.NewReader("5\r\nhel")) + req.Header.Set("Content-Encoding", "aws-chunked") + req.Header.Set("X-Amz-Content-Sha256", s3StreamingUnsignedPayloadTrailer) + req.Header.Set("X-Amz-Decoded-Content-Length", "5") + server.handle(rec, req) + require.Equal(t, http.StatusBadRequest, rec.Code) + require.Contains(t, rec.Body.String(), "InvalidRequest") + + rec = httptest.NewRecorder() + server.handle(rec, newS3TestRequest(http.MethodGet, "/bkt-truncated/obj.bin", nil)) + require.Equal(t, http.StatusNotFound, rec.Code) +} + func TestS3Server_PutObjectStreamingRejectsSignedPayloadMarker(t *testing.T) { t.Parallel() @@ -484,6 +531,34 @@ func TestS3Server_UploadPartStreamingWithTrailerChecksum(t *testing.T) { require.Contains(t, rec.Body.String(), "BadDigest") } +func TestS3Server_UploadPartStreamingRejectsTruncatedChunk(t *testing.T) { + t.Parallel() + + st := store.NewMVCCStore() + server := NewS3Server(nil, "", st, newLocalAdapterCoordinator(st), nil) + + rec := httptest.NewRecorder() + server.handle(rec, newS3TestRequest(http.MethodPut, "/bkt-part-truncated", nil)) + require.Equal(t, http.StatusOK, rec.Code) + + rec = httptest.NewRecorder() + server.handle(rec, newS3TestRequest(http.MethodPost, "/bkt-part-truncated/blob.bin?uploads=", nil)) + require.Equal(t, http.StatusOK, rec.Code) + var init s3InitiateMultipartUploadResult + require.NoError(t, xml.Unmarshal(rec.Body.Bytes(), &init)) + + rec = httptest.NewRecorder() + req := newS3TestRequest(http.MethodPut, + fmt.Sprintf("/bkt-part-truncated/blob.bin?uploadId=%s&partNumber=1", init.UploadId), + strings.NewReader("5\r\nhel")) + req.Header.Set("Content-Encoding", "aws-chunked") + req.Header.Set("X-Amz-Content-Sha256", s3StreamingUnsignedPayloadTrailer) + req.Header.Set("X-Amz-Decoded-Content-Length", "5") + server.handle(rec, req) + require.Equal(t, http.StatusBadRequest, rec.Code) + require.Contains(t, rec.Body.String(), "InvalidRequest") +} + func TestS3Server_ProxiesFollowerRequestsBeforeAuth(t *testing.T) { t.Parallel() From 1df92444d24a580a4b0976dedc0fcb0f3cb6cb01 Mon Sep 17 00:00:00 2001 From: bootjp Date: Wed, 8 Jul 2026 11:41:27 +0900 Subject: [PATCH 6/8] s3: preserve protocol errors before admission --- adapter/s3.go | 77 ++++++++-------- adapter/s3_admission_test.go | 165 ++++++++++++++++++++++++++++++++++- adapter/s3_test.go | 26 +++++- 3 files changed, 226 insertions(+), 42 deletions(-) diff --git a/adapter/s3.go b/adapter/s3.go index ebbe4f6b9..7557c3d0f 100644 --- a/adapter/s3.go +++ b/adapter/s3.go @@ -872,14 +872,14 @@ func (s *S3Server) putObject(w http.ResponseWriter, r *http.Request, bucket stri expectedPayloadSHA := normalizeS3PayloadHash(r.Header.Get("X-Amz-Content-Sha256")) validatePayloadSHA := expectedPayloadSHA != "" && !isS3PayloadMarker(expectedPayloadSHA) admissionProtocol := s3PutAdmissionProtocolForPayload(expectedPayloadSHA) - if !s.admitS3PutRequest(w, r, bucket, objectKey, s3MaxObjectSizeBytes, "object exceeds maximum allowed size") { - return - } streamBody, bodyErr := prepareStreamingPutBody(w, r, s3MaxObjectSizeBytes, expectedPayloadSHA, "object exceeds maximum allowed size") if bodyErr != nil { writeS3Error(w, bodyErr.Status, bodyErr.Code, bodyErr.Message, bucket, objectKey) return } + if !s.admitS3PutRequest(w, r, bucket, objectKey, s3MaxObjectSizeBytes, "object exceeds maximum allowed size") { + return + } part := s3ObjectPart{PartNo: 1} sizeBytes := int64(0) chunkNo := uint64(0) @@ -932,11 +932,21 @@ func (s *S3Server) putObject(w http.ResponseWriter, r *http.Request, bucket stri return } } - readBuf, exactLength := nextReadBuffer() + readBuf, finalRead := nextReadBuffer() if len(readBuf) == 0 { break } - n, readErr := readS3PutChunk(r.Body, readBuf, !exactLength) + allowFinalPartial := admissionProtocol != s3PutAdmissionProtocolFixed + n, readErr := readS3PutChunk(r.Body, readBuf, allowFinalPartial, finalRead) + if readErr != nil && !errors.Is(readErr, io.EOF) { + s.cleanupManifestBlobs(r.Context(), bucket, meta.Generation, objectKey, uploadedManifest()) + if be, ok := classifyS3BodyReadErr(readErr, "object exceeds maximum allowed size"); ok { + writeS3Error(w, be.Status, be.Code, be.Message, bucket, objectKey) + return + } + writeS3InternalError(w, readErr) + return + } if n > 0 { releaseAdmission, admissionErr := s.acquireS3PutAdmission(r.Context(), int64(n), admissionProtocol) if admissionErr != nil { @@ -979,15 +989,6 @@ func (s *S3Server) putObject(w http.ResponseWriter, r *http.Request, bucket stri if errors.Is(readErr, io.EOF) { break } - if readErr != nil { - s.cleanupManifestBlobs(r.Context(), bucket, meta.Generation, objectKey, uploadedManifest()) - if be, ok := classifyS3BodyReadErr(readErr, "object exceeds maximum allowed size"); ok { - writeS3Error(w, be.Status, be.Code, be.Message, bucket, objectKey) - return - } - writeS3InternalError(w, readErr) - return - } } if err := flushBatch(); err != nil { s.cleanupManifestBlobs(r.Context(), bucket, meta.Generation, objectKey, uploadedManifest()) @@ -1412,14 +1413,14 @@ func (s *S3Server) uploadPart(w http.ResponseWriter, r *http.Request, bucket str partPayloadSHA := normalizeS3PayloadHash(r.Header.Get("X-Amz-Content-Sha256")) admissionProtocol := s3PutAdmissionProtocolForPayload(partPayloadSHA) - if !s.admitS3PutRequest(w, r, bucket, objectKey, s3MaxPartSizeBytes, "part exceeds maximum allowed size") { - return - } partStreamBody, bodyErr := prepareStreamingPutBody(w, r, s3MaxPartSizeBytes, partPayloadSHA, "part exceeds maximum allowed size") if bodyErr != nil { writeS3Error(w, bodyErr.Status, bodyErr.Code, bodyErr.Message, bucket, objectKey) return } + if !s.admitS3PutRequest(w, r, bucket, objectKey, s3MaxPartSizeBytes, "part exceeds maximum allowed size") { + return + } // Pre-allocate the part's commit timestamp before writing any blob chunks so // that the same version identifier is used for every chunk in this attempt. @@ -1484,23 +1485,24 @@ func (s *S3Server) uploadPart(w http.ResponseWriter, r *http.Request, bucket str return } } - readBuf, exactLength := nextReadBuffer() + readBuf, finalRead := nextReadBuffer() if len(readBuf) == 0 { break } - n, readErr := readS3PutChunk(r.Body, readBuf, !exactLength) + allowFinalPartial := admissionProtocol != s3PutAdmissionProtocolFixed + n, readErr := readS3PutChunk(r.Body, readBuf, allowFinalPartial, finalRead) + if readErr != nil && !errors.Is(readErr, io.EOF) { + if be, ok := classifyS3BodyReadErr(readErr, "part exceeds maximum allowed size"); ok { + writeS3Error(w, be.Status, be.Code, be.Message, bucket, objectKey) + return + } + writeS3InternalError(w, readErr) + return + } if n == 0 { if errors.Is(readErr, io.EOF) { break } - if readErr != nil { - if be, ok := classifyS3BodyReadErr(readErr, "part exceeds maximum allowed size"); ok { - writeS3Error(w, be.Status, be.Code, be.Message, bucket, objectKey) - return - } - writeS3InternalError(w, readErr) - return - } continue } releaseAdmission, admissionErr := s.acquireS3PutAdmission(r.Context(), int64(n), admissionProtocol) @@ -1534,14 +1536,6 @@ func (s *S3Server) uploadPart(w http.ResponseWriter, r *http.Request, bucket str if errors.Is(readErr, io.EOF) { break } - if readErr != nil { - if be, ok := classifyS3BodyReadErr(readErr, "part exceeds maximum allowed size"); ok { - writeS3Error(w, be.Status, be.Code, be.Message, bucket, objectKey) - return - } - writeS3InternalError(w, readErr) - return - } } if err := flushBatch(); err != nil { writeS3InternalError(w, err) @@ -2546,13 +2540,13 @@ func nextS3PutReadBuffer(buf []byte, protocol string, contentLength int64, readB if remaining <= 0 { return nil, true } - if remaining < int64(len(buf)) { + if remaining <= int64(len(buf)) { return buf[:int(remaining)], true } - return buf, true + return buf, false } -func readS3PutChunk(r io.Reader, buf []byte, allowFinalPartial bool) (int, error) { +func readS3PutChunk(r io.Reader, buf []byte, allowFinalPartial bool, allowEOFWithFullBuffer bool) (int, error) { total := 0 for total < len(buf) { n, err := r.Read(buf[total:]) @@ -2560,7 +2554,7 @@ func readS3PutChunk(r io.Reader, buf []byte, allowFinalPartial bool) (int, error total += n } if errors.Is(err, io.EOF) { - return s3PutChunkEOFResult(total, len(buf), allowFinalPartial) + return s3PutChunkEOFResult(total, len(buf), allowFinalPartial, allowEOFWithFullBuffer) } if err != nil { return total, errors.WithStack(err) @@ -2572,7 +2566,10 @@ func readS3PutChunk(r io.Reader, buf []byte, allowFinalPartial bool) (int, error return total, nil } -func s3PutChunkEOFResult(total, bufLen int, allowFinalPartial bool) (int, error) { +func s3PutChunkEOFResult(total, bufLen int, allowFinalPartial bool, allowEOFWithFullBuffer bool) (int, error) { + if allowEOFWithFullBuffer && total == bufLen { + return total, io.EOF + } if !allowFinalPartial && (total > 0 || bufLen > 0) { return total, errS3PutIncompleteBody } diff --git a/adapter/s3_admission_test.go b/adapter/s3_admission_test.go index 5d85a122e..d0153f882 100644 --- a/adapter/s3_admission_test.go +++ b/adapter/s3_admission_test.go @@ -207,10 +207,42 @@ func TestS3Server_PutObjectAdmissionRejectsChunkedOverS3LimitBeforeSlowDown(t *t require.Equal(t, http.StatusRequestEntityTooLarge, rec.Code, rec.Body.String()) require.Contains(t, rec.Body.String(), "EntityTooLarge") - require.Equal(t, 1, observer.rejections[s3PutAdmissionStagePrereserve+"|"+s3PutAdmissionProtocolChunked]) + require.Zero(t, observer.rejections[s3PutAdmissionStagePrereserve+"|"+s3PutAdmissionProtocolChunked]) require.Zero(t, observer.rejections[s3PutAdmissionStagePerBatch+"|"+s3PutAdmissionProtocolChunked]) } +func TestS3Server_PutObjectAdmissionValidatesProtocolBeforeSlowDown(t *testing.T) { + t.Parallel() + + server, observer := newS3AdmissionTestServer(t, time.Second) + createS3AdmissionTestBucket(t, server, "admit-protocol") + + release, err := server.putAdmission.acquire(context.Background(), s3ChunkSize) + require.NoError(t, err) + defer release() + + rec := httptest.NewRecorder() + req := newS3TestRequest(http.MethodPut, "/admit-protocol/signed.bin", strings.NewReader("ignored")) + req.Header.Set("Content-Encoding", "aws-chunked") + req.Header.Set("X-Amz-Content-Sha256", s3StreamingSignedPayload) + req.Header.Set("X-Amz-Decoded-Content-Length", "7") + server.handle(rec, req) + require.Equal(t, http.StatusNotImplemented, rec.Code, rec.Body.String()) + require.Contains(t, rec.Body.String(), "NotImplemented") + + rec = httptest.NewRecorder() + req = newS3TestRequest(http.MethodPut, "/admit-protocol/malformed.bin", nil) + req.Body = io.NopCloser(strings.NewReader("5\r\nhello\r\n0\r\n\r\n")) + req.ContentLength = -1 + req.Header.Set("Content-Encoding", "aws-chunked") + req.Header.Set("X-Amz-Content-Sha256", sha256Hex("hello")) + server.handle(rec, req) + require.Equal(t, http.StatusBadRequest, rec.Code, rec.Body.String()) + require.Contains(t, rec.Body.String(), "InvalidRequest") + require.Zero(t, observer.rejections[s3PutAdmissionStagePrereserve+"|"+s3PutAdmissionProtocolChunked]) + require.Zero(t, observer.rejections[s3PutAdmissionStagePrereserve+"|"+s3PutAdmissionProtocolFixed]) +} + func TestS3Server_PutObjectAdmissionRequiresContentLengthForPlainPut(t *testing.T) { t.Parallel() @@ -282,6 +314,28 @@ func TestS3Server_PutObjectRejectsTruncatedFixedLengthBody(t *testing.T) { require.Equal(t, http.StatusNotFound, rec.Code) } +func TestS3Server_PutObjectAdmissionReportsReadErrorBeforeSlowDown(t *testing.T) { + t.Parallel() + + server, observer := newS3AdmissionTestServer(t, 5*time.Millisecond) + createS3AdmissionTestBucket(t, server, "admit-read-error") + + reader := &s3AdmissionFillingErrorReader{ + server: server, + payload: []byte("bad"), + err: newAwsChunkedError(io.ErrUnexpectedEOF), + } + defer reader.releaseHeld() + rec := httptest.NewRecorder() + req := newS3TestRequest(http.MethodPut, "/admit-read-error/body.bin", reader) + req.ContentLength = int64(len(reader.payload)) + server.handle(rec, req) + + require.Equal(t, http.StatusBadRequest, rec.Code, rec.Body.String()) + require.Contains(t, rec.Body.String(), "InvalidRequest") + require.Zero(t, observer.rejections[s3PutAdmissionStagePerBatch+"|"+s3PutAdmissionProtocolFixed]) +} + func TestS3Server_PutObjectAdmissionChunkedStreamsPastSingleChunkBudget(t *testing.T) { t.Parallel() @@ -308,6 +362,46 @@ func TestS3Server_PutObjectAdmissionChunkedStreamsPastSingleChunkBudget(t *testi require.Equal(t, string(payload), rec.Body.String()) } +func TestS3Server_UploadPartAdmissionValidatesProtocolBeforeSlowDown(t *testing.T) { + t.Parallel() + + server, observer := newS3AdmissionTestServer(t, time.Second) + uploadID := initS3AdmissionMultipartUpload(t, server, "admit-part-protocol", "object.bin") + + release, err := server.putAdmission.acquire(context.Background(), s3ChunkSize) + require.NoError(t, err) + defer release() + + rec := httptest.NewRecorder() + req := newS3TestRequest( + http.MethodPut, + fmt.Sprintf("/admit-part-protocol/object.bin?uploadId=%s&partNumber=1", uploadID), + strings.NewReader("ignored"), + ) + req.Header.Set("Content-Encoding", "aws-chunked") + req.Header.Set("X-Amz-Content-Sha256", s3StreamingSignedPayload) + req.Header.Set("X-Amz-Decoded-Content-Length", "7") + server.handle(rec, req) + require.Equal(t, http.StatusNotImplemented, rec.Code, rec.Body.String()) + require.Contains(t, rec.Body.String(), "NotImplemented") + + rec = httptest.NewRecorder() + req = newS3TestRequest( + http.MethodPut, + fmt.Sprintf("/admit-part-protocol/object.bin?uploadId=%s&partNumber=2", uploadID), + nil, + ) + req.Body = io.NopCloser(strings.NewReader("5\r\nhello\r\n0\r\n\r\n")) + req.ContentLength = -1 + req.Header.Set("Content-Encoding", "aws-chunked") + req.Header.Set("X-Amz-Content-Sha256", sha256Hex("hello")) + server.handle(rec, req) + require.Equal(t, http.StatusBadRequest, rec.Code, rec.Body.String()) + require.Contains(t, rec.Body.String(), "InvalidRequest") + require.Zero(t, observer.rejections[s3PutAdmissionStagePrereserve+"|"+s3PutAdmissionProtocolChunked]) + require.Zero(t, observer.rejections[s3PutAdmissionStagePrereserve+"|"+s3PutAdmissionProtocolFixed]) +} + func TestS3Server_UploadPartAdmissionAllowsFixedLengthLargerThanBudget(t *testing.T) { t.Parallel() @@ -341,6 +435,62 @@ func TestS3Server_UploadPartAdmissionAllowsFixedLengthLargerThanBudget(t *testin require.Zero(t, server.putAdmission.inflight.Load()) } +func TestS3Server_UploadPartAdmissionReportsReadErrorBeforeSlowDown(t *testing.T) { + t.Parallel() + + server, observer := newS3AdmissionTestServer(t, 5*time.Millisecond) + uploadID := initS3AdmissionMultipartUpload(t, server, "admit-part-read-error", "object.bin") + + reader := &s3AdmissionFillingErrorReader{ + server: server, + payload: []byte("bad"), + err: newAwsChunkedError(io.ErrUnexpectedEOF), + } + defer reader.releaseHeld() + rec := httptest.NewRecorder() + req := newS3TestRequest( + http.MethodPut, + fmt.Sprintf("/admit-part-read-error/object.bin?uploadId=%s&partNumber=1", uploadID), + reader, + ) + req.ContentLength = int64(len(reader.payload)) + server.handle(rec, req) + + require.Equal(t, http.StatusBadRequest, rec.Code, rec.Body.String()) + require.Contains(t, rec.Body.String(), "InvalidRequest") + require.Zero(t, observer.rejections[s3PutAdmissionStagePerBatch+"|"+s3PutAdmissionProtocolFixed]) +} + +type s3AdmissionFillingErrorReader struct { + server *S3Server + payload []byte + err error + release func() + done bool +} + +func (r *s3AdmissionFillingErrorReader) Read(p []byte) (int, error) { + if r.done { + return 0, io.EOF + } + if r.server != nil && r.server.putAdmission != nil { + release, err := r.server.putAdmission.acquire(context.Background(), s3ChunkSize) + if err == nil { + r.release = release + } + } + r.done = true + return copy(p, r.payload), r.err +} + +func (r *s3AdmissionFillingErrorReader) releaseHeld() { + if r.release == nil { + return + } + r.release() + r.release = nil +} + func newS3AdmissionTestServer(t *testing.T, timeout time.Duration) (*S3Server, *recordingS3PutAdmissionObserver) { t.Helper() @@ -365,3 +515,16 @@ func createS3AdmissionTestBucket(t *testing.T, server *S3Server, bucket string) server.handle(rec, newS3TestRequest(http.MethodPut, "/"+bucket, nil)) require.Equal(t, http.StatusOK, rec.Code, rec.Body.String()) } + +func initS3AdmissionMultipartUpload(t *testing.T, server *S3Server, bucket, object string) string { + t.Helper() + + createS3AdmissionTestBucket(t, server, bucket) + rec := httptest.NewRecorder() + server.handle(rec, newS3TestRequest(http.MethodPost, fmt.Sprintf("/%s/%s?uploads=", bucket, object), nil)) + require.Equal(t, http.StatusOK, rec.Code, rec.Body.String()) + var initResult s3InitiateMultipartUploadResult + require.NoError(t, xml.Unmarshal(rec.Body.Bytes(), &initResult)) + require.NotEmpty(t, initResult.UploadId) + return initResult.UploadId +} diff --git a/adapter/s3_test.go b/adapter/s3_test.go index 7287a4b3c..0c9faa61b 100644 --- a/adapter/s3_test.go +++ b/adapter/s3_test.go @@ -388,13 +388,37 @@ func TestReadS3PutChunkPreservesFullBufferReaderError(t *testing.T) { buf := make([]byte, 4) reader := &s3FullBufferErrorReader{err: newAwsChunkedError(io.ErrUnexpectedEOF)} - n, err := readS3PutChunk(reader, buf, true) + n, err := readS3PutChunk(reader, buf, true, false) require.Equal(t, len(buf), n) var chunkedErr *awsChunkedError require.ErrorAs(t, err, &chunkedErr) } +func TestReadS3PutChunkAcceptsFinalFullBufferEOF(t *testing.T) { + t.Parallel() + + buf := make([]byte, 4) + reader := &s3FullBufferErrorReader{err: io.EOF} + + n, err := readS3PutChunk(reader, buf, false, true) + + require.Equal(t, len(buf), n) + require.ErrorIs(t, err, io.EOF) +} + +func TestReadS3PutChunkRejectsNonFinalFullBufferEOF(t *testing.T) { + t.Parallel() + + buf := make([]byte, 4) + reader := &s3FullBufferErrorReader{err: io.EOF} + + n, err := readS3PutChunk(reader, buf, false, false) + + require.Equal(t, len(buf), n) + require.ErrorIs(t, err, errS3PutIncompleteBody) +} + type s3FullBufferErrorReader struct { err error } From a15c8d8099a8fc3adc6e5c0505f080eded7de493 Mon Sep 17 00:00:00 2001 From: bootjp Date: Wed, 8 Jul 2026 12:31:53 +0900 Subject: [PATCH 7/8] s3: reserve admission before body reads --- adapter/s3.go | 42 ++++++---- adapter/s3_admission.go | 8 +- adapter/s3_admission_test.go | 145 +++++++++++++++++++++++++++++++++-- adapter/s3_chunked.go | 7 ++ 4 files changed, 179 insertions(+), 23 deletions(-) diff --git a/adapter/s3.go b/adapter/s3.go index 7557c3d0f..2f6366a46 100644 --- a/adapter/s3.go +++ b/adapter/s3.go @@ -936,9 +936,16 @@ func (s *S3Server) putObject(w http.ResponseWriter, r *http.Request, bucket stri if len(readBuf) == 0 { break } - allowFinalPartial := admissionProtocol != s3PutAdmissionProtocolFixed + releaseAdmission, admissionErr := s.acquireS3PutAdmissionForRead(r.Context(), len(readBuf), admissionProtocol, streamBody, sizeBytes) + if admissionErr != nil { + s.cleanupManifestBlobs(r.Context(), bucket, meta.Generation, objectKey, uploadedManifest()) + writeS3AdmissionError(w, bucket, objectKey, s.s3PutAdmissionRetryAfter()) + return + } + allowFinalPartial := s3PutReadAllowsFinalPartial(admissionProtocol, r.ContentLength) n, readErr := readS3PutChunk(r.Body, readBuf, allowFinalPartial, finalRead) if readErr != nil && !errors.Is(readErr, io.EOF) { + releaseAdmission() s.cleanupManifestBlobs(r.Context(), bucket, meta.Generation, objectKey, uploadedManifest()) if be, ok := classifyS3BodyReadErr(readErr, "object exceeds maximum allowed size"); ok { writeS3Error(w, be.Status, be.Code, be.Message, bucket, objectKey) @@ -948,12 +955,6 @@ func (s *S3Server) putObject(w http.ResponseWriter, r *http.Request, bucket stri return } if n > 0 { - releaseAdmission, admissionErr := s.acquireS3PutAdmission(r.Context(), int64(n), admissionProtocol) - if admissionErr != nil { - s.cleanupManifestBlobs(r.Context(), bucket, meta.Generation, objectKey, uploadedManifest()) - writeS3AdmissionError(w, bucket, objectKey, s.s3PutAdmissionRetryAfter()) - return - } pendingAdmission = append(pendingAdmission, releaseAdmission) chunk := append([]byte(nil), readBuf[:n]...) if _, err := hasher.Write(chunk); err != nil { @@ -985,6 +986,8 @@ func (s *S3Server) putObject(w http.ResponseWriter, r *http.Request, bucket stri } sizeBytes += int64(n) chunkNo++ + } else { + releaseAdmission() } if errors.Is(readErr, io.EOF) { break @@ -1489,9 +1492,15 @@ func (s *S3Server) uploadPart(w http.ResponseWriter, r *http.Request, bucket str if len(readBuf) == 0 { break } - allowFinalPartial := admissionProtocol != s3PutAdmissionProtocolFixed + releaseAdmission, admissionErr := s.acquireS3PutAdmissionForRead(r.Context(), len(readBuf), admissionProtocol, partStreamBody, sizeBytes) + if admissionErr != nil { + writeS3AdmissionError(w, bucket, objectKey, s.s3PutAdmissionRetryAfter()) + return + } + allowFinalPartial := s3PutReadAllowsFinalPartial(admissionProtocol, r.ContentLength) n, readErr := readS3PutChunk(r.Body, readBuf, allowFinalPartial, finalRead) if readErr != nil && !errors.Is(readErr, io.EOF) { + releaseAdmission() if be, ok := classifyS3BodyReadErr(readErr, "part exceeds maximum allowed size"); ok { writeS3Error(w, be.Status, be.Code, be.Message, bucket, objectKey) return @@ -1500,16 +1509,12 @@ func (s *S3Server) uploadPart(w http.ResponseWriter, r *http.Request, bucket str return } if n == 0 { + releaseAdmission() if errors.Is(readErr, io.EOF) { break } continue } - releaseAdmission, admissionErr := s.acquireS3PutAdmission(r.Context(), int64(n), admissionProtocol) - if admissionErr != nil { - writeS3AdmissionError(w, bucket, objectKey, s.s3PutAdmissionRetryAfter()) - return - } pendingAdmission = append(pendingAdmission, releaseAdmission) chunk := append([]byte(nil), readBuf[:n]...) if _, err := hasher.Write(chunk); err != nil { @@ -2546,6 +2551,17 @@ func nextS3PutReadBuffer(buf []byte, protocol string, contentLength int64, readB return buf, false } +func (s *S3Server) acquireS3PutAdmissionForRead(ctx context.Context, bytes int, protocol string, body *s3StreamingBody, readBytes int64) (func(), error) { + if body.decodedLengthReached(readBytes) { + return func() {}, nil + } + return s.acquireS3PutAdmission(ctx, int64(bytes), protocol) +} + +func s3PutReadAllowsFinalPartial(protocol string, contentLength int64) bool { + return protocol != s3PutAdmissionProtocolFixed || contentLength < 0 +} + func readS3PutChunk(r io.Reader, buf []byte, allowFinalPartial bool, allowEOFWithFullBuffer bool) (int, error) { total := 0 for total < len(buf) { diff --git a/adapter/s3_admission.go b/adapter/s3_admission.go index 21933c835..f40ad5b81 100644 --- a/adapter/s3_admission.go +++ b/adapter/s3_admission.go @@ -16,6 +16,7 @@ import ( const ( s3PutAdmissionDefaultMaxInflightBytes = 256 << 20 s3PutAdmissionDefaultTimeout = 30 * time.Second + s3PutAdmissionDefaultRetryAfter = time.Second s3PutAdmissionDisableMaxInflightBytes = int64(1<<63 - 1) s3PutAdmissionMaxInflightEnv = "ELASTICKV_S3_PUT_ADMISSION_MAX_INFLIGHT_BYTES" @@ -234,7 +235,7 @@ func (s *S3Server) admitS3PutRequest(w http.ResponseWriter, r *http.Request, buc } if err := s.putAdmission.peekHeadroom(bytes); err != nil { s.observeS3PutAdmissionRejection(s3PutAdmissionStagePrereserve, protocol) - writeS3AdmissionError(w, bucket, objectKey, s.putAdmission.timeout) + writeS3AdmissionError(w, bucket, objectKey, s.s3PutAdmissionRetryAfter()) return false } return true @@ -324,10 +325,7 @@ func (s *S3Server) shouldFlushS3PutBatchBeforeRead(protocol string, pendingAdmis } func (s *S3Server) s3PutAdmissionRetryAfter() time.Duration { - if s == nil || s.putAdmission == nil { - return s3PutAdmissionDefaultTimeout - } - return s.putAdmission.timeout + return s3PutAdmissionDefaultRetryAfter } func writeS3AdmissionError(w http.ResponseWriter, bucket, objectKey string, retryAfter time.Duration) { diff --git a/adapter/s3_admission_test.go b/adapter/s3_admission_test.go index d0153f882..72d24f357 100644 --- a/adapter/s3_admission_test.go +++ b/adapter/s3_admission_test.go @@ -279,6 +279,43 @@ func TestS3Server_PutObjectAdmissionAllowsSmallRequestAndReleasesBudget(t *testi require.Equal(t, "ok", rec.Body.String()) } +func TestS3Server_PutObjectAdmissionAcquiresBeforeBodyRead(t *testing.T) { + t.Parallel() + + server, observer := newS3AdmissionTestServer(t, time.Second) + createS3AdmissionTestBucket(t, server, "admit-before-read") + + reader := &s3AdmissionObservedReader{ + server: server, + payload: []byte("bounded"), + } + rec := httptest.NewRecorder() + req := newS3TestRequest(http.MethodPut, "/admit-before-read/object.bin", reader) + req.ContentLength = int64(len(reader.payload)) + server.handle(rec, req) + + require.Equal(t, http.StatusOK, rec.Code, rec.Body.String()) + require.EqualValues(t, s3ChunkSize, reader.firstReadInflight) + require.Zero(t, observer.lastInflight) +} + +func TestS3Server_PutObjectAdmissionRetryAfterIgnoresWaitTimeout(t *testing.T) { + t.Parallel() + + server, _ := newS3AdmissionTestServer(t, 45*time.Second) + createS3AdmissionTestBucket(t, server, "admit-retry-after") + release, err := server.putAdmission.acquire(context.Background(), s3ChunkSize) + require.NoError(t, err) + defer release() + + rec := httptest.NewRecorder() + server.handle(rec, newS3TestRequest(http.MethodPut, "/admit-retry-after/blocked.bin", strings.NewReader("blocked"))) + + require.Equal(t, http.StatusServiceUnavailable, rec.Code, rec.Body.String()) + require.Contains(t, rec.Body.String(), "SlowDown") + require.Equal(t, "1", rec.Header().Get("Retry-After")) +} + func TestS3Server_PutObjectAdmissionExactFixedLengthChunkDoesNotChargeEOF(t *testing.T) { t.Parallel() @@ -362,11 +399,91 @@ func TestS3Server_PutObjectAdmissionChunkedStreamsPastSingleChunkBudget(t *testi require.Equal(t, string(payload), rec.Body.String()) } +func TestS3Server_PutObjectAdmissionDisabledAllowsPlainChunkedEOF(t *testing.T) { + t.Parallel() + + st := store.NewMVCCStore() + server := NewS3Server( + nil, + "", + st, + newLocalAdapterCoordinator(st), + nil, + withS3PutAdmissionForTest(s3PutAdmissionDisableMaxInflightBytes, time.Second), + ) + createS3AdmissionTestBucket(t, server, "admit-disabled") + + payload := "plain chunked body" + rec := httptest.NewRecorder() + req := newS3TestRequest(http.MethodPut, "/admit-disabled/chunked.bin", strings.NewReader(payload)) + req.ContentLength = -1 + server.handle(rec, req) + + require.Equal(t, http.StatusOK, rec.Code, rec.Body.String()) + + rec = httptest.NewRecorder() + server.handle(rec, newS3TestRequest(http.MethodGet, "/admit-disabled/chunked.bin", nil)) + require.Equal(t, http.StatusOK, rec.Code) + require.Equal(t, payload, rec.Body.String()) +} + +func TestS3Server_UploadPartAdmissionAcquiresBeforeBodyRead(t *testing.T) { + t.Parallel() + + server, observer := newS3AdmissionTestServer(t, time.Second) + uploadID := initS3AdmissionMultipartUpload(t, server, "admit-part-before-read") + + reader := &s3AdmissionObservedReader{ + server: server, + payload: []byte("part"), + } + rec := httptest.NewRecorder() + req := newS3TestRequest( + http.MethodPut, + fmt.Sprintf("/admit-part-before-read/object.bin?uploadId=%s&partNumber=1", uploadID), + reader, + ) + req.ContentLength = int64(len(reader.payload)) + server.handle(rec, req) + + require.Equal(t, http.StatusOK, rec.Code, rec.Body.String()) + require.EqualValues(t, s3ChunkSize, reader.firstReadInflight) + require.Zero(t, observer.lastInflight) +} + +func TestS3Server_UploadPartAdmissionDisabledAllowsPlainChunkedEOF(t *testing.T) { + t.Parallel() + + st := store.NewMVCCStore() + server := NewS3Server( + nil, + "", + st, + newLocalAdapterCoordinator(st), + nil, + withS3PutAdmissionForTest(s3PutAdmissionDisableMaxInflightBytes, time.Second), + ) + uploadID := initS3AdmissionMultipartUpload(t, server, "admit-part-disabled") + + payload := "plain chunked part" + rec := httptest.NewRecorder() + req := newS3TestRequest( + http.MethodPut, + fmt.Sprintf("/admit-part-disabled/object.bin?uploadId=%s&partNumber=1", uploadID), + strings.NewReader(payload), + ) + req.ContentLength = -1 + server.handle(rec, req) + + require.Equal(t, http.StatusOK, rec.Code, rec.Body.String()) + require.Equal(t, quoteS3ETag(md5Hex(payload)), rec.Header().Get("ETag")) +} + func TestS3Server_UploadPartAdmissionValidatesProtocolBeforeSlowDown(t *testing.T) { t.Parallel() server, observer := newS3AdmissionTestServer(t, time.Second) - uploadID := initS3AdmissionMultipartUpload(t, server, "admit-part-protocol", "object.bin") + uploadID := initS3AdmissionMultipartUpload(t, server, "admit-part-protocol") release, err := server.putAdmission.acquire(context.Background(), s3ChunkSize) require.NoError(t, err) @@ -439,7 +556,7 @@ func TestS3Server_UploadPartAdmissionReportsReadErrorBeforeSlowDown(t *testing.T t.Parallel() server, observer := newS3AdmissionTestServer(t, 5*time.Millisecond) - uploadID := initS3AdmissionMultipartUpload(t, server, "admit-part-read-error", "object.bin") + uploadID := initS3AdmissionMultipartUpload(t, server, "admit-part-read-error") reader := &s3AdmissionFillingErrorReader{ server: server, @@ -469,11 +586,29 @@ type s3AdmissionFillingErrorReader struct { done bool } -func (r *s3AdmissionFillingErrorReader) Read(p []byte) (int, error) { +type s3AdmissionObservedReader struct { + server *S3Server + payload []byte + firstReadInflight int64 + done bool +} + +func (r *s3AdmissionObservedReader) Read(p []byte) (int, error) { if r.done { return 0, io.EOF } if r.server != nil && r.server.putAdmission != nil { + r.firstReadInflight = r.server.putAdmission.inflight.Load() + } + r.done = true + return copy(p, r.payload), io.EOF +} + +func (r *s3AdmissionFillingErrorReader) Read(p []byte) (int, error) { + if r.done { + return 0, io.EOF + } + if r.server != nil && r.server.putAdmission != nil && r.server.putAdmission.inflight.Load() == 0 { release, err := r.server.putAdmission.acquire(context.Background(), s3ChunkSize) if err == nil { r.release = release @@ -516,12 +651,12 @@ func createS3AdmissionTestBucket(t *testing.T, server *S3Server, bucket string) require.Equal(t, http.StatusOK, rec.Code, rec.Body.String()) } -func initS3AdmissionMultipartUpload(t *testing.T, server *S3Server, bucket, object string) string { +func initS3AdmissionMultipartUpload(t *testing.T, server *S3Server, bucket string) string { t.Helper() createS3AdmissionTestBucket(t, server, bucket) rec := httptest.NewRecorder() - server.handle(rec, newS3TestRequest(http.MethodPost, fmt.Sprintf("/%s/%s?uploads=", bucket, object), nil)) + server.handle(rec, newS3TestRequest(http.MethodPost, fmt.Sprintf("/%s/object.bin?uploads=", bucket), nil)) require.Equal(t, http.StatusOK, rec.Code, rec.Body.String()) var initResult s3InitiateMultipartUploadResult require.NoError(t, xml.Unmarshal(rec.Body.Bytes(), &initResult)) diff --git a/adapter/s3_chunked.go b/adapter/s3_chunked.go index 6f20b5deb..d2bcdea60 100644 --- a/adapter/s3_chunked.go +++ b/adapter/s3_chunked.go @@ -282,6 +282,13 @@ func (s *s3StreamingBody) writeDecoded(p []byte) { _, _ = s.trailerHash.Write(p) } +func (s *s3StreamingBody) decodedLengthReached(readBytes int64) bool { + return s != nil && + s.reader != nil && + s.reader.declaredDecoded >= 0 && + readBytes >= s.reader.declaredDecoded +} + // verifyTrailer must be called after the body has been read to EOF. It // returns a non-nil error when the client advertised a trailer checksum via // X-Amz-Trailer and the value that arrived in the chunked trailer does not From b050baf1325fec50de97c18efb71ac0716426cf2 Mon Sep 17 00:00:00 2001 From: bootjp Date: Wed, 8 Jul 2026 13:12:51 +0900 Subject: [PATCH 8/8] s3: map short fixed reads to incomplete body --- adapter/s3.go | 2 +- adapter/s3_admission_test.go | 53 ++++++++++++++++++++++++++++++++++++ 2 files changed, 54 insertions(+), 1 deletion(-) diff --git a/adapter/s3.go b/adapter/s3.go index 2f6366a46..cc33a0999 100644 --- a/adapter/s3.go +++ b/adapter/s3.go @@ -2731,7 +2731,7 @@ func classifyS3BodyReadErr(err error, tooLargeMessage string) (*s3PutBodyError, if errors.As(err, &chunkedErr) { return &s3PutBodyError{Status: http.StatusBadRequest, Code: "InvalidRequest", Message: chunkedErr.Error()}, true } - if errors.Is(err, errS3PutIncompleteBody) { + if errors.Is(err, errS3PutIncompleteBody) || errors.Is(err, io.ErrUnexpectedEOF) { return &s3PutBodyError{Status: http.StatusBadRequest, Code: "IncompleteBody", Message: errS3PutIncompleteBody.Error()}, true } return nil, false diff --git a/adapter/s3_admission_test.go b/adapter/s3_admission_test.go index 72d24f357..c63ace6b0 100644 --- a/adapter/s3_admission_test.go +++ b/adapter/s3_admission_test.go @@ -351,6 +351,26 @@ func TestS3Server_PutObjectRejectsTruncatedFixedLengthBody(t *testing.T) { require.Equal(t, http.StatusNotFound, rec.Code) } +func TestS3Server_PutObjectRejectsUnexpectedEOFFixedLengthBody(t *testing.T) { + t.Parallel() + + server, observer := newS3AdmissionTestServer(t, time.Second) + createS3AdmissionTestBucket(t, server, "admit-unexpected") + + rec := httptest.NewRecorder() + req := newS3TestRequest(http.MethodPut, "/admit-unexpected/short.bin", &s3UnexpectedEOFReader{payload: []byte("short")}) + req.ContentLength = int64(len("short") + 1) + server.handle(rec, req) + + require.Equal(t, http.StatusBadRequest, rec.Code, rec.Body.String()) + require.Contains(t, rec.Body.String(), "IncompleteBody") + require.Zero(t, observer.lastInflight) + + rec = httptest.NewRecorder() + server.handle(rec, newS3TestRequest(http.MethodGet, "/admit-unexpected/short.bin", nil)) + require.Equal(t, http.StatusNotFound, rec.Code) +} + func TestS3Server_PutObjectAdmissionReportsReadErrorBeforeSlowDown(t *testing.T) { t.Parallel() @@ -578,6 +598,26 @@ func TestS3Server_UploadPartAdmissionReportsReadErrorBeforeSlowDown(t *testing.T require.Zero(t, observer.rejections[s3PutAdmissionStagePerBatch+"|"+s3PutAdmissionProtocolFixed]) } +func TestS3Server_UploadPartRejectsUnexpectedEOFFixedLengthBody(t *testing.T) { + t.Parallel() + + server, observer := newS3AdmissionTestServer(t, time.Second) + uploadID := initS3AdmissionMultipartUpload(t, server, "admit-part-unexpected") + + rec := httptest.NewRecorder() + req := newS3TestRequest( + http.MethodPut, + fmt.Sprintf("/admit-part-unexpected/object.bin?uploadId=%s&partNumber=1", uploadID), + &s3UnexpectedEOFReader{payload: []byte("short")}, + ) + req.ContentLength = int64(len("short") + 1) + server.handle(rec, req) + + require.Equal(t, http.StatusBadRequest, rec.Code, rec.Body.String()) + require.Contains(t, rec.Body.String(), "IncompleteBody") + require.Zero(t, observer.lastInflight) +} + type s3AdmissionFillingErrorReader struct { server *S3Server payload []byte @@ -586,6 +626,11 @@ type s3AdmissionFillingErrorReader struct { done bool } +type s3UnexpectedEOFReader struct { + payload []byte + done bool +} + type s3AdmissionObservedReader struct { server *S3Server payload []byte @@ -618,6 +663,14 @@ func (r *s3AdmissionFillingErrorReader) Read(p []byte) (int, error) { return copy(p, r.payload), r.err } +func (r *s3UnexpectedEOFReader) Read(p []byte) (int, error) { + if r.done { + return 0, io.EOF + } + r.done = true + return copy(p, r.payload), io.ErrUnexpectedEOF +} + func (r *s3AdmissionFillingErrorReader) releaseHeld() { if r.release == nil { return