unique.patch

diff --git a/src/unique/handle.go b/src/unique/handle.go
index 520ab70f8c..113eb90fd9 100644
--- a/src/unique/handle.go
+++ b/src/unique/handle.go
@@ -9,6 +9,8 @@ import (
 	isync "internal/sync"
 	"runtime"
 	"sync"
+	"sync/atomic"
+	"time"
 	"unsafe"
 	"weak"
 )
@@ -48,6 +50,10 @@ func Make[T comparable](value T) Handle[T] {
 	}
 	m := ma.(*uniqueMap[T])
 
+	for removesInFlight.Load() > 0 {
+		time.Sleep(1 * time.Millisecond)
+	}
+
 	// Keep around any values we allocate for insertion. There
 	// are a few different ways we can race with other threads
 	// and create values that we might discard. By keeping
@@ -84,7 +90,7 @@ func Make[T comparable](value T) Handle[T] {
 		// Try to remove it and start over.
 		m.CompareAndDelete(value, wp)
 	}
-	runtime.KeepAlive(toInsert)
+	runtime.AddCleanup(toInsert, m.removeKeyIfWeak, value)
 	return Handle[T]{ptr}
 }
 
@@ -126,23 +132,36 @@ func addUniqueMap[T comparable](typ *abi.Type) *uniqueMap[T] {
 	m := &uniqueMap[T]{cloneSeq: makeCloneSeq(typ)}
 	a, loaded := uniqueMaps.LoadOrStore(typ, m)
 	if !loaded {
-		// Add a cleanup function for the new map.
 		cleanupFuncsMu.Lock()
 		cleanupFuncs = append(cleanupFuncs, func() {
-			// Delete all the entries whose weak references are nil and clean up
-			// deleted entries.
-			m.All()(func(key T, wp weak.Pointer[T]) bool {
-				if wp.Value() == nil {
-					m.CompareAndDelete(key, wp)
-				}
-				return true
-			})
+			println("XXX: unique cleanup: in-flight=", removesInFlight.Load(), "; deleted=", uniqueDeleted.Load())
 		})
 		cleanupFuncsMu.Unlock()
 	}
 	return a.(*uniqueMap[T])
 }
 
+var removesInFlight, uniqueDeleted atomic.Int64
+
+func (m *uniqueMap[T]) removeKeyIfWeak(key T) {
+	removesInFlight.Add(1)
+	defer removesInFlight.Add(-1)
+	wp, ok := m.Load(key)
+	if !ok {
+		println("XXX: unique failed to Load key")
+		return
+	}
+	if wp.Value() != nil {
+		println("XXX: unique delet not weak")
+		return
+	}
+	if m.CompareAndDelete(key, wp) {
+		uniqueDeleted.Add(1)
+	} else {
+		println("XXX: unique failed to delete key")
+	}
+}
+
 // setupMake is used to perform initial setup for unique.Make.
 var setupMake sync.Once