accounts/keystore: add Close method and improve finalizer documentation

accounts/keystore/keystore.go

@@ -69,6 +69,7 @@ type KeyStore struct {
 	updateFeed  event.Feed              // Event feed to notify wallet additions/removals
 	updateScope event.SubscriptionScope // Subscription scope tracking current live listeners
 	updating    bool                    // Whether the event notification loop is running
+	closed      bool                    // Whether the KeyStore has been closed
 
 	mu       sync.RWMutex
 	importMu sync.Mutex // Import Mutex locks the import to prevent two insertions from racing
@@ -96,11 +97,14 @@ func (ks *KeyStore) init(keydir string) {
 	ks.unlocked = make(map[common.Address]*unlocked)
 	ks.cache, ks.changes = newAccountCache(keydir)
 
-	// TODO: In order for this finalizer to work, there must be no references
-	// to ks. addressCache doesn't keep a reference but unlocked keys do,
-	// so the finalizer will not trigger until all timed unlocks have expired.
+	// Set up finalizer to close the account cache when KeyStore is garbage collected.
+	// The finalizer will only trigger when there are no more references to ks.
+	// Unlocked keys with active timers may keep references to ks, so call Close()
+	// explicitly to ensure proper cleanup before the KeyStore is no longer needed.
 	runtime.AddCleanup(ks, func(c *accountCache) {
-		c.close()
+		if c != nil {
+			c.close()
+		}
 	}, ks.cache)
 
 	// Create the initial list of wallets from the cache
@@ -500,6 +504,39 @@ func (ks *KeyStore) isUpdating() bool {
 	return ks.updating
 }
 
+// Close properly shuts down the KeyStore by clearing all unlocked keys
+// and stopping any active timers. This ensures that the finalizer can
+// trigger properly when the KeyStore is no longer referenced.
+func (ks *KeyStore) Close() error {
+	ks.mu.Lock()
+	defer ks.mu.Unlock()
+
+	// Check if already closed
+	if ks.closed {
+		return nil
+	}
+
+	// Clear all unlocked keys and stop their timers
+	for addr, u := range ks.unlocked {
+		if u.abort != nil {
+			select {
+			case <-u.abort:
+				// Channel already closed
+			default:
+				close(u.abort)
+			}
+		}
+		zeroKey(u.PrivateKey)
+		delete(ks.unlocked, addr)
+	}
+
+	// Mark as closed - don't close cache here to avoid double-close
+	// The finalizer will handle cache cleanup
+	ks.closed = true
+
+	return nil
+}
+
 // zeroKey zeroes a private key in memory.
 func zeroKey(k *ecdsa.PrivateKey) {
 	b := k.D.Bits()

accounts/keystore/keystore_test.go

@@ -457,6 +457,51 @@ func checkEvents(t *testing.T, want []walletEvent, have []walletEvent) {
 	}
 }
 
+func TestKeyStoreClose(t *testing.T) {
+	t.Parallel()
+	_, ks := tmpKeyStore(t)
+
+	// Create a new account
+	a, err := ks.NewAccount("foo")
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Unlock the account with a timeout
+	err = ks.TimedUnlock(a, "foo", 10*time.Second)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Verify the account is unlocked
+	ks.mu.RLock()
+	unlocked := len(ks.unlocked)
+	ks.mu.RUnlock()
+	if unlocked != 1 {
+		t.Errorf("expected 1 unlocked account, got %d", unlocked)
+	}
+
+	// Close the keystore
+	err = ks.Close()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Verify all accounts are locked after close
+	ks.mu.RLock()
+	unlocked = len(ks.unlocked)
+	ks.mu.RUnlock()
+	if unlocked != 0 {
+		t.Errorf("expected 0 unlocked accounts after close, got %d", unlocked)
+	}
+
+	// Verify that the account cache is closed
+	if ks.cache != nil {
+		// The cache should be closed, but we can't easily test this without
+		// exposing internal state. The finalizer will handle the cleanup.
+	}
+}
+
 func tmpKeyStore(t *testing.T) (string, *KeyStore) {
 	d := t.TempDir()
 	return d, NewKeyStore(d, veryLightScryptN, veryLightScryptP)