From c2ed73e03f7175ffb75f159fd10736106aafd154 Mon Sep 17 00:00:00 2001 From: Huy Nguyen Date: Fri, 7 Nov 2025 14:31:59 +1030 Subject: [PATCH] TitleDb loading is handled for multithreading Fix some memory leaks --- TinfoilVibeServer.sln.DotSettings.user | 1 + TinfoilVibeServer/Services/ArchiveHandler.cs | 7 +- .../Services/IndexBuilderService.cs | 14 ++- TinfoilVibeServer/Services/NSPExtractor.cs | 8 +- TinfoilVibeServer/Services/SnapshotService.cs | 22 +++-- .../Services/TitleDatabaseService.cs | 85 ++++++++++++++----- 6 files changed, 97 insertions(+), 40 deletions(-) diff --git a/TinfoilVibeServer.sln.DotSettings.user b/TinfoilVibeServer.sln.DotSettings.user index a8484bd..5a582b6 100644 --- a/TinfoilVibeServer.sln.DotSettings.user +++ b/TinfoilVibeServer.sln.DotSettings.user @@ -52,6 +52,7 @@ ForceIncluded ForceIncluded ForceIncluded + ForceIncluded ForceIncluded ForceIncluded ForceIncluded diff --git a/TinfoilVibeServer/Services/ArchiveHandler.cs b/TinfoilVibeServer/Services/ArchiveHandler.cs index 69a8ae4..dd1ab9b 100644 --- a/TinfoilVibeServer/Services/ArchiveHandler.cs +++ b/TinfoilVibeServer/Services/ArchiveHandler.cs @@ -114,6 +114,7 @@ public sealed class ArchiveHandler : IArchiveHandler var entryCount = 0; try { + // todo: handle and skip multipart archives using var archive = RarArchive.Open(path); entryCount = archive.Entries.Count; foreach (var entry in archive.Entries) @@ -122,10 +123,14 @@ public sealed class ArchiveHandler : IArchiveHandler { try { - using var streamWrapper = new SeekableBufferedStream(entry.OpenEntryStream(), entry.Size, 64 * 1024 * 1024, false); + using var streamWrapper = new SeekableBufferedStream(entry.OpenEntryStream(), entry.Size, 64 * 1024 * 1024, true); var title = _nspExtractor.ExtractFromStream(streamWrapper); if (title != null) titles.Add((entry.Key, entry.Size, title)); } + catch (IncompleteArchiveException incompleteArchiveException) + { + _logger.LogWarning("Incomplete archive {Archive}: {Exception}", path, incompleteArchiveException.Message); + } catch (Exception e) { if (e.Message.StartsWith("Failed to extract NSP")) diff --git a/TinfoilVibeServer/Services/IndexBuilderService.cs b/TinfoilVibeServer/Services/IndexBuilderService.cs index f5de0e2..8921f20 100644 --- a/TinfoilVibeServer/Services/IndexBuilderService.cs +++ b/TinfoilVibeServer/Services/IndexBuilderService.cs @@ -1,13 +1,7 @@ -using System; -using System.Collections.Generic; -using System.Globalization; -using System.Linq; -using System.Security.Cryptography; +using System.Security.Cryptography; using System.Text.Json; using System.Text.RegularExpressions; using LibHac.Ncm; -using Microsoft.Extensions.Configuration; -using Microsoft.Extensions.Logging; using TinfoilVibeServer.Models; namespace TinfoilVibeServer.Services; @@ -86,7 +80,7 @@ public sealed class IndexBuilderService: IHostedService { var titleId = title.TitleId; var name = _titleDb.TryGetTitle(titleId, out var t) - ? t.Name + ? t?.Name : "Unknown"; var versionNumberParsed = (title.ContentMetaType == ContentMetaType.Application) ? title.Version : title.Version * 0x10000; @@ -96,7 +90,7 @@ public sealed class IndexBuilderService: IHostedService { // Patch should use the application title name name = _titleDb.TryGetTitle(title.ApplicationTitle, out var appTitle) - ? appTitle.Name + ? appTitle?.Name : "Unknown"; //_logger.LogInformation("Patch title {TitleId} uses application title {ApplicationTitle}", titleId, title.ApplicationTitle); } @@ -129,7 +123,9 @@ public sealed class IndexBuilderService: IHostedService private IndexCache? LoadCache() { if (!File.Exists(_cachePath)) return null; + _lock.Wait(); var json = File.ReadAllText(_cachePath); + _lock.Release(); _logger.LogInformation("Loaded index cache from {Path}", _cachePath); return JsonSerializer.Deserialize(json); } diff --git a/TinfoilVibeServer/Services/NSPExtractor.cs b/TinfoilVibeServer/Services/NSPExtractor.cs index d51f323..1f3f6d6 100644 --- a/TinfoilVibeServer/Services/NSPExtractor.cs +++ b/TinfoilVibeServer/Services/NSPExtractor.cs @@ -186,13 +186,15 @@ namespace TinfoilVibeServer.Services var storage = new StreamStorage(stream, true); var partition = new PartitionFileSystem(); partition.Initialize(storage).ThrowIfFailure(); + _logger.LogInformation("PFS0 found"); return true; } - catch (Exception e) + catch { - _logger.LogError("Failed to extract NSP: {Exception}", e.Message); - return false; + // ignored } + + return false; } private bool IsXciFileSystem(Stream stream) { diff --git a/TinfoilVibeServer/Services/SnapshotService.cs b/TinfoilVibeServer/Services/SnapshotService.cs index a6c1962..3e0c0e4 100644 --- a/TinfoilVibeServer/Services/SnapshotService.cs +++ b/TinfoilVibeServer/Services/SnapshotService.cs @@ -543,9 +543,16 @@ public sealed class SnapshotService : IDisposable, ISnapshotService, IHostedServ public void RebuildSnapshot() { - // Build a fresh snapshot and persist it. - BuildSnapshotAsync(); // private method inside the same class - PersistSnapshotAsync(); // private method inside the same class + // 1️⃣ Flush the old in‑memory snapshot + _cache.Clear(); + _hashCache.Clear(); + _archiveLookup.Clear(); + _sizeLookup.Clear(); + //_failedAttempts.Clear(); // if you keep per‑user counters + + // 2️⃣ Re‑build from disk again + BuildSnapshotAsync().Wait(); // synchronous – we already own the lock + PersistSnapshotAsync().Wait(); // same SnapshotRebuilt?.Invoke(this, EventArgs.Empty); } #endregion @@ -602,7 +609,6 @@ public sealed class SnapshotService : IDisposable, ISnapshotService, IHostedServ var ext = Path.GetExtension(filePath).ToLowerInvariant(); if (ext is not ".nsp" and not ".xci" and not ".xcz") { - // Open the NSP/XCI with LibHac and read the first stream. // The first stream is the first entry returned by GetContentInfos(). try @@ -612,12 +618,16 @@ public sealed class SnapshotService : IDisposable, ISnapshotService, IHostedServ var first = reader.GetEntries().FirstOrDefault(); if (first == null) return ComputeFullHash(filePath); - return _nspExtractor.ExtractHashFromStream(first.Stream); + using var firstStream = first.Stream; + var hash = _nspExtractor.ExtractHashFromStream(firstStream); + return hash; } catch { // On error, fall back to the full file hash - return ComputeFullHash(filePath); + using var fs = new FileStream(filePath, FileMode.Open, FileAccess.Read, FileShare.Read); + var ncaMetadataWithHash = _nspExtractor.ExtractFromStream(fs); + return ncaMetadataWithHash?.Hash ?? string.Empty; } } else diff --git a/TinfoilVibeServer/Services/TitleDatabaseService.cs b/TinfoilVibeServer/Services/TitleDatabaseService.cs index d29340a..53bab92 100644 --- a/TinfoilVibeServer/Services/TitleDatabaseService.cs +++ b/TinfoilVibeServer/Services/TitleDatabaseService.cs @@ -36,7 +36,9 @@ public sealed class TitleDatabaseService : IHostedService // Regex to find a 16‑digit hex TitleId in a filename private static readonly Regex TitleIdRegex = new( @"([0-9a-fA-F]{8})([0-9a-fA-F]{8})", RegexOptions.Compiled); - + private readonly SemaphoreSlim _reloadLock = new(1, 1); // protects reload logic + private Task? _reloadTask; // the currently running or last finished reload + #endregion #region ctor @@ -107,19 +109,12 @@ public sealed class TitleDatabaseService : IHostedService { var ttlSec = _options.CurrentValue.TtlSeconds; _logger.LogInformation("Reloading title database cache (TTL={TTL}s).", ttlSec); - + var entryOptions = new MemoryCacheEntryOptions() - .SetSlidingExpiration(TimeSpan.FromSeconds(ttlSec)).RegisterPostEvictionCallback((key, value, reason, - state) => - { - _logger.LogInformation("Cache eviction: {Key} ({Reason})", key, reason); - }); // <‑‑ sliding! + .SetSlidingExpiration(TimeSpan.FromSeconds(ttlSec)) + .RegisterPostEvictionCallback((key, value, reason, state) => { _logger.LogInformation("Cache eviction: {Key} ({Reason})", key, reason); }); var dict = await LoadFromDiskAsync() ?? new Dictionary(); - - // Set the new entry – the sliding expiration will now - // automatically move 30 s (or whatever you configured) forward - // every time the entry is accessed via Get/Set. _cache.Set(CacheKey, dict, entryOptions); _logger.LogInformation("Title DB reloaded – {Count} items cached (TTL={TTL}s).", dict.Count, ttlSec); @@ -142,21 +137,20 @@ public sealed class TitleDatabaseService : IHostedService /* ---------------------------------------------------------------- */ public async Task> GetAllAsync() { - if (!_cache.TryGetValue(CacheKey, out Dictionary? dict)) - { - await ReloadCacheAsync(); // cache miss → load from disk - _cache.TryGetValue(CacheKey, out dict); - } - return dict!; // Get() has already slid the entry + await EnsureCacheLoadedAsync().ConfigureAwait(false); + + // The cache entry is guaranteed to exist now + _cache.TryGetValue(CacheKey, out Dictionary? dict); + return dict!; } public async Task GetAsync(string titleId) { - var all = await GetAllAsync(); // slides the entry + var all = await GetAllAsync(); // slides the entry all.TryGetValue(titleId, out var dto); return dto; } - + /* ---------------------------------------------------------------- */ /* 3️⃣ Persist to disk & notify snapshot service /* ---------------------------------------------------------------- */ @@ -165,18 +159,23 @@ public sealed class TitleDatabaseService : IHostedService // Trigger a rebuild so SnapshotService (and any other listeners) // can pick up the new snapshot. //_snapshotService.RebuildSnapshot(); + await Task.CompletedTask; } /* ---------------------------------------------------------------- */ /* 4️⃣ Dispose /* ---------------------------------------------------------------- */ - public void Dispose() => _snapshotService.SnapshotRebuilt -= SnapshotServiceOnSnapshotRebuilt; + public void Dispose() + { + _snapshotService.SnapshotRebuilt -= SnapshotServiceOnSnapshotRebuilt; + _reloadLock.Dispose(); + } private async void SnapshotServiceOnSnapshotRebuilt(object? o, EventArgs eventArgs) { try { - await ReloadCacheAsync(); + await EnsureCacheLoadedAsync(); } catch (Exception e) { @@ -197,7 +196,51 @@ public sealed class TitleDatabaseService : IHostedService #endregion #region Private helpers + /// + /// Makes sure the cache is loaded. If a reload is already in progress, + /// the caller simply awaits the same task. Otherwise it starts a new reload. + /// + private async Task EnsureCacheLoadedAsync() + { + // Fast path – a cache entry exists, no work needed + if (_cache.TryGetValue(CacheKey, out _)) + return; + // Fast path – a reload is already underway + var existingTask = _reloadTask; + if (existingTask != null) + { + await existingTask.ConfigureAwait(false); + return; + } + + // Slow path – we need to start the reload + await _reloadLock.WaitAsync().ConfigureAwait(false); + try + { + // Double‑check after acquiring the lock + if (_cache.TryGetValue(CacheKey, out _)) + return; + + // If another thread started the reload while we were waiting, use it + existingTask = _reloadTask; + if (existingTask != null) + { + await existingTask.ConfigureAwait(false); + return; + } + + // Create the shared task + _reloadTask = ReloadCacheAsync(); + await _reloadTask.ConfigureAwait(false); + } + finally + { + _reloadTask = null; // reset for the next miss + _reloadLock.Release(); + } + } + /// /// Downloads the JSON file from GitHub (raw) if it does not exist /// or the cached copy is older than the remote one.