From 71ade1bc8cefbe3507e76210245f61300a31d2a3 Mon Sep 17 00:00:00 2001
From: Yaohui Liu <AsakusaRinne@gmail.com>
Date: Sat, 13 May 2023 21:45:16 +0800
Subject: [PATCH] fix: error when using graph in multi-threads.

---
 src/TensorFlowNET.Core/Device/DeviceSpec.cs   |  5 ++-
 .../Basics/ThreadSafeTest.cs                  | 41 +++++++++++++++++++
 2 files changed, 44 insertions(+), 2 deletions(-)
 create mode 100644 test/TensorFlowNET.UnitTest/Basics/ThreadSafeTest.cs

diff --git a/src/TensorFlowNET.Core/Device/DeviceSpec.cs b/src/TensorFlowNET.Core/Device/DeviceSpec.cs
index f4ea8cf0..255191cb 100644
--- a/src/TensorFlowNET.Core/Device/DeviceSpec.cs
+++ b/src/TensorFlowNET.Core/Device/DeviceSpec.cs
@@ -1,4 +1,5 @@
 using System;
+using System.Collections.Concurrent;
 using System.Collections.Generic;
 using System.Text;
 using System.Threading.Tasks;
@@ -7,8 +8,8 @@ namespace Tensorflow.Device
 {
     public class DeviceSpec
     {
-        private static Dictionary<string, Components> _STRING_TO_COMPONENTS_CACHE = new();
-        private static Dictionary<Components, string> _COMPONENTS_TO_STRING_CACHE = new();
+        private static ConcurrentDictionary<string, Components> _STRING_TO_COMPONENTS_CACHE = new();
+        private static ConcurrentDictionary<Components, string> _COMPONENTS_TO_STRING_CACHE = new();
         private string _job;
         private int _replica;
         private int _task;
diff --git a/test/TensorFlowNET.UnitTest/Basics/ThreadSafeTest.cs b/test/TensorFlowNET.UnitTest/Basics/ThreadSafeTest.cs
new file mode 100644
index 00000000..6a633448
--- /dev/null
+++ b/test/TensorFlowNET.UnitTest/Basics/ThreadSafeTest.cs
@@ -0,0 +1,41 @@
+using Microsoft.VisualStudio.TestTools.UnitTesting;
+using System;
+using System.Collections.Generic;
+using System.IO;
+using System.Linq;
+using System.Text;
+using System.Threading;
+using System.Threading.Tasks;
+using Tensorflow;
+using static Tensorflow.Binding;
+
+namespace TensorFlowNET.UnitTest.Basics
+{
+    [TestClass]
+    public class ThreadSafeTest
+    {
+        [TestMethod]
+        public void GraphWithMultiThreads()
+        {
+            List<Thread> threads = new List<Thread>();
+
+            const int THREADS_COUNT = 5;
+
+            for (int t = 0; t < THREADS_COUNT; t++)
+            {
+                Thread thread = new Thread(() =>
+                {
+                    Graph g = new Graph();
+                    Session session = new Session(g);
+                    session.as_default();
+                    var input = tf.placeholder(tf.int32, shape: new Shape(6));
+                    var op = tf.reshape(input, new int[] { 2, 3 });
+                });
+                thread.Start();
+                threads.Add(thread);
+            }
+
+            threads.ForEach(t => t.Join());
+        }
+    }
+}