-
Notifications
You must be signed in to change notification settings - Fork 539
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Preallocate byte slices from prom pool (#679)
* Preallocate byte slices from prom pool Signed-off-by: Annanay <annanayagarwal@gmail.com> * make fmt Signed-off-by: Annanay <annanayagarwal@gmail.com> * Add tests around prealloc Signed-off-by: Annanay <annanayagarwal@gmail.com> * Add CHANGELOG, comment out test and add reason Signed-off-by: Annanay <annanayagarwal@gmail.com> * Adjust pool bucket sizes Signed-off-by: Annanay <annanayagarwal@gmail.com>
- Loading branch information
Showing
10 changed files
with
232 additions
and
54 deletions.
There are no files selected for viewing
This file contains 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 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 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 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 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 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 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 |
|---|---|---|
| @@ -0,0 +1,46 @@ | ||
| package tempopb | ||
|
|
||
| import ( | ||
| "github.com/prometheus/prometheus/pkg/pool" | ||
| ) | ||
|
|
||
| var ( | ||
| // buckets: [0.5KiB, 1KiB, 2KiB, 4KiB, 8KiB, 16KiB] | ||
| bytePool = pool.New(500, 16_000, 2, func(size int) interface{} { return make([]byte, 0, size) }) | ||
| ) | ||
|
|
||
| // PreallocRequest is a (repeated bytes requests) which preallocs slices on Unmarshal. | ||
| type PreallocRequest struct { | ||
| Request []byte | ||
| } | ||
|
|
||
| // Unmarshal implements proto.Message. | ||
| func (r *PreallocRequest) Unmarshal(dAtA []byte) error { | ||
| r.Request = bytePool.Get(len(dAtA)).([]byte) | ||
| r.Request = r.Request[:len(dAtA)] | ||
| copy(r.Request, dAtA) | ||
| return nil | ||
| } | ||
|
|
||
| // MarshalTo implements proto.Marshaller. | ||
| // returned int is not used | ||
| func (r *PreallocRequest) MarshalTo(dAtA []byte) (int, error) { | ||
| copy(dAtA[:], r.Request[:]) | ||
| return len(r.Request), nil | ||
| } | ||
|
|
||
| // Size implements proto.Sizer. | ||
| func (r *PreallocRequest) Size() (n int) { | ||
| if r == nil { | ||
| return 0 | ||
| } | ||
| return len(r.Request) | ||
| } | ||
|
|
||
| // ReuseRequest puts the byte slice back into bytePool for reuse. | ||
| func ReuseRequest(req *PushBytesRequest) { | ||
| for i := range req.Requests { | ||
| // We want to preserve the underlying allocated memory, [:0] helps us retains the cap() of the slice | ||
| bytePool.Put(req.Requests[i].Request[:0]) | ||
| } | ||
| } |
This file contains 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 |
|---|---|---|
| @@ -0,0 +1,116 @@ | ||
| package tempopb | ||
|
|
||
| import ( | ||
| "math/rand" | ||
| "testing" | ||
|
|
||
| "github.com/stretchr/testify/assert" | ||
| ) | ||
|
|
||
| func TestUnmarshal(t *testing.T) { | ||
| var dummyData = make([]byte, 10) | ||
| rand.Read(dummyData) | ||
|
|
||
| preallocReq := &PreallocRequest{} | ||
| err := preallocReq.Unmarshal(dummyData) | ||
| assert.NoError(t, err) | ||
|
|
||
| assert.Equal(t, dummyData, preallocReq.Request) | ||
| } | ||
|
|
||
| func TestMarshal(t *testing.T) { | ||
| preallocReq := &PreallocRequest{ | ||
| Request: make([]byte, 10), | ||
| } | ||
| rand.Read(preallocReq.Request) | ||
|
|
||
| var dummyData = make([]byte, 10) | ||
| _, err := preallocReq.MarshalTo(dummyData) | ||
| assert.NoError(t, err) | ||
|
|
||
| assert.Equal(t, preallocReq.Request, dummyData) | ||
| } | ||
|
|
||
| func TestSize(t *testing.T) { | ||
| preallocReq := &PreallocRequest{ | ||
| Request: make([]byte, 10), | ||
| } | ||
| assert.Equal(t, 10, preallocReq.Size()) | ||
| } | ||
|
|
||
| /* The prometheus pool pkg is a wrapper around sync.Pool | ||
| From the comments on sync.Pool pkg: | ||
| // Get selects an arbitrary item from the Pool, removes it from the | ||
| // Pool, and returns it to the caller. | ||
| // Get may choose to ignore the pool and treat it as empty. | ||
| // Callers should not assume any relation between values passed to Put and | ||
| // the values returned by Get. | ||
| And for those reasons, the test below is rendered flaky. However, it should have little impact in a production environment. | ||
| Commenting it out but retaining as part of the package as an indicator that the logic is tested. | ||
| func TestReuseRequest(t *testing.T) { | ||
| tests := []struct { | ||
| name string | ||
| donate int | ||
| request int | ||
| expectedEqual bool | ||
| }{ | ||
| { | ||
| name: "same size", | ||
| donate: 1500, | ||
| request: 1500, | ||
| expectedEqual: true, | ||
| }, | ||
| { | ||
| name: "larger donate - same bucket", | ||
| donate: 1600, | ||
| request: 1500, | ||
| expectedEqual: true, | ||
| }, | ||
| { | ||
| name: "larger donate - different bucket", | ||
| donate: 2100, | ||
| request: 1500, | ||
| }, | ||
| } | ||
| for _, tt := range tests { | ||
| t.Run(tt.name, func(t *testing.T) { | ||
| // create push requests of known size | ||
| req := MakeBytesRequestWithSize(tt.donate) | ||
| assert.Len(t, req.Requests, 1) | ||
| expectedAddr := &req.Requests[0].Request[0] | ||
| // "donate" to bytePool | ||
| ReuseRequest(req) | ||
| // unmarshal a new request | ||
| var dummyData = make([]byte, tt.request) | ||
| preallocReq := &PreallocRequest{} | ||
| assert.NoError(t, preallocReq.Unmarshal(dummyData)) | ||
| actualAddr := &preallocReq.Request[0] | ||
| if tt.expectedEqual { | ||
| assert.Equal(t, expectedAddr, actualAddr) | ||
| } else { | ||
| assert.NotEqual(t, expectedAddr, actualAddr) | ||
| } | ||
| }) | ||
| } | ||
| } | ||
| func MakeBytesRequestWithSize(maxBytes int) *PushBytesRequest { | ||
| reqBytes := make([]byte, maxBytes) | ||
| rand.Read(reqBytes) | ||
| return &PushBytesRequest{ | ||
| Requests: []PreallocRequest{ | ||
| { | ||
| Request: reqBytes, | ||
| }, | ||
| }, | ||
| } | ||
| } | ||
| */ |
Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.
Oops, something went wrong.
Oops, something went wrong.