Skip to content

Commit 211c043

Browse files
yaeldMSTomFinley
authored andcommitted
Create a shorter temp file name for model loading. (#397)
Create a shorter temp file name for model loading, as well as remove the potential for a race condition among multiple openings by using the creation of a lock file.
1 parent f94203e commit 211c043

File tree

1 file changed

+28
-11
lines changed

1 file changed

+28
-11
lines changed

src/Microsoft.ML.Data/Model/Repository.cs

+28-11
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,6 @@
66
using System.Collections.Generic;
77
using System.IO;
88
using System.IO.Compression;
9-
using Microsoft.ML.Runtime.Data;
109
using Microsoft.ML.Runtime.Internal.Utilities;
1110

1211
namespace Microsoft.ML.Runtime.Model
@@ -73,7 +72,7 @@ public void Dispose()
7372
}
7473
}
7574

76-
// These are the open entries that may contain streams into our _dirTemp.
75+
// These are the open entries that may contain streams into our DirTemp.
7776
private List<Entry> _open;
7877

7978
private bool _disposed;
@@ -108,19 +107,37 @@ internal Repository(bool needDir, IExceptionContext ectx)
108107
PathMap = new Dictionary<string, string>();
109108
_open = new List<Entry>();
110109
if (needDir)
111-
{
112-
DirTemp = GetTempPath();
113-
Directory.CreateDirectory(DirTemp);
114-
}
110+
DirTemp = GetShortTempDir();
115111
else
116112
GC.SuppressFinalize(this);
117113
}
118114

119-
// REVIEW: This should use host environment functionality.
120-
private static string GetTempPath()
115+
private static string GetShortTempDir()
116+
{
117+
var rnd = RandomUtils.Create();
118+
string path;
119+
do
120+
{
121+
path = Path.Combine(Path.GetTempPath(), "TLC_" + rnd.Next().ToString("X"));
122+
path = Path.GetFullPath(path);
123+
Directory.CreateDirectory(path);
124+
}
125+
while (!EnsureDirectory(path));
126+
return path;
127+
}
128+
129+
private static bool EnsureDirectory(string path)
121130
{
122-
Guid guid = Guid.NewGuid();
123-
return Path.GetFullPath(Path.Combine(Path.GetTempPath(), "TLC_" + guid.ToString()));
131+
path = Path.GetFullPath(Path.Combine(path, ".lock"));
132+
try
133+
{
134+
using (var stream = new FileStream(path, FileMode.CreateNew))
135+
return true;
136+
}
137+
catch
138+
{
139+
return false;
140+
}
124141
}
125142

126143
~Repository()
@@ -232,7 +249,7 @@ protected void GetPath(out string pathEnt, out string pathTemp, string dir, stri
232249
_ectx.CheckParam(!name.Contains(".."), nameof(name));
233250

234251
// The gymnastics below are meant to deal with bad invocations including absolute paths, etc.
235-
// That's why we go through it even if _dirTemp is null.
252+
// That's why we go through it even if DirTemp is null.
236253
string root = Path.GetFullPath(DirTemp ?? @"x:\dummy");
237254
string entityPath = Path.Combine(root, dir ?? "", name);
238255
entityPath = Path.GetFullPath(entityPath);

0 commit comments

Comments
 (0)