UDF bug fix caused by ThreadStatic BroadcastVariablesRegistry (#551)

Author: Niharikadutta

src/csharp/Microsoft.Spark.E2ETest/UdfTests/UdfSimpleTypesTests.cs

@@ -5,6 +5,7 @@
 using System;
 using System.Collections.Generic;
 using System.Linq;
+using System.Threading;
 using Microsoft.Spark.Sql;
 using Microsoft.Spark.Sql.Types;
 using Xunit;
@@ -166,5 +167,29 @@ public void TestUdfWithReturnAsTimestampType()
                 }
             }
         }
+
+        /// <summary>
+        /// Test to validate UDFs defined in separate threads work.
+        /// </summary>
+        [Fact]
+        public void TestUdfWithMultipleThreads()
+        {
+            try
+            {
+                void DefineUdf() => Udf<string, string>(str => str);
+
+                // Define a UDF in the main thread.
+                Udf<string, string>(str => str);
+
+                // Verify a UDF can be defined in a separate thread.
+                Thread t = new Thread(DefineUdf);
+                t.Start();
+                t.Join();
+            }
+            catch (Exception)
+            {
+                Assert.True(false);
+            }
+        }
     }
 }

src/csharp/Microsoft.Spark/Broadcast.cs

@@ -4,6 +4,7 @@
 using System.IO;
 using System.Runtime.Serialization;
 using System.Runtime.Serialization.Formatters.Binary;
+using System.Threading;
 using Microsoft.Spark.Interop;
 using Microsoft.Spark.Interop.Ipc;
 using Microsoft.Spark.Services;
@@ -261,28 +262,26 @@ internal static void Remove(long bid)
     /// </summary>
     internal static class JvmBroadcastRegistry
     {
-        [ThreadStatic]
-        private static readonly List<JvmObjectReference> s_jvmBroadcastVariables =
-            new List<JvmObjectReference>();
+        private static ThreadLocal<List<JvmObjectReference>> s_jvmBroadcastVariables = 
+            new ThreadLocal<List<JvmObjectReference>>(() => new List<JvmObjectReference>());
 
         /// <summary>
         /// Adds a JVMObjectReference object of type <see cref="Broadcast{T}"/> to the list.
         /// </summary>
         /// <param name="broadcastJvmObject">JVMObjectReference of the Broadcast variable</param>
         internal static void Add(JvmObjectReference broadcastJvmObject) =>
-            s_jvmBroadcastVariables.Add(broadcastJvmObject);
+            s_jvmBroadcastVariables.Value.Add(broadcastJvmObject);
 
         /// <summary>
         /// Clears s_jvmBroadcastVariables of all the JVMObjectReference objects of type
         /// <see cref="Broadcast{T}"/>.
         /// </summary>
-        internal static void Clear() => s_jvmBroadcastVariables.Clear();
+        internal static void Clear() => s_jvmBroadcastVariables.Value.Clear();
 
         /// <summary>
         /// Returns the static member s_jvmBroadcastVariables.
         /// </summary>
         /// <returns>A list of all broadcast objects of type <see cref="JvmObjectReference"/></returns>
-        internal static List<JvmObjectReference> GetAll() => s_jvmBroadcastVariables;
+        internal static List<JvmObjectReference> GetAll() => s_jvmBroadcastVariables.Value;
     }
 }
-