[Java.Interop] Fix JNI global-reference ownership leaks

external/Java.Interop/src/Java.Interop/Java.Interop/JniPeerMembers.JniInstanceMethods.cs

@@ -1,7 +1,6 @@
 #nullable enable
 
 using System;
-using System.Collections.Concurrent;
 
 namespace Java.Interop
 {
@@ -39,16 +38,16 @@ internal    JniType                                 JniPeerType {
 
 		readonly Type                                       DeclaringType;
 
-		ConcurrentDictionary<string, JniMethodInfo>?             instanceMethods;
-		ConcurrentDictionary<Type, JniInstanceMethods>?          subclassConstructors;
+		JniValueCache<string, JniMethodInfo>?                     instanceMethods;
+		JniValueCache<Type, JniInstanceMethods>?                 subclassConstructors;
 
-		ConcurrentDictionary<string, JniMethodInfo>               InstanceMethods      => GetOrCreate (ref instanceMethods, 3);
-		ConcurrentDictionary<Type, JniInstanceMethods>            SubclassConstructors => GetOrCreate (ref subclassConstructors, 1);
+		JniValueCache<string, JniMethodInfo>                      InstanceMethods      => JniValueCache<string, JniMethodInfo>.GetOrCreate (ref instanceMethods, 1, 3, static value => value.StaticRedirect?.Dispose ());
+		JniValueCache<Type, JniInstanceMethods>                  SubclassConstructors => JniValueCache<Type, JniInstanceMethods>.GetOrCreate (ref subclassConstructors, 1, 1, static value => value.Dispose ());
 
 		internal void Dispose ()
 		{
-			Clear (ref instanceMethods);
-			Clear (ref subclassConstructors, static value => value.Dispose ());
+			JniValueCache<string, JniMethodInfo>.Dispose (ref instanceMethods);
+			JniValueCache<Type, JniInstanceMethods>.Dispose (ref subclassConstructors);
 
 			if (jniPeerType != null)
 				jniPeerType.Dispose ();
@@ -111,15 +110,20 @@ JniMethodInfo GetMethodInfo (ReadOnlySpan<char> method, ReadOnlySpan<char> signa
 				var methodName = newMethod.Value.TargetJniMethodName is string name ? name.AsSpan () : method;
 				var methodSig  = newMethod.Value.TargetJniMethodSignature is string sig ? sig.AsSpan () : signature;
 
-				using var t = new JniType (typeName);
-				if (newMethod.Value.TargetJniMethodInstanceToStatic &&
-						t.TryGetStaticMethod (methodName, methodSig, out m)) {
-					m.ParameterCount = newMethod.Value.TargetJniMethodParameterCount;
-					m.StaticRedirect = new JniType (typeName);
-					return m;
-				}
-				if (t.TryGetInstanceMethod (methodName, methodSig, out m)) {
-					return m;
+				JniType? t = new JniType (typeName);
+				try {
+					if (newMethod.Value.TargetJniMethodInstanceToStatic &&
+							t.TryGetStaticMethod (methodName, methodSig, out m)) {
+						m.ParameterCount = newMethod.Value.TargetJniMethodParameterCount;
+						m.StaticRedirect = t;
+						t = null;
+						return m;
+					}
+					if (t.TryGetInstanceMethod (methodName, methodSig, out m)) {
+						return m;
+					}
+				} finally {
+					t?.Dispose ();
 				}
 				Console.Error.WriteLine ($"warning: For declared method `{Members.JniPeerTypeName}.{method}.{signature}`, could not find requested method `{typeName}.{methodName}.{methodSig}`!");
 			}

external/Java.Interop/src/Java.Interop/Java.Interop/JniPeerMembers.JniStaticMethods.cs

@@ -1,7 +1,6 @@
 #nullable enable
 
 using System;
-using System.Collections.Concurrent;
 
 namespace Java.Interop
 {
@@ -15,13 +14,13 @@ internal JniStaticMethods (JniPeerMembers members)
 
 		internal    readonly    JniPeerMembers              Members;
 
-		ConcurrentDictionary<string, JniMethodInfo>? staticMethods;
+		JniValueCache<string, JniMethodInfo>? staticMethods;
 
-		ConcurrentDictionary<string, JniMethodInfo> StaticMethods => GetOrCreate (ref staticMethods, 3);
+		JniValueCache<string, JniMethodInfo> StaticMethods => JniValueCache<string, JniMethodInfo>.GetOrCreate (ref staticMethods, 1, 3, static value => value.StaticRedirect?.Dispose ());
 
 		internal void Dispose ()
 		{
-			Clear (ref staticMethods);
+			JniValueCache<string, JniMethodInfo>.Dispose (ref staticMethods);
 		}
 
 		public JniMethodInfo GetMethodInfo (string encodedMember)
@@ -38,12 +37,20 @@ JniMethodInfo GetMethodInfo (ReadOnlySpan<char> method, ReadOnlySpan<char> signa
 			var m              = (JniMethodInfo?) null;
 			var newMethod      = JniEnvironment.Runtime.TypeManager.GetReplacementMethodInfo (Members.JniPeerTypeName, method, signature);
 			if (newMethod.HasValue) {
-				using var t = new JniType (newMethod.Value.TargetJniType ?? Members.JniPeerTypeName);
-				if (t.TryGetStaticMethod (
-						newMethod.Value.TargetJniMethodName is string name ? name.AsSpan () : method,
-						newMethod.Value.TargetJniMethodSignature is string sig ? sig.AsSpan () : signature,
-						out m)) {
-					return m;
+				JniType? t = new JniType (newMethod.Value.TargetJniType ?? Members.JniPeerTypeName);
+				try {
+					if (t.TryGetStaticMethod (
+							newMethod.Value.TargetJniMethodName is string name ? name.AsSpan () : method,
+							newMethod.Value.TargetJniMethodSignature is string sig ? sig.AsSpan () : signature,
+							out m)) {
+						if (!JniEnvironment.Types.IsSameObject (t.PeerReference, Members.JniPeerType.PeerReference)) {
+							m.StaticRedirect = t;
+							t = null;
+						}
+						return m;
+					}
+				} finally {
+					t?.Dispose ();
 				}
 			}
 			if (Members.JniPeerType.TryGetStaticMethod (method, signature, out m)) {
@@ -72,23 +79,28 @@ JniType GetMethodDeclaringType (JniMethodInfo method)
 			if (fallbackTypes == null) {
 				return null;
 			}
-			foreach (var ft in fallbackTypes) {
-				JniType? t = null;
-				try {
+			JniType? t = null;
+			try {
+				JniMethodInfo? m = null;
+				foreach (var ft in fallbackTypes) {
 					if (!JniType.TryParse (ft, out t)) {
 						continue;
 					}
-					if (t.TryGetStaticMethod (method, signature, out var m)) {
-						m.StaticRedirect    = t;
-						t                   = null;
-						return m;
+					if (t.TryGetStaticMethod (method, signature, out m)) {
+						break;
 					}
+					t.Dispose ();
+					t = null;
 				}
-				finally {
-					t?.Dispose ();
+				if (m != null) {
+					// Transfer ownership only after the fallback enumerator has been disposed.
+					m.StaticRedirect = t;
+					t = null;
 				}
+				return m;
+			} finally {
+				t?.Dispose ();
 			}
-			return null;
 		}
 
 		public unsafe void InvokeVoidMethod (string encodedMember, JniArgumentValue* parameters)

external/Java.Interop/src/Java.Interop/Java.Interop/JniPeerMembers.JniValueCache.cs

@@ -0,0 +1,92 @@
+#nullable enable
+
+using System;
+using System.Collections;
+using System.Collections.Concurrent;
+using System.Collections.Generic;
+using System.Diagnostics.CodeAnalysis;
+using System.Threading;
+
+namespace Java.Interop {
+
+	partial class JniPeerMembers {
+
+		private sealed class JniValueCache<TKey, TValue> : IDisposable, IReadOnlyDictionary<TKey, TValue>
+			where TKey : notnull
+			where TValue : class
+		{
+
+			readonly ConcurrentDictionary<TKey, TValue> values;
+			readonly Action<TValue> dispose;
+
+			public JniValueCache (int concurrencyLevel, int capacity, Action<TValue> dispose)
+			{
+				values       = new ConcurrentDictionary<TKey, TValue> (concurrencyLevel, capacity);
+				this.dispose = dispose;
+			}
+
+			public int Count => values.Count;
+			public IEnumerable<TKey> Keys => values.Keys;
+			public IEnumerable<TValue> Values => values.Values;
+			public TValue this [TKey key] => values [key];
+
+			public bool ContainsKey (TKey key) => values.ContainsKey (key);
+			public bool TryGetValue (TKey key, [MaybeNullWhen (false)] out TValue value) => values.TryGetValue (key, out value);
+
+			public IEnumerator<KeyValuePair<TKey, TValue>> GetEnumerator () => values.GetEnumerator ();
+			IEnumerator IEnumerable.GetEnumerator () => GetEnumerator ();
+
+			internal static JniValueCache<TKey, TValue> GetOrCreate (ref JniValueCache<TKey, TValue>? cache, int concurrencyLevel, int capacity, Action<TValue> dispose)
+			{
+				var value = Volatile.Read (ref cache);
+				if (value != null)
+					return value;
+
+				var candidate = new JniValueCache<TKey, TValue> (concurrencyLevel, capacity, dispose);
+				var existing = Interlocked.CompareExchange (ref cache, candidate, null);
+				if (existing == null)
+					return candidate;
+
+				candidate.Dispose ();
+				return existing;
+			}
+
+			internal static void Dispose (ref JniValueCache<TKey, TValue>? cache)
+			{
+				Interlocked.Exchange (ref cache, null)?.Dispose ();
+			}
+
+			public TValue GetOrAdd (TKey key, Func<TKey, TValue> factory)
+			{
+				return GetOrAdd (key, static (key, factory) => factory (key), factory);
+			}
+
+			public TValue GetOrAdd<TArg> (TKey key, Func<TKey, TArg, TValue> factory, TArg argument)
+			{
+				if (values.TryGetValue (key, out var value))
+					return value;
+
+				// ConcurrentDictionary may invoke a GetOrAdd factory multiple times and discard
+				// losing values. Construct explicitly so an unpublished owner can be disposed.
+				// JNI lookup can also reenter this cache, so do not lock construction.
+				TValue? candidate = factory (key, argument);
+				try {
+					value = values.GetOrAdd (key, candidate);
+					if (ReferenceEquals (value, candidate))
+						candidate = null;
+					return value;
+				} finally {
+					if (candidate != null)
+						dispose (candidate);
+				}
+			}
+
+			public void Dispose ()
+			{
+				foreach (var value in values.Values)
+					dispose (value);
+				values.Clear ();
+			}
+		}
+	}
+}

external/Java.Interop/src/Java.Interop/Java.Interop/JniPeerMembers.cs

@@ -136,14 +136,14 @@ static void Clear<TKey, TValue> (ref ConcurrentDictionary<TKey, TValue>? diction
 
 		protected virtual void Dispose (bool disposing)
 		{
-			if (!disposing || jniPeerType == null)
+			if (!disposing)
 				return;
 
 			instanceMethods.Dispose ();
 			instanceFields.Dispose ();
 			staticMethods.Dispose ();
 			staticFields.Dispose ();
-			jniPeerType.Dispose ();
+			jniPeerType?.Dispose ();
 
 			jniPeerType     = null;
 		}

external/Java.Interop/src/Java.Interop/Java.Interop/JniType.cs

@@ -5,7 +5,6 @@
 using System.Diagnostics;
 using System.Diagnostics.CodeAnalysis;
 using System.Collections.Generic;
-using System.Linq;
 using System.Reflection;
 using System.Runtime.CompilerServices;
 using System.Runtime.InteropServices;
@@ -122,6 +121,12 @@ public void Dispose ()
 			JniObjectReference.Dispose (ref peerReference);
 		}
 
+		internal void DisposeUnlessRegisteredWithRuntime ()
+		{
+			if (!registered)
+				Dispose ();
+		}
+
 		public JniType? GetSuperclass ()
 		{
 			AssertValid ();
@@ -151,30 +156,55 @@ public bool IsInstanceOfType (JniObjectReference value)
 			return JniEnvironment.Types.IsInstanceOf (value, PeerReference);
 		}
 
-#pragma warning disable 0414
-		// This isn't used anywhere; it's just present so that the GC won't collect the referenced delegates.
-		JniNativeMethodRegistration[]? methods;
-#pragma warning restore 0414
+		object? nativeMethodsLock;
+		// Retains delegates from every batch JNI may have partially registered.
+		List<JniNativeMethodRegistration[]>? methods;
+
+		object GetNativeMethodsLock ()
+		{
+			var value = Volatile.Read (ref nativeMethodsLock);
+			if (value != null)
+				return value;
+
+			var candidate = new object ();
+			return Interlocked.CompareExchange (ref nativeMethodsLock, candidate, null) ?? candidate;
+		}
 
+		/// <remarks>
+		/// Once a non-empty registration is requested, the runtime retains this type and its
+		/// delegates until unregistration or disposal, even if registration throws: JNI may
+		/// have registered part of the batch.
+		/// </remarks>
 		[RequiresDynamicCode ("Native method registration via JniNativeMethodRegistration[] requires dynamic code generation. Use the blittable RegisterNatives(JniObjectReference, ReadOnlySpan<JniNativeMethod>) overload with statically-compiled function pointers for Native AOT compatibility.")]
 		public void RegisterNativeMethods (params JniNativeMethodRegistration[] methods)
 		{
 			AssertValid ();
 
 			if (methods == null)
 				throw new ArgumentNullException (nameof (methods));
+			if (methods.Length == 0)
+				return;
 
-			JniEnvironment.Types.RegisterNatives (PeerReference, methods, checked ((int)methods.Length));
-			// Prevents method delegates from being GC'd so long as this type remains
-			this.methods = methods;
-			RegisterWithRuntime ();
+			lock (GetNativeMethodsLock ()) {
+				// Retain each batch before calling RegisterNatives: JNI stores only the
+				// unmanaged function pointers and may publish part of a batch before throwing.
+				// Storing it afterward could therefore leave callable pointers to collected
+				// delegates.
+				this.methods ??= new List<JniNativeMethodRegistration[]> ();
+				this.methods.Add (methods);
+				RegisterWithRuntime ();
+				JniEnvironment.Types.RegisterNatives (PeerReference, methods, methods.Length);
+			}
 		}
 
 		public void UnregisterNativeMethods ()
 		{
 			AssertValid ();
 
-			JniEnvironment.Types.UnregisterNatives (PeerReference);
+			lock (GetNativeMethodsLock ()) {
+				JniEnvironment.Types.UnregisterNatives (PeerReference);
+				methods = null;
+			}
 		}
 
 		public JniMethodInfo GetConstructor (string signature)

external/Java.Interop/src/Java.Interop/Java.Interop/ManagedPeer.cs

@@ -282,30 +282,32 @@ static unsafe void RegisterNativeMembers (
 
 			try {
 				var r_nativeClass   = new JniObjectReference (n_nativeClass);
-#pragma warning disable CA2000
+#pragma warning disable CA2000 // Disposed below unless native registration transfers ownership to the runtime.
 				var nativeClass     = new JniType (ref r_nativeClass, JniObjectReferenceOptions.Copy);
 #pragma warning restore CA2000
+				try {
+					var methodsRef              = new JniObjectReference (n_methods);
 
-				var methodsRef              = new JniObjectReference (n_methods);
-
-				var typeSig                 = new JniTypeSignature (nativeClass.Name);
-				var type                    = GetTypeFromSignature (JniEnvironment.Runtime.TypeManager, typeSig);
+					var typeSig                 = new JniTypeSignature (nativeClass.Name);
+					var type                    = GetTypeFromSignature (JniEnvironment.Runtime.TypeManager, typeSig);
 
-				int methodsLength           = JniEnvironment.Strings.GetStringLength (methodsRef);
-				var methodsChars            = JniEnvironment.Strings.GetStringChars (methodsRef, null);
-				var methods                 = new ReadOnlySpan<char>(methodsChars, methodsLength);
-				try {
-					JniEnvironment.Runtime.TypeManager.RegisterNativeMembers (nativeClass, type, methods);
-				}
-				catch (Exception e) {
-					throw new NotSupportedException (
-							$"Unable to register native members for Java type `{nativeClass.Name}` <=> managed type `{type?.AssemblyQualifiedName}`.",
-							e);
-				}
-				finally {
-					JniEnvironment.Strings.ReleaseStringChars (methodsRef, methodsChars);
+					int methodsLength           = JniEnvironment.Strings.GetStringLength (methodsRef);
+					var methodsChars            = JniEnvironment.Strings.GetStringChars (methodsRef, null);
+					var methods                 = new ReadOnlySpan<char>(methodsChars, methodsLength);
+					try {
+						JniEnvironment.Runtime.TypeManager.RegisterNativeMembers (nativeClass, type, methods);
+					}
+					catch (Exception e) {
+						throw new NotSupportedException (
+								$"Unable to register native members for Java type `{nativeClass.Name}` <=> managed type `{type?.AssemblyQualifiedName}`.",
+								e);
+					}
+					finally {
+						JniEnvironment.Strings.ReleaseStringChars (methodsRef, methodsChars);
+					}
+				} finally {
+					nativeClass.DisposeUnlessRegisteredWithRuntime ();
 				}
-
 			}
 			catch (Exception e) {
 				__r?.OnUserUnhandledException (ref envp, e);