This is an automated email from the ASF dual-hosted git repository. wilfred-s pushed a commit to branch branch-1.9 in repository https://gitbox.apache.org/repos/asf/yunikorn-core.git
commit 922bc77e92cbdee7c3ea086912fc6e8344e44b50 Author: Ompal singh <[email protected]> AuthorDate: Wed Jun 10 12:53:48 2026 +1000 [YUNIKORN-2249] Add compression to REST API (#1090) Add compression to REST API and compress output when larger than a single frame. Based on MTU of 1500, compress data when it exceeds 1400 bytes. Closes: #1090 Signed-off-by: Wilfred Spiegelenburg <[email protected]> (cherry picked from commit 77b42be35b1b0cf8dfce5c364e0fbf1b22209290) --- pkg/webservice/gzip.go | 176 +++++++++++++++++++++++++++++++ pkg/webservice/webservice.go | 2 +- pkg/webservice/webservice_test.go | 216 ++++++++++++++++++++++++++++++++++++++ 3 files changed, 393 insertions(+), 1 deletion(-) diff --git a/pkg/webservice/gzip.go b/pkg/webservice/gzip.go new file mode 100644 index 00000000..6fc5debd --- /dev/null +++ b/pkg/webservice/gzip.go @@ -0,0 +1,176 @@ +/* + Licensed to the Apache Software Foundation (ASF) under one + or more contributor license agreements. See the NOTICE file + distributed with this work for additional information + regarding copyright ownership. The ASF licenses this file + to you under the Apache License, Version 2.0 (the + "License"); you may not use this file except in compliance + with the License. You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. +*/ + +package webservice + +import ( + "bytes" + "compress/gzip" + "net/http" + "strings" + + "go.uber.org/zap" + + "github.com/apache/yunikorn-core/pkg/log" +) + +// minCompressionSize is the minimum response body size in bytes required before gzip +// compression is applied. Responses smaller than this threshold are sent uncompressed +// because the gzip framing overhead would exceed any size savings. The value is set +// below the typical Ethernet MTU (1500 bytes) to account for TCP/IP and HTTP header +// overhead, ensuring a small response body fits within a single network packet. +var minCompressionSize = 1400 + +// deferredGzipResponseWriter buffers the first minCompressionSize bytes of the +// response body and defers the compression decision until either the buffer is full +// (switch to gzip streaming) or the handler returns (send raw if still under threshold). +// This avoids the memory cost of buffering the entire response while still skipping +// gzip for small payloads where the overhead outweighs the savings. +type deferredGzipResponseWriter struct { + http.ResponseWriter + buf bytes.Buffer + gz *gzip.Writer + statusCode int + decided bool + useGzip bool +} + +// WriteHeader captures the status code until the compression decision is made, then +// forwards it to the underlying ResponseWriter. +func (d *deferredGzipResponseWriter) WriteHeader(code int) { + if d.decided { + d.ResponseWriter.WriteHeader(code) + return + } + d.statusCode = code +} + +// Write buffers incoming bytes until the threshold is reached, at which point it +// commits to gzip streaming. After the decision is made, writes go directly to the +// chosen writer. +func (d *deferredGzipResponseWriter) Write(b []byte) (int, error) { + if d.decided { + if d.useGzip { + return d.gz.Write(b) + } + return d.ResponseWriter.Write(b) + } + + n, err := d.buf.Write(b) + if d.buf.Len() >= minCompressionSize { + d.switchToGzip() + } + return n, err +} + +// switchToGzip commits to gzip encoding: sets response headers, flushes the buffered +// bytes through the gzip writer, and marks the decision as final. +func (d *deferredGzipResponseWriter) switchToGzip() { + d.decided = true + d.useGzip = true + d.ResponseWriter.Header().Set("Content-Encoding", "gzip") + d.ResponseWriter.Header().Add("Vary", "Accept-Encoding") + if d.statusCode != 0 { + d.ResponseWriter.WriteHeader(d.statusCode) + } + if _, err := d.gz.Write(d.buf.Bytes()); err != nil { + log.Log(log.REST).Error("failed to write buffered bytes to gzip writer", + zap.Error(err)) + } + d.buf.Reset() +} + +// finalize is called via defer after the handler returns. If no compression decision +// was made yet (response stayed below threshold), the buffered bytes are written raw. +// The gzip writer is closed if compression was used. +func (d *deferredGzipResponseWriter) finalize() { + if !d.decided { + d.decided = true + if d.statusCode != 0 { + d.ResponseWriter.WriteHeader(d.statusCode) + } + if _, err := d.ResponseWriter.Write(d.buf.Bytes()); err != nil { + log.Log(log.REST).Error("failed to write buffered response bytes", + zap.Error(err)) + } + } + if d.useGzip { + if closeErr := d.gz.Close(); closeErr != nil { + log.Log(log.REST).Warn("failed to flush and close gzip writer", + zap.Error(closeErr)) + } + } +} + +// compressResponse is an HTTP middleware that compresses the response body using +// gzip when the client declares gzip support in the Accept-Encoding request header +// (per RFC 9110 §12.5.3). Compression is only applied when the response body exceeds +// minCompressionSize bytes; smaller responses are sent uncompressed to avoid the +// gzip framing overhead exceeding the size savings. +// +// Responses are always served uncompressed when the client has not requested gzip, +// ensuring full backwards compatibility. +// +// The event-stream endpoint (/ws/v1/events/stream) is excluded because it uses +// server-sent events with http.Flusher, which requires writing directly to the +// underlying connection. +func compressResponse(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/ws/v1/events/stream" || !clientAcceptsGzip(r) { + next.ServeHTTP(w, r) + return + } + + gz, err := gzip.NewWriterLevel(w, gzip.BestSpeed) + if err != nil { + // gzip.BestSpeed is always a valid level, so this branch is unreachable in + // practice. Fall back to an uncompressed response rather than failing the call. + log.Log(log.REST).Error("failed to create gzip writer, sending uncompressed response", + zap.Error(err)) + next.ServeHTTP(w, r) + return + } + + dw := &deferredGzipResponseWriter{ + ResponseWriter: w, + gz: gz, + } + defer dw.finalize() + next.ServeHTTP(dw, r) + }) +} + +// clientAcceptsGzip reports whether the request's Accept-Encoding header lists gzip +// as an acceptable content encoding. The q=0 case (gzip explicitly excluded) is +// handled so that "gzip;q=0" correctly returns false. +func clientAcceptsGzip(r *http.Request) bool { + for _, token := range strings.Split(r.Header.Get("Accept-Encoding"), ",") { + token = strings.TrimSpace(token) + parts := strings.SplitN(token, ";", 2) + if strings.ToLower(strings.TrimSpace(parts[0])) != "gzip" { + continue + } + // "gzip" with no q-value means q=1.0 — acceptable. + if len(parts) == 1 { + return true + } + // "gzip;q=0" means explicitly not acceptable; any other q-value is acceptable. + return strings.TrimSpace(parts[1]) != "q=0" + } + return false +} diff --git a/pkg/webservice/webservice.go b/pkg/webservice/webservice.go index 57587630..6c53c7fa 100644 --- a/pkg/webservice/webservice.go +++ b/pkg/webservice/webservice.go @@ -67,7 +67,7 @@ func (m *WebService) StartWebApp() { router := newRouter() m.httpServer = &http.Server{ Addr: ":9080", - Handler: router, + Handler: compressResponse(router), ReadHeaderTimeout: 10 * time.Second, } diff --git a/pkg/webservice/webservice_test.go b/pkg/webservice/webservice_test.go index 868df959..128023e6 100644 --- a/pkg/webservice/webservice_test.go +++ b/pkg/webservice/webservice_test.go @@ -19,9 +19,12 @@ package webservice import ( + "bytes" + "compress/gzip" "fmt" "io" "net/http" + "strings" "testing" "gotest.tools/v3/assert" @@ -151,3 +154,216 @@ func Test_HeaderChecks(t *testing.T) { }) } } + +// newTestClient returns an HTTP client with automatic gzip decompression disabled, +// so that tests can inspect raw response headers and body bytes. +func newTestClient() *http.Client { + return &http.Client{ + Transport: &http.Transport{DisableCompression: true}, + } +} + +func Test_GzipCompression(t *testing.T) { + // Lower the threshold to 0 so that even the small responses returned by the + // empty-context endpoints are compressed. This test focuses on the Accept-Encoding + // negotiation logic, not the threshold behaviour (see Test_GzipMinCompressionSize). + orig := minCompressionSize + minCompressionSize = 0 + t.Cleanup(func() { minCompressionSize = orig }) + + s := NewWebApp(&scheduler.ClusterContext{}, nil) + s.StartWebApp() + defer func() { + if err := s.StopWebApp(); err != nil { + t.Fatal("failed to stop webapp") + } + }() + + client := newTestClient() + + tests := []struct { + name string + acceptEncoding string + wantGzip bool + }{ + { + name: "gzip requested — response must be gzip-encoded", + acceptEncoding: "gzip", + wantGzip: true, + }, + { + name: "gzip with quality factor — response must be gzip-encoded", + acceptEncoding: "gzip;q=0.9", + wantGzip: true, + }, + { + name: "gzip explicitly excluded — response must be uncompressed", + acceptEncoding: "gzip;q=0", + wantGzip: false, + }, + { + name: "no Accept-Encoding — response must be uncompressed", + acceptEncoding: "", + wantGzip: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + req, err := http.NewRequest(http.MethodGet, base+"/ws/v1/clusters", nil) + assert.NilError(t, err, "unexpected error creating request") + if tt.acceptEncoding != "" { + req.Header.Set("Accept-Encoding", tt.acceptEncoding) + } + + //nolint:gosec + resp, err := client.Do(req) + assert.NilError(t, err, "unexpected error executing request") + assert.Equal(t, resp.StatusCode, http.StatusOK, "expected 200 OK") + + body, err := io.ReadAll(resp.Body) + _ = resp.Body.Close() + assert.NilError(t, err, "unexpected error reading response body") + + if tt.wantGzip { + assert.Equal(t, resp.Header.Get("Content-Encoding"), "gzip", + "expected Content-Encoding: gzip in response headers") + assert.Equal(t, resp.Header.Get("Vary"), "Accept-Encoding", + "expected Vary: Accept-Encoding in response headers") + + gr, gzErr := gzip.NewReader(bytes.NewReader(body)) + assert.NilError(t, gzErr, "expected body to be valid gzip data") + decompressed, readErr := io.ReadAll(gr) + assert.NilError(t, readErr, "unexpected error decompressing body") + _ = gr.Close() + assert.Assert(t, len(decompressed) > 0, "decompressed body must not be empty") + } else { + assert.Equal(t, resp.Header.Get("Content-Encoding"), "", + "expected no Content-Encoding header for uncompressed response") + assert.Assert(t, len(body) > 0, "response body must not be empty") + } + }) + } +} + +func Test_GzipMinCompressionSize(t *testing.T) { + s := NewWebApp(&scheduler.ClusterContext{}, nil) + s.StartWebApp() + defer func() { + if err := s.StopWebApp(); err != nil { + t.Fatal("failed to stop webapp") + } + }() + + client := newTestClient() + + // With the default threshold (1500 bytes), the empty-context /ws/v1/clusters + // endpoint returns a small JSON payload well below the threshold. Even when the + // client requests gzip, the response must be sent uncompressed because compression + // would add more bytes than it saves. + req, err := http.NewRequest(http.MethodGet, base+"/ws/v1/clusters", nil) + assert.NilError(t, err, "unexpected error creating request") + req.Header.Set("Accept-Encoding", "gzip") + + //nolint:gosec + resp, err := client.Do(req) + assert.NilError(t, err, "unexpected error executing request") + assert.Equal(t, resp.StatusCode, http.StatusOK, "expected 200 OK") + + body, err := io.ReadAll(resp.Body) + _ = resp.Body.Close() + assert.NilError(t, err, "unexpected error reading response body") + + assert.Assert(t, len(body) < minCompressionSize, + "test pre-condition: response must be smaller than minCompressionSize") + assert.Equal(t, resp.Header.Get("Content-Encoding"), "", + "small response must not be gzip-encoded even when gzip is requested") + assert.Equal(t, resp.Header.Get("Vary"), "", + "Vary header must not be set when response is not compressed") +} + +func Test_GzipClientAcceptsGzip(t *testing.T) { + tests := []struct { + name string + acceptEncoding string + want bool + }{ + {"empty header", "", false}, + {"gzip only", "gzip", true}, + {"gzip with whitespace", " gzip ", true}, + {"gzip, deflate", "gzip, deflate", true}, + {"deflate only", "deflate", false}, + {"gzip;q=0.5", "gzip;q=0.5", true}, + {"gzip;q=0 (excluded)", "gzip;q=0", false}, + {"uppercase GZIP", "GZIP", true}, + {"identity only", "identity", false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + req, err := http.NewRequest(http.MethodGet, "/", nil) + assert.NilError(t, err) + if tt.acceptEncoding != "" { + req.Header.Set("Accept-Encoding", tt.acceptEncoding) + } + got := clientAcceptsGzip(req) + assert.Equal(t, got, tt.want, "clientAcceptsGzip(%q)", tt.acceptEncoding) + }) + } +} + +func Test_GzipExcludesEventStream(t *testing.T) { + s := NewWebApp(&scheduler.ClusterContext{}, nil) + s.StartWebApp() + defer func() { + if err := s.StopWebApp(); err != nil { + t.Fatal("failed to stop webapp") + } + }() + + client := newTestClient() + + // The event-stream endpoint must never receive a gzip Content-Encoding header, + // even when the client sends Accept-Encoding: gzip, because it uses SSE / Flusher. + req, err := http.NewRequest(http.MethodGet, base+"/ws/v1/events/stream", nil) + assert.NilError(t, err, "unexpected error creating request") + req.Header.Set("Accept-Encoding", "gzip") + + //nolint:gosec + resp, err := client.Do(req) + assert.NilError(t, err, "unexpected error executing request") + _ = resp.Body.Close() + + assert.Equal(t, resp.Header.Get("Content-Encoding"), "", + "event-stream endpoint must not be gzip-encoded") +} + +func Test_GzipVaryHeaderNotDuplicated(t *testing.T) { + orig := minCompressionSize + minCompressionSize = 0 + t.Cleanup(func() { minCompressionSize = orig }) + + s := NewWebApp(&scheduler.ClusterContext{}, nil) + s.StartWebApp() + defer func() { + if err := s.StopWebApp(); err != nil { + t.Fatal("failed to stop webapp") + } + }() + + client := newTestClient() + + req, err := http.NewRequest(http.MethodGet, base+"/ws/v1/clusters", nil) + assert.NilError(t, err) + req.Header.Set("Accept-Encoding", "gzip") + + //nolint:gosec + resp, err := client.Do(req) + assert.NilError(t, err, "unexpected error executing request") + _ = resp.Body.Close() + + vary := resp.Header["Vary"] + // Vary header should be present exactly once with value "Accept-Encoding". + assert.Equal(t, len(vary), 1, "expected exactly one Vary header value") + assert.Equal(t, strings.TrimSpace(vary[0]), "Accept-Encoding", "unexpected Vary header value") +} --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
