Skip to content

Commit

Permalink
unique: don't retain uncloned input as key
Browse files Browse the repository at this point in the history
Currently the unique package tries to clone strings that get stored in
its internal map to avoid retaining large strings.

However, this falls over entirely due to the fact that the original
string is *still* stored in the map as a key. Whoops. Fix this by
storing the cloned value in the map instead.

This change also adds a test which fails without this change.

Change-Id: I1a6bb68ed79b869ea12ab6be061a5ae4b4377ddb
Reviewed-on: https://go-review.googlesource.com/c/go/+/610738
Reviewed-by: Michael Pratt <mpratt@google.com>
Auto-Submit: Michael Knyszek <mknyszek@google.com>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
  • Loading branch information
mknyszek authored and gopherbot committed Sep 4, 2024
1 parent 79fd633 commit 21ac23a
Show file tree
Hide file tree
Showing 2 changed files with 26 additions and 3 deletions.
7 changes: 4 additions & 3 deletions src/unique/handle.go
Original file line number Diff line number Diff line change
Expand Up @@ -50,21 +50,22 @@ func Make[T comparable](value T) Handle[T] {
toInsert *T // Keep this around to keep it alive.
toInsertWeak weak.Pointer[T]
)
newValue := func() weak.Pointer[T] {
newValue := func() (T, weak.Pointer[T]) {
if toInsert == nil {
toInsert = new(T)
*toInsert = clone(value, &m.cloneSeq)
toInsertWeak = weak.Make(toInsert)
}
return toInsertWeak
return *toInsert, toInsertWeak
}
var ptr *T
for {
// Check the map.
wp, ok := m.Load(value)
if !ok {
// Try to insert a new value into the map.
wp, _ = m.LoadOrStore(value, newValue())
k, v := newValue()
wp, _ = m.LoadOrStore(k, v)
}
// Now that we're sure there's a value in the map, let's
// try to get the pointer we need out of it.
Expand Down
22 changes: 22 additions & 0 deletions src/unique/handle_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,10 @@ import (
"internal/abi"
"reflect"
"runtime"
"strings"
"testing"
"time"
"unsafe"
)

// Set up special types. Because the internal maps are sharded by type,
Expand Down Expand Up @@ -110,3 +113,22 @@ func checkMapsFor[T comparable](t *testing.T, value T) {
}
t.Errorf("failed to drain internal maps of %v", value)
}

func TestMakeClonesStrings(t *testing.T) {
s := strings.Clone("abcdefghijklmnopqrstuvwxyz") // N.B. Must be big enough to not be tiny-allocated.
ran := make(chan bool)
runtime.SetFinalizer(unsafe.StringData(s), func(_ *byte) {
ran <- true
})
h := Make(s)

// Clean up s (hopefully) and run the finalizer.
runtime.GC()

select {
case <-time.After(1 * time.Second):
t.Fatal("string was improperly retained")
case <-ran:
}
runtime.KeepAlive(h)
}

0 comments on commit 21ac23a

Please sign in to comment.