fix: [v1] Backport derived.go GC fix (unsubscribe derived refreshables when unreachable)

refreshable/derived.go

@@ -0,0 +1,82 @@
+// Copyright (c) 2021 Palantir Technologies. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package refreshable
+
+import (
+	"runtime"
+	"sync"
+	"sync/atomic"
+)
+
+// cleanupState tracks the lifecycle of a derivedRefreshable and is responsible
+// for calling unsubscribe functions when the derived is no longer needed.
+// Cleanup fires when the derivedRefreshable is garbage collected AND all
+// user subscriptions on it have been removed.
+type cleanupState struct {
+	subCount atomic.Int32
+	gcDone   atomic.Bool
+	unsubs   []func()
+	once     sync.Once
+}
+
+func (s *cleanupState) tryCleanup() {
+	if s.subCount.Load() > 0 || !s.gcDone.Load() {
+		return
+	}
+	s.once.Do(func() {
+		for _, unsub := range s.unsubs {
+			unsub()
+		}
+	})
+}
+
+// derivedRefreshable wraps an inner Refreshable and manages cleanup of parent
+// subscriptions when the derived is garbage collected. The refs field holds
+// references to upstream objects (e.g. parent derivedRefreshable wrappers) to
+// prevent them from being GC'd prematurely in chained Map() scenarios.
+type derivedRefreshable struct {
+	inner Refreshable
+	state *cleanupState
+	refs  []any
+}
+
+func newDerivedRefreshable(inner Refreshable, unsubs ...func()) *derivedRefreshable {
+	state := &cleanupState{unsubs: unsubs}
+	d := &derivedRefreshable{
+		inner: inner,
+		state: state,
+	}
+	runtime.AddCleanup(d, func(s *cleanupState) {
+		s.gcDone.Store(true)
+		s.tryCleanup()
+	}, state)
+	return d
+}
+
+func (d *derivedRefreshable) Current() interface{} {
+	return d.inner.Current()
+}
+
+func (d *derivedRefreshable) Subscribe(consumer func(interface{})) (unsubscribe func()) {
+	d.state.subCount.Add(1)
+	innerUnsub := d.inner.Subscribe(consumer)
+	state := d.state
+	var once sync.Once
+	return func() {
+		once.Do(func() {
+			innerUnsub()
+			state.subCount.Add(-1)
+			state.tryCleanup()
+		})
+	}
+}
+
+func (d *derivedRefreshable) Map(mapFn func(interface{}) interface{}) Refreshable {
+	result := d.inner.Map(mapFn)
+	if dr, ok := result.(*derivedRefreshable); ok {
+		dr.refs = append(dr.refs, d)
+	}
+	return result
+}

refreshable/derived_test.go

@@ -0,0 +1,202 @@
+// Copyright (c) 2021 Palantir Technologies. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package refreshable_test
+
+import (
+	"runtime"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/palantir/pkg/refreshable"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// awaitGCCleanup runs GC repeatedly until condition returns true or the timeout expires.
+func awaitGCCleanup(t *testing.T, condition func() bool) {
+	t.Helper()
+	deadline := time.Now().Add(5 * time.Second)
+	for time.Now().Before(deadline) {
+		runtime.GC()
+		time.Sleep(10 * time.Millisecond)
+		if condition() {
+			return
+		}
+	}
+	t.Fatal("timed out waiting for GC cleanup")
+}
+
+func TestMapGCCleanup(t *testing.T) {
+	parent := refreshable.NewDefaultRefreshable(1)
+	var mapCalls atomic.Int32
+	mapped := parent.Map(func(i interface{}) interface{} {
+		mapCalls.Add(1)
+		return i.(int) * 2
+	})
+
+	// Verify the mapped refreshable works.
+	assert.Equal(t, 2, mapped.Current())
+	require.NoError(t, parent.Update(5))
+	assert.Equal(t, 10, mapped.Current())
+
+	// Drop the derived reference and wait for GC cleanup.
+	runtime.KeepAlive(mapped)
+	mapped = nil //nolint:ineffassign
+
+	updateVal := 100
+	awaitGCCleanup(t, func() bool {
+		mapCalls.Store(0)
+		updateVal++
+		_ = parent.Update(updateVal)
+		return mapCalls.Load() == 0
+	})
+}
+
+func TestValidatingGCCleanup(t *testing.T) {
+	parent := refreshable.NewDefaultRefreshable(1)
+	var validateCalls atomic.Int32
+	vr, err := refreshable.NewValidatingRefreshable(parent, func(i interface{}) error {
+		validateCalls.Add(1)
+		return nil
+	})
+	require.NoError(t, err)
+	assert.Equal(t, 1, vr.Current())
+
+	// Drop the validating refreshable and wait for GC cleanup.
+	runtime.KeepAlive(vr)
+	vr = nil //nolint:ineffassign
+
+	updateVal := 100
+	awaitGCCleanup(t, func() bool {
+		validateCalls.Store(0)
+		updateVal++
+		_ = parent.Update(updateVal)
+		return validateCalls.Load() == 0
+	})
+}
+
+func TestMapValidatingGCCleanup(t *testing.T) {
+	parent := refreshable.NewDefaultRefreshable("hello")
+	var mapCalls atomic.Int32
+	vr, err := refreshable.NewMapValidatingRefreshable(parent, func(i interface{}) (interface{}, error) {
+		mapCalls.Add(1)
+		return len(i.(string)), nil
+	})
+	require.NoError(t, err)
+	assert.Equal(t, 5, vr.Current())
+
+	// Drop the validating refreshable and wait for GC cleanup.
+	runtime.KeepAlive(vr)
+	vr = nil //nolint:ineffassign
+
+	vals := []string{"a", "bb", "ccc", "dddd", "eeeee", "ffffff"}
+	idx := 0
+	awaitGCCleanup(t, func() bool {
+		mapCalls.Store(0)
+		_ = parent.Update(vals[idx%len(vals)])
+		idx++
+		return mapCalls.Load() == 0
+	})
+}
+
+func TestMapGCWithActiveSubscriber(t *testing.T) {
+	parent := refreshable.NewDefaultRefreshable(1)
+	mapped := parent.Map(func(i interface{}) interface{} {
+		return i.(int) * 2
+	})
+
+	// Subscribe to the derived refreshable.
+	var latest atomic.Value
+	unsub := mapped.Subscribe(func(i interface{}) {
+		latest.Store(i)
+	})
+
+	require.NoError(t, parent.Update(5))
+	assert.Equal(t, 10, latest.Load())
+
+	// Drop the derived reference but keep the subscription.
+	runtime.KeepAlive(mapped)
+	mapped = nil //nolint:ineffassign
+
+	// Run GC — the subscription should keep updates flowing.
+	for i := 0; i < 5; i++ {
+		runtime.GC()
+		time.Sleep(10 * time.Millisecond)
+	}
+
+	require.NoError(t, parent.Update(7))
+	assert.Equal(t, 14, latest.Load(), "subscription should still receive updates after derived is GC'd")
+
+	// Now unsubscribe — cleanup should fire since gcDone is true and subCount reaches 0.
+	unsub()
+
+	// Verify updates no longer flow. Use a mapFn-based check by creating a new mapped
+	// refreshable to observe parent subscriber behavior indirectly.
+	var mapCalls atomic.Int32
+	probe := parent.Map(func(i interface{}) interface{} {
+		mapCalls.Add(1)
+		return i
+	})
+	_ = probe // keep alive
+
+	// The original subscription's cleanup should have already fired synchronously.
+	// Verify the parent still works for new subscribers.
+	mapCalls.Store(0)
+	require.NoError(t, parent.Update(99))
+	assert.Greater(t, mapCalls.Load(), int32(0), "parent should still notify new subscribers")
+	runtime.KeepAlive(probe)
+}
+
+func TestDerivedMapChain(t *testing.T) {
+	parent := refreshable.NewDefaultRefreshable(1)
+	var middleCalls, finalCalls atomic.Int32
+
+	middle := parent.Map(func(i interface{}) interface{} {
+		middleCalls.Add(1)
+		return i.(int) * 2
+	})
+	final := middle.Map(func(i interface{}) interface{} {
+		finalCalls.Add(1)
+		return i.(int) + 100
+	})
+
+	assert.Equal(t, 102, final.Current())
+
+	require.NoError(t, parent.Update(5))
+	assert.Equal(t, 110, final.Current())
+
+	// Drop the middle reference but keep the final.
+	// The chain should stay alive because final.refs holds middle.
+	middleCalls.Store(0)
+	finalCalls.Store(0)
+	runtime.KeepAlive(middle)
+	middle = nil //nolint:ineffassign
+
+	for i := 0; i < 5; i++ {
+		runtime.GC()
+		time.Sleep(10 * time.Millisecond)
+	}
+
+	middleCalls.Store(0)
+	finalCalls.Store(0)
+	require.NoError(t, parent.Update(10))
+	assert.Equal(t, 120, final.Current())
+	assert.Greater(t, middleCalls.Load(), int32(0), "middle mapFn should still be called")
+	assert.Greater(t, finalCalls.Load(), int32(0), "final mapFn should still be called")
+
+	// Now drop the final reference — entire chain should be cleaned up.
+	runtime.KeepAlive(final)
+	final = nil //nolint:ineffassign
+
+	updateVal := 100
+	awaitGCCleanup(t, func() bool {
+		middleCalls.Store(0)
+		finalCalls.Store(0)
+		updateVal++
+		_ = parent.Update(updateVal)
+		return middleCalls.Load() == 0 && finalCalls.Load() == 0
+	})
+}

refreshable/go.mod

@@ -1,6 +1,6 @@
 module github.com/palantir/pkg/refreshable
 
-go 1.20
+go 1.24
 
 require (
 	github.com/palantir/pkg v1.1.0

refreshable/refreshable_default.go

@@ -81,8 +81,10 @@ func (d *DefaultRefreshable) unsubscribe(consumerFnPtr *func(interface{})) {
 
 func (d *DefaultRefreshable) Map(mapFn func(interface{}) interface{}) Refreshable {
 	newRefreshable := NewDefaultRefreshable(mapFn(d.Current()))
-	d.Subscribe(func(updatedVal interface{}) {
+	unsub := d.Subscribe(func(updatedVal interface{}) {
 		_ = newRefreshable.Update(mapFn(updatedVal))
 	})
-	return newRefreshable
+	derived := newDerivedRefreshable(newRefreshable, unsub)
+	derived.refs = append(derived.refs, d)
+	return derived
 }

refreshable/refreshable_validating.go

@@ -66,30 +66,32 @@ func newValidatingRefreshable(origRefreshable Refreshable, validatingFn func(int
 
 	var lastValidateErr atomic.Value
 	lastValidateErr.Store(errorWrapper{})
-	v := ValidatingRefreshable{
-		Refreshable:     validatedRefreshable,
-		lastValidateErr: &lastValidateErr,
-	}
 
 	updateValueFn := func(i interface{}) {
 		mappedVal, err := validatingFn(i)
 		if err != nil {
-			v.lastValidateErr.Store(errorWrapper{err})
+			lastValidateErr.Store(errorWrapper{err})
 			return
 		}
 		if storeMappedVal {
 			err = validatedRefreshable.Update(mappedVal)
 		} else {
 			err = validatedRefreshable.Update(i)
 		}
-		v.lastValidateErr.Store(errorWrapper{err: err})
+		lastValidateErr.Store(errorWrapper{err: err})
 	}
 
-	origRefreshable.Subscribe(updateValueFn)
+	unsub := origRefreshable.Subscribe(updateValueFn)
 
 	// manually update value after performing subscription. This ensures that, if the current value changed between when
 	// it was fetched earlier in the function and when the subscription was performed, it is properly captured.
 	updateValueFn(origRefreshable.Current())
 
+	derived := newDerivedRefreshable(validatedRefreshable, unsub)
+	derived.refs = append(derived.refs, origRefreshable)
+	v := ValidatingRefreshable{
+		Refreshable:     derived,
+		lastValidateErr: &lastValidateErr,
+	}
 	return &v, nil
 }

refreshable/v2.go

@@ -25,10 +25,10 @@ func ToV2[T any](v1 Refreshable) refreshablev2.Refreshable[T] {
 // FromV2 converts from a v1 Refreshable created by this package to v2 supporting type safety via generics.
 func FromV2[T any](v2 refreshablev2.Refreshable[T]) Refreshable {
 	v1 := NewDefaultRefreshable(v2.Current())
-	v2.Subscribe(func(i T) {
+	unsub := v2.Subscribe(func(i T) {
 		if err := v1.Update(i); err != nil {
 			panic(err)
 		}
 	})
-	return v1
+	return newDerivedRefreshable(v1, unsub)
 }