Add mechanism to detect leaked references to mimirpb buffers #13609
+408
−23
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -6,27 +6,57 @@ import ( | |
| "bytes" | ||
| "fmt" | ||
| "math" | ||
| "syscall" | ||
| "testing" | ||
| "unsafe" | ||
|
|
||
| gogoproto "github.com/gogo/protobuf/proto" | ||
| "github.com/prometheus/prometheus/model/histogram" | ||
| "go.uber.org/atomic" | ||
| "google.golang.org/grpc/encoding" | ||
| "google.golang.org/grpc/encoding/proto" | ||
| "google.golang.org/grpc/mem" | ||
| protobufproto "google.golang.org/protobuf/proto" | ||
| "google.golang.org/protobuf/protoadapt" | ||
| ) | ||
|
|
||
| type CustomCodecConfig struct { | ||
| InstrumentRefLeaksPct float64 | ||
| } | ||
|
|
||
| var baseCodecV2Name = encoding.GetCodecV2(proto.Name).Name() | ||
|
|
||
| func (cfg CustomCodecConfig) codec() *codecV2 { | ||
| c := &codecV2{} | ||
| if cfg.InstrumentRefLeaksPct > 0 { | ||
| c.instrumentRefLeaksOneIn = uint64(math.Trunc(100 / cfg.InstrumentRefLeaksPct)) | ||
tcard marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| } | ||
| return c | ||
| } | ||
tcard marked this conversation as resolved.
Show resolved
Hide resolved
|
||
|
|
||
| var globalCodec encoding.CodecV2 | ||
|
|
||
| func (cfg CustomCodecConfig) RegisterGlobally() { | ||
| globalCodec = cfg.codec() | ||
| encoding.RegisterCodecV2(globalCodec) | ||
| } | ||
|
|
||
| func init() { | ||
| config := CustomCodecConfig{} | ||
| if testing.Testing() { | ||
| // Instrument all buffers when testing. | ||
| config.InstrumentRefLeaksPct = 100 | ||
| } | ||
| config.RegisterGlobally() | ||
| } | ||
|
|
||
| // codecV2 customizes gRPC marshalling and unmarshalling. | ||
| // We customize marshalling in order to use optimized paths when possible. | ||
| // We customize unmarshalling in order to use an optimized path when possible, | ||
| // and to retain the unmarshalling buffer when necessary. | ||
| type codecV2 struct { | ||
| instrumentRefLeaksOneIn uint64 | ||
| unmarshaledWithBufferRefCount atomic.Uint64 | ||
| } | ||
|
|
||
| var _ encoding.CodecV2 = &codecV2{} | ||
|
|
@@ -113,11 +143,47 @@ func unmarshalSlicePoolSizes() []int { | |
| return sizes | ||
| } | ||
|
|
||
| // Unmarshal unmarshals an object using the global codec. Prefer this over | ||
| // calling the Unmarshal method directly, as it will take advantage of leak | ||
| // detection. | ||
| func Unmarshal(data []byte, v gogoproto.Unmarshaler) error { | ||
| return globalCodec.Unmarshal(mem.BufferSlice{mem.SliceBuffer(data)}, v) | ||
| } | ||
|
|
||
| var pageSize = syscall.Getpagesize() | ||
|
|
||
| // Unmarshal customizes gRPC unmarshalling. | ||
| // If v implements MessageWithBufferRef, its SetBuffer method is called with the unmarshalling buffer and the buffer's reference count gets incremented. | ||
| // The latter means that v's FreeBuffer method should be called when v is no longer used. | ||
| func (c *codecV2) Unmarshal(data mem.BufferSlice, v any) error { | ||
| holder, isBufferHolder := v.(MessageWithBufferRef) | ||
| instrumentLeaks := data.Len() > 0 && isBufferHolder && c.instrumentRefLeaksOneIn > 0 && c.unmarshaledWithBufferRefCount.Add(1)%c.instrumentRefLeaksOneIn == 0 | ||
|
|
||
| var buf mem.Buffer | ||
| if instrumentLeaks { | ||
| // Allocate separate pages for this buffer. We'll detect ref leaks by | ||
| // munmaping the pages on Free, after which trying to access them will | ||
| // segfault. | ||
| pageAlignedLen := roundUpToMultiple(data.Len(), pageSize) | ||
| b, err := syscall.Mmap(0, 0, pageAlignedLen, syscall.PROT_READ|syscall.PROT_WRITE, syscall.MAP_PRIVATE|syscall.MAP_ANON) | ||
| if err != nil { | ||
| panic(fmt.Errorf("mmap: %w", err)) | ||
| } | ||
| b = b[:data.Len()] | ||
| data.CopyTo(b) | ||
| buf = mem.NewBuffer(&b, nil) | ||
| instrumentedBuf := &instrumentLeaksBuf{Buffer: buf} | ||
| instrumentedBuf.refCount.Inc() | ||
| buf = instrumentedBuf | ||
| } else if len(data) == 1 { | ||
| // BufferSlice.MaterializeToBuffer already has this behavior when | ||
| // len(data) == 1, but we reproduce it here for explicitness and for | ||
| // ensuring forward-compatibility. | ||
| data[0].Ref() | ||
| buf = data[0] | ||
| } else { | ||
| buf = data.MaterializeToBuffer(unmarshalSlicePool) | ||
| } | ||
| // Decrement buf's reference count. Note though that if v implements MessageWithBufferRef, | ||
| // we increase buf's reference count first so it doesn't go to zero. | ||
| defer buf.Free() | ||
|
|
@@ -133,7 +199,7 @@ func (c *codecV2) Unmarshal(data mem.BufferSlice, v any) error { | |
| } | ||
| } | ||
|
|
||
| if isBufferHolder { | ||
| buf.Ref() | ||
| holder.SetBuffer(buf) | ||
| } | ||
|
|
@@ -142,7 +208,7 @@ func (c *codecV2) Unmarshal(data mem.BufferSlice, v any) error { | |
| } | ||
|
|
||
| func (c *codecV2) Name() string { | ||
| return baseCodecV2Name | ||
| } | ||
|
|
||
| // MessageWithBufferRef is an unmarshalling buffer retaining protobuf message. | ||
|
|
@@ -158,6 +224,10 @@ type BufferHolder struct { | |
| buffer mem.Buffer | ||
| } | ||
|
|
||
| func (m *BufferHolder) Buffer() mem.Buffer { | ||
| return m.buffer | ||
| } | ||
|
|
||
| func (m *BufferHolder) SetBuffer(buf mem.Buffer) { | ||
| m.buffer = buf | ||
| } | ||
|
|
@@ -171,6 +241,30 @@ func (m *BufferHolder) FreeBuffer() { | |
|
|
||
| var _ MessageWithBufferRef = &BufferHolder{} | ||
|
|
||
| type instrumentLeaksBuf struct { | ||
| mem.Buffer | ||
| refCount atomic.Int64 | ||
| } | ||
|
|
||
| func (b *instrumentLeaksBuf) Ref() { | ||
| b.Buffer.Ref() | ||
| b.refCount.Inc() | ||
| } | ||
|
|
||
| func (b *instrumentLeaksBuf) Free() { | ||
| b.Buffer.Free() | ||
|
|
||
| if b.refCount.Dec() == 0 { | ||
| buf := b.ReadOnlyData() | ||
| ptr := unsafe.SliceData(buf) | ||
| allPages := unsafe.Slice(ptr, roundUpToMultiple(len(buf), pageSize)) | ||
| err := syscall.Munmap(allPages) | ||
|
||
| if err != nil { | ||
| panic(fmt.Errorf("munmap: %w", err)) | ||
| } | ||
| } | ||
| } | ||
tcard marked this conversation as resolved.
Show resolved
Hide resolved
|
||
|
|
||
| // MinTimestamp returns the minimum timestamp (milliseconds) among all series | ||
| // in the WriteRequest. Returns math.MaxInt64 if the request is empty. | ||
| func (m *WriteRequest) MinTimestamp() int64 { | ||
|
|
@@ -534,3 +628,7 @@ type orderAwareMetricMetadata struct { | |
| // order is the 0-based index of this metadata object in a wider metadata array. | ||
| order int | ||
| } | ||
|
|
||
| func roundUpToMultiple(n, of int) int { | ||
| return ((n + of - 1) / of) * of | ||
| } | ||
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.