Fix data race in BucketedBytes pool

CHANGELOG.md

@@ -24,6 +24,7 @@ We use *breaking :warning:* to mark changes that are not backward compatible (re
 - [#4663](https://github.com/thanos-io/thanos/pull/4663) Fetcher: Fix discovered data races.
 - [#4754](https://github.com/thanos-io/thanos/pull/4754) Query: Fix possible panic on stores endpoint.
 - [#4753](https://github.com/thanos-io/thanos/pull/4753) Store: validate block sync concurrency parameter
+- [#4792](https://github.com/thanos-io/thanos/pull/4792) Store: Fix data race in BucketedBytes pool.
 
 ## [v0.23.1](https://github.com/thanos-io/thanos/tree/release-0.23) - 2021.10.1
 

pkg/pool/pool.go

@@ -107,8 +107,9 @@ func (p *BucketedBytes) Put(b *[]byte) {
 		return
 	}
 
+	sz := cap(*b)
 	for i, bktSize := range p.sizes {
-		if cap(*b) > bktSize {
+		if sz > bktSize {
 			continue
 		}
 		*b = (*b)[:0]
@@ -118,13 +119,11 @@ func (p *BucketedBytes) Put(b *[]byte) {
 
 	p.mtx.Lock()
 	defer p.mtx.Unlock()
-
 	// We could assume here that our users will not make the slices larger
 	// but lets be on the safe side to avoid an underflow of p.usedTotal.
-	sz := uint64(cap(*b))
-	if sz >= p.usedTotal {
+	if uint64(sz) >= p.usedTotal {
 		p.usedTotal = 0
 	} else {
-		p.usedTotal -= sz
+		p.usedTotal -= uint64(sz)
 	}
 }

pkg/pool/pool_test.go

@@ -4,8 +4,7 @@
 package pool
 
 import (
-	"bytes"
-	"fmt"
+	"strings"
 	"sync"
 	"testing"
 	"time"
@@ -71,52 +70,57 @@ func TestRacePutGet(t *testing.T) {
 
 	s := sync.WaitGroup{}
 
-	// Start two goroutines: they always Get and Put two byte slices
-	// to which they write 'foo' / 'barbazbaz' and check if the data is still
+	const goroutines = 100
+
+	// Start multiple goroutines: they always Get and Put two byte slices
+	// to which they write their contents and check if the data is still
 	// there after writing it, before putting it back.
-	errs := make(chan error, 2)
-	stop := make(chan bool, 2)
+	errs := make(chan error, goroutines)
+	stop := make(chan struct{})
 
-	f := func(txt string) {
+	f := func(txt string, grow bool) {
 		defer s.Done()
 		for {
 			select {
 			case <-stop:
 				return
 			default:
-				c, err := chunkPool.Get(3)
-				if err != nil {
-					errs <- errors.Wrapf(err, "goroutine %s", txt)
-					return
-				}
-
-				buf := bytes.NewBuffer(*c)
-
-				_, err = fmt.Fprintf(buf, "%s", txt)
+				c, err := chunkPool.Get(len(txt))
 				if err != nil {
 					errs <- errors.Wrapf(err, "goroutine %s", txt)
 					return
 				}
 
-				if buf.String() != txt {
+				*c = append(*c, txt...)
+				if string(*c) != txt {
 					errs <- errors.New("expected to get the data just written")
 					return
 				}
+				if grow {
+					*c = append(*c, txt...)
+					*c = append(*c, txt...)
+					if string(*c) != txt+txt+txt {
+						errs <- errors.New("expected to get the data just written")
+						return
+					}
+				}
 
-				b := buf.Bytes()
-				chunkPool.Put(&b)
+				chunkPool.Put(c)
 			}
 		}
 	}
 
-	s.Add(2)
-	go f("foo")
-	go f("barbazbaz")
-
-	time.Sleep(5 * time.Second)
-	stop <- true
-	stop <- true
+	for i := 0; i < goroutines; i++ {
+		s.Add(1)
+		// make sure we start multiple goroutines with same len buf requirements, to hit same pools
+		s := strings.Repeat(string(byte(i)), i%10)
+		// some of the goroutines will append more elements to the provided slice
+		grow := i%2 == 0
+		go f(s, grow)
+	}
 
+	time.Sleep(1 * time.Second)
+	close(stop)
 	s.Wait()
 	select {
 	case err := <-errs: