Skip to content

Commit f0078b3

Browse files
authored
Merge branch 'main' into feat/civisibility-trimpath-source-paths
2 parents 8dec13d + 2743f8e commit f0078b3

2 files changed

Lines changed: 112 additions & 64 deletions

File tree

profiler/compression.go

Lines changed: 29 additions & 53 deletions
Original file line numberDiff line numberDiff line change
@@ -184,15 +184,15 @@ func (b *compressionPipelineBuilder) Build(in compression, out compression) (com
184184
if err != nil {
185185
return nil, err
186186
}
187-
return newGzipRecompressor(gzipOut), nil
187+
return newGzipDecodingRecompressor(gzipOut), nil
188188
}
189189

190190
if in.algorithm == compressionAlgorithmGzip && out.algorithm == compressionAlgorithmZstd {
191191
encoder, err := b.getZstdEncoder(getZstdLevelOrDefault(out.level))
192192
if err != nil {
193193
return nil, err
194194
}
195-
return newZstdRecompressor(encoder), nil
195+
return newGzipDecodingRecompressor(encoder), nil
196196
}
197197

198198
return nil, fmt.Errorf("unsupported recompression: %s -> %s", in, out)
@@ -228,78 +228,54 @@ func (r *passthroughCompressor) Close() error {
228228
return nil
229229
}
230230

231-
func newGzipRecompressor(gzipOut *kgzip.Writer) *gzipRecompressor {
232-
return &gzipRecompressor{gzipOut: gzipOut, err: make(chan error)}
231+
// newGzipDecodingRecompressor returns a recompressor that decompresses
232+
// gzip-compressed input and re-encodes it using the supplied output
233+
// compressor (e.g. a gzip writer at a different level, or a zstd encoder).
234+
// Callers must always pair Reset with Close.
235+
func newGzipDecodingRecompressor(out compressor) *gzipDecodingRecompressor {
236+
return &gzipDecodingRecompressor{out: out, err: make(chan error)}
233237
}
234238

235-
type gzipRecompressor struct {
239+
type gzipDecodingRecompressor struct {
236240
// err synchronizes finishing writes after closing pw and reports any
237241
// error during recompression
238-
err chan error
239-
pw io.WriteCloser
240-
gzipOut *kgzip.Writer
242+
err chan error
243+
pw io.WriteCloser
244+
out compressor
241245
}
242246

243-
func (r *gzipRecompressor) Reset(w io.Writer) {
244-
r.gzipOut.Reset(w)
247+
func (r *gzipDecodingRecompressor) Reset(w io.Writer) {
248+
r.out.Reset(w)
245249
pr, pw := io.Pipe()
246250
go func() {
247-
gzr, err := kgzip.NewReader(pr)
248-
if err != nil {
249-
r.err <- err
250-
return
251-
}
252-
_, err = io.Copy(r.gzipOut, gzr)
253-
r.err <- err
254-
}()
255-
r.pw = pw
256-
}
251+
// Always close pr with the resulting error so any pending
252+
// or future pw.Write call returns it instead of blocking on
253+
// a reader that has gone away.
254+
var finalErr error
255+
defer func() {
256+
pr.CloseWithError(finalErr)
257+
r.err <- finalErr
258+
}()
257259

258-
func (r *gzipRecompressor) Write(p []byte) (int, error) {
259-
return r.pw.Write(p)
260-
}
261-
262-
func (r *gzipRecompressor) Close() error {
263-
r.pw.Close()
264-
err := <-r.err
265-
return cmp.Or(err, r.gzipOut.Close())
266-
}
267-
268-
func newZstdRecompressor(encoder *sharedZstdEncoder) *zstdRecompressor {
269-
return &zstdRecompressor{zstdOut: encoder, err: make(chan error)}
270-
}
271-
272-
type zstdRecompressor struct {
273-
// err synchronizes finishing writes after closing pw and reports any
274-
// error during recompression
275-
err chan error
276-
pw io.WriteCloser
277-
zstdOut *sharedZstdEncoder
278-
}
279-
280-
func (r *zstdRecompressor) Reset(w io.Writer) {
281-
r.zstdOut.Reset(w)
282-
pr, pw := io.Pipe()
283-
go func() {
284260
gzr, err := kgzip.NewReader(pr)
285261
if err != nil {
286-
r.err <- err
262+
finalErr = err
287263
return
288264
}
289-
_, err = io.Copy(r.zstdOut, gzr)
290-
r.err <- err
265+
_, copyErr := io.Copy(r.out, gzr)
266+
closeErr := gzr.Close()
267+
finalErr = cmp.Or(copyErr, closeErr)
291268
}()
292269
r.pw = pw
293270
}
294271

295-
func (r *zstdRecompressor) Write(p []byte) (int, error) {
272+
func (r *gzipDecodingRecompressor) Write(p []byte) (int, error) {
296273
return r.pw.Write(p)
297274
}
298275

299-
func (r *zstdRecompressor) Close() error {
276+
func (r *gzipDecodingRecompressor) Close() error {
300277
r.pw.Close()
301-
err := <-r.err
302-
return cmp.Or(err, r.zstdOut.Close())
278+
return cmp.Or(<-r.err, r.out.Close())
303279
}
304280

305281
// newSharedZstdEncoder creates a new shared Zstd encoder with the given level.

profiler/compression_test.go

Lines changed: 83 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -7,11 +7,11 @@ package profiler
77

88
import (
99
"bytes"
10-
"compress/gzip"
1110
"fmt"
1211
"io"
1312
"os"
1413
"testing"
14+
"testing/synctest"
1515
"time"
1616

1717
kgzip "github.com/klauspost/compress/gzip"
@@ -58,7 +58,79 @@ func TestNewCompressionPipeline(t *testing.T) {
5858
}
5959
}
6060

61-
// checkZstdLevel checks that data is zstd-compressed with the given level
61+
// TestGzipDecodingRecompressorInvalidInput verifies that the recompressor
62+
// surfaces an error (instead of deadlocking) when its input is not a valid
63+
// gzip stream. A naive implementation spawns a goroutine that calls
64+
// kgzip.NewReader, which fails on a bad header. Without explicitly closing
65+
// the read end of the pipe, the goroutine exits while the caller's pw.Write
66+
// blocks forever waiting for a reader that no longer exists.
67+
//
68+
// The test runs inside a testing/synctest bubble so that deadlock detection
69+
// is deterministic instead of relying on a wall-clock timeout: synctest.Wait
70+
// returns once every goroutine in the bubble is durably blocked or has
71+
// exited, and we then assert that Write actually returned.
72+
func TestGzipDecodingRecompressorInvalidInput(t *testing.T) {
73+
synctest.Test(t, func(t *testing.T) {
74+
gzipOut, err := kgzip.NewWriterLevel(nil, gzip6Compression.level)
75+
require.NoError(t, err)
76+
r := newGzipDecodingRecompressor(gzipOut)
77+
r.Reset(io.Discard)
78+
// Always close so the synctest bubble drains cleanly even
79+
// when the recompressor is deadlocked: pw.Close unblocks
80+
// any pending Write, and <-r.err unblocks the goroutine
81+
// started by Reset.
82+
defer r.Close()
83+
84+
// Non-gzip data large enough to overflow whatever buffer
85+
// kgzip.NewReader uses internally for header parsing. With
86+
// io.Pipe being unbuffered, any write left over after the
87+
// recompressor's goroutine errors out will block forever
88+
// absent a fix.
89+
data := bytes.Repeat([]byte("not a gzip stream\n"), 4096)
90+
91+
writeDone := make(chan error, 1)
92+
go func() {
93+
_, werr := r.Write(data)
94+
writeDone <- werr
95+
}()
96+
97+
// Block until every goroutine in the bubble is either
98+
// durably blocked or has exited. With the fix in place,
99+
// the recompressor's goroutine closes pr with an error,
100+
// Write returns that error, and writeDone receives a
101+
// value. Without the fix, Write stays blocked on the
102+
// pipe and writeDone is empty.
103+
synctest.Wait()
104+
105+
select {
106+
case werr := <-writeDone:
107+
require.Error(t, werr, "expected an error for non-gzip input")
108+
default:
109+
t.Error("deadlock: Write did not return after recompressor goroutine exited")
110+
}
111+
})
112+
}
113+
114+
// checkGzipLevel checks that data is gzip-compressed with the given level.
115+
func checkGzipLevel(t *testing.T, data []byte, level int) {
116+
t.Helper()
117+
require.NotEmpty(t, data)
118+
gr, err := kgzip.NewReader(bytes.NewReader(data))
119+
require.NoError(t, err)
120+
in := new(bytes.Buffer)
121+
_, err = io.Copy(in, gr)
122+
require.NoError(t, err)
123+
require.NoError(t, gr.Close())
124+
out := new(bytes.Buffer)
125+
gw, err := kgzip.NewWriterLevel(out, level)
126+
require.NoError(t, err)
127+
_, err = io.Copy(gw, in)
128+
require.NoError(t, err)
129+
require.NoError(t, gw.Close())
130+
require.Equal(t, data, out.Bytes())
131+
}
132+
133+
// checkZstdLevel checks that data is zstd-compressed with the given level.
62134
func checkZstdLevel(t *testing.T, data []byte, level zstd.EncoderLevel) {
63135
t.Helper()
64136
require.NotEmpty(t, data)
@@ -95,19 +167,19 @@ func TestDebugCompressionEnv(t *testing.T) {
95167
t.Run("explicit-gzip", func(t *testing.T) {
96168
t.Setenv("DD_PROFILING_DEBUG_COMPRESSION_SETTINGS", "gzip")
97169
p := startTestProfiler(t, 1, WithProfileTypes(HeapProfile, BlockProfile), WithPeriod(time.Millisecond)).ReceiveProfile(t)
98-
r, err := gzip.NewReader(bytes.NewReader(p.attachments["delta-heap.pprof"]))
99-
require.NoError(t, err)
100-
_, err = io.Copy(io.Discard, r)
101-
require.NoError(t, err)
170+
checkGzipLevel(t, p.attachments["delta-heap.pprof"], gzip6Compression.level)
102171
})
103172

104173
t.Run("explicit-gzip-already-gzipped-input", func(t *testing.T) {
105174
t.Setenv("DD_PROFILING_DEBUG_COMPRESSION_SETTINGS", "gzip")
106175
p := startTestProfiler(t, 1, WithProfileTypes(CPUProfile), WithPeriod(time.Millisecond)).ReceiveProfile(t)
107-
r, err := gzip.NewReader(bytes.NewReader(p.attachments["cpu.pprof"]))
108-
require.NoError(t, err)
109-
_, err = io.Copy(io.Discard, r)
110-
require.NoError(t, err)
176+
checkGzipLevel(t, p.attachments["cpu.pprof"], gzip6Compression.level)
177+
})
178+
179+
t.Run("gzip-1", func(t *testing.T) {
180+
t.Setenv("DD_PROFILING_DEBUG_COMPRESSION_SETTINGS", "gzip-1")
181+
p := startTestProfiler(t, 1, WithProfileTypes(HeapProfile), WithPeriod(time.Millisecond)).ReceiveProfile(t)
182+
checkGzipLevel(t, p.attachments["delta-heap.pprof"], gzip1Compression.level)
111183
})
112184

113185
t.Run("zstd-delta", func(t *testing.T) {
@@ -192,7 +264,7 @@ func BenchmarkRecompression(b *testing.B) {
192264
if err != nil {
193265
b.Fatal(err)
194266
}
195-
z := newZstdRecompressor(encoder)
267+
z := newGzipDecodingRecompressor(encoder)
196268
z.Reset(io.Discard)
197269
if _, err := z.Write(data); err != nil {
198270
b.Fatal(err)

0 commit comments

Comments
 (0)