Share tenancy-neutral identity-map state across ForTenant sessions (#4947)

src/DocumentDbTests/Bugs/Bug_4947_for_tenant_tenantless_documents.cs

@@ -0,0 +1,161 @@
+using System;
+using System.Linq;
+using System.Threading.Tasks;
+using Marten;
+using Marten.Testing.Harness;
+using Shouldly;
+using Xunit;
+
+namespace DocumentDbTests.Bugs;
+
+public class Bug_4947_for_tenant_tenantless_documents: BugIntegrationContext
+{
+    public class GlobalDoc
+    {
+        public Guid Id { get; set; }
+        public string Label { get; set; }
+    }
+
+    public class TenantedDoc
+    {
+        public Guid Id { get; set; }
+        public string Label { get; set; }
+    }
+
+    [Fact]
+    public async Task pending_tenantless_document_is_visible_through_for_tenant()
+    {
+        // The exact scenario from GH-4947: identity session, Store() a tenancy-neutral
+        // document (not yet committed), then read it back through ForTenant()
+        await using var session = theStore.IdentitySession();
+
+        var doc = new GlobalDoc { Id = Guid.NewGuid(), Label = "foo" };
+        session.Store(doc);
+
+        var loaded = await session.ForTenant("bar").LoadAsync<GlobalDoc>(doc.Id);
+
+        loaded.ShouldNotBeNull();
+        loaded.ShouldBeSameAs(doc);
+    }
+
+    [Fact]
+    public async Task pending_tenantless_document_is_visible_through_load_many_for_tenant()
+    {
+        await using var session = theStore.IdentitySession();
+
+        var one = new GlobalDoc { Id = Guid.NewGuid(), Label = "one" };
+        var two = new GlobalDoc { Id = Guid.NewGuid(), Label = "two" };
+        session.Store(one, two);
+
+        var loaded = await session.ForTenant("bar").LoadManyAsync<GlobalDoc>(one.Id, two.Id);
+
+        loaded.Count.ShouldBe(2);
+        loaded.OrderBy(x => x.Label).Select(x => x.Label).ShouldBe(new[] { "one", "two" });
+    }
+
+    [Fact]
+    public async Task committed_tenantless_document_resolves_to_the_same_instance_across_tenants()
+    {
+        var id = Guid.NewGuid();
+
+        await using (var seed = theStore.LightweightSession())
+        {
+            seed.Store(new GlobalDoc { Id = id, Label = "global" });
+            await seed.SaveChangesAsync();
+        }
+
+        await using var session = theStore.IdentitySession();
+
+        var viaA = await session.ForTenant("a").LoadAsync<GlobalDoc>(id);
+        viaA.ShouldNotBeNull();
+        viaA.Label.ShouldBe("global");
+
+        var viaB = await session.ForTenant("b").LoadAsync<GlobalDoc>(id);
+        viaB.ShouldNotBeNull();
+
+        // One document, one identity map entry -- every ForTenant view of the same identity
+        // session (and the session itself) must resolve the same tracked instance
+        viaB.ShouldBeSameAs(viaA);
+
+        var direct = await session.LoadAsync<GlobalDoc>(id);
+        direct.ShouldBeSameAs(viaA);
+    }
+
+    [Fact]
+    public async Task tenantless_document_stored_through_for_tenant_is_visible_to_the_parent_session()
+    {
+        await using var session = theStore.IdentitySession();
+
+        var doc = new GlobalDoc { Id = Guid.NewGuid(), Label = "foo" };
+        session.ForTenant("a").Store(doc);
+
+        (await session.LoadAsync<GlobalDoc>(doc.Id)).ShouldBeSameAs(doc);
+        (await session.ForTenant("b").LoadAsync<GlobalDoc>(doc.Id)).ShouldBeSameAs(doc);
+    }
+
+    [Fact]
+    public async Task conjoined_documents_stay_isolated_per_tenant()
+    {
+        // Guard rail for #4801: a conjoined document must NOT be shared across ForTenant views
+        StoreOptions(opts => opts.Schema.For<TenantedDoc>().MultiTenanted());
+
+        var id = Guid.NewGuid();
+
+        await using (var seed = theStore.LightweightSession("a"))
+        {
+            seed.Store(new TenantedDoc { Id = id, Label = "tenant-a" });
+            await seed.SaveChangesAsync();
+        }
+
+        await using (var seed = theStore.LightweightSession("b"))
+        {
+            seed.Store(new TenantedDoc { Id = id, Label = "tenant-b" });
+            await seed.SaveChangesAsync();
+        }
+
+        await using var session = theStore.IdentitySession("a");
+
+        var a = await session.ForTenant("a").LoadAsync<TenantedDoc>(id);
+        a.Label.ShouldBe("tenant-a");
+
+        var b = await session.ForTenant("b").LoadAsync<TenantedDoc>(id);
+        b.Label.ShouldBe("tenant-b");
+
+        b.ShouldNotBeSameAs(a);
+    }
+
+    [Fact]
+    public async Task mixed_store_shares_only_the_tenancy_neutral_document()
+    {
+        StoreOptions(opts => opts.Schema.For<TenantedDoc>().MultiTenanted());
+
+        var globalId = Guid.NewGuid();
+        var tenantedId = Guid.NewGuid();
+
+        await using (var seed = theStore.LightweightSession("a"))
+        {
+            seed.Store(new GlobalDoc { Id = globalId, Label = "global" });
+            seed.Store(new TenantedDoc { Id = tenantedId, Label = "tenant-a" });
+            await seed.SaveChangesAsync();
+        }
+
+        await using (var seed = theStore.LightweightSession("b"))
+        {
+            seed.Store(new TenantedDoc { Id = tenantedId, Label = "tenant-b" });
+            await seed.SaveChangesAsync();
+        }
+
+        await using var session = theStore.IdentitySession("a");
+
+        var globalViaA = await session.ForTenant("a").LoadAsync<GlobalDoc>(globalId);
+        var tenantedViaA = await session.ForTenant("a").LoadAsync<TenantedDoc>(tenantedId);
+
+        var globalViaB = await session.ForTenant("b").LoadAsync<GlobalDoc>(globalId);
+        var tenantedViaB = await session.ForTenant("b").LoadAsync<TenantedDoc>(tenantedId);
+
+        globalViaB.ShouldBeSameAs(globalViaA);
+
+        tenantedViaA.Label.ShouldBe("tenant-a");
+        tenantedViaB.Label.ShouldBe("tenant-b");
+    }
+}

src/Marten/Internal/Sessions/NestedTenantQuerySession.cs

@@ -28,6 +28,8 @@ internal NestedTenantQuerySession(QuerySession parent, Tenant tenant): base((Doc
 
     protected internal override IDocumentStorage<T> selectStorage<T>(DocumentProvider<T> provider)
     {
-        return _parent.selectStorage(provider);
+        var storage = _parent.selectStorage(provider);
+        this.ShareTenantNeutralStateWith(_parent, storage);
+        return storage;
     }
 }

src/Marten/Internal/Sessions/NestedTenantSession.cs

@@ -78,7 +78,9 @@ protected internal override void resetDirtyChecking()
 
     protected internal override IDocumentStorage<T> selectStorage<T>(DocumentProvider<T> provider)
     {
-        return _parent.selectStorage(provider);
+        var storage = _parent.selectStorage(provider);
+        this.ShareTenantNeutralStateWith(_parent, storage);
+        return storage;
     }
 
     public override void Dispose()

src/Marten/Internal/Sessions/NestedTenantSessionState.cs

@@ -0,0 +1,48 @@
+#nullable enable
+using Marten.Internal.Storage;
+using Marten.Storage;
+
+namespace Marten.Internal.Sessions;
+
+internal static class NestedTenantSessionState
+{
+    /// <summary>
+    ///     #4947 — a nested ForTenant() session keeps its own identity map / version tracker so that
+    ///     conjoined documents (where the same id means a different document per tenant) cannot bleed
+    ///     across tenants (#4801). Tenancy-neutral documents are the opposite case: there is exactly
+    ///     one document per id for the whole database, so the parent session and every ForTenant()
+    ///     view of it must see the very same tracked instance. Alias the parent's identity-map (and
+    ///     version) entry for this document type into the nested session the first time the type is
+    ///     used through the nested session — before any load/store can create a competing entry.
+    /// </summary>
+    internal static void ShareTenantNeutralStateWith<T>(this QuerySession nested, QuerySession parent,
+        IDocumentStorage<T> storage) where T : notnull
+    {
+        // Same tenant as the parent: the whole map/tracker is already shared wholesale
+        if (ReferenceEquals(nested.ItemMap, parent.ItemMap))
+        {
+            return;
+        }
+
+        // Only identity-mapped (identity map + dirty checking) storage tracks documents in the session
+        if (storage is not ISharedTenantNeutralSessionState shared)
+        {
+            return;
+        }
+
+        // Conjoined documents are genuinely per-tenant — keep them isolated (#4801)
+        if (shared.TenancyStyle == TenancyStyle.Conjoined)
+        {
+            return;
+        }
+
+        // Under database-per-tenant the same id in another tenant's database is a *different*
+        // document even for a tenancy-neutral doc type, so only share within one database.
+        if (!ReferenceEquals(nested.Database, parent.Database))
+        {
+            return;
+        }
+
+        shared.ShareTenantNeutralStateWith(parent, nested);
+    }
+}

src/Marten/Internal/Storage/IdentityMapDocumentStorage.cs

@@ -8,11 +8,28 @@
 using Marten.Internal.CodeGeneration;
 using Marten.Linq.Selectors;
 using Marten.Schema;
+using Marten.Storage;
 using Npgsql;
 
 namespace Marten.Internal.Storage;
 
-public abstract class IdentityMapDocumentStorage<T, TId>: DocumentStorage<T, TId> where T: notnull where TId: notnull
+/// <summary>
+///     #4947 — seam that lets a nested ForTenant() session share the parent session's identity map
+///     (and version state) for a *single document type* rather than all-or-nothing per session.
+/// </summary>
+internal interface ISharedTenantNeutralSessionState
+{
+    TenancyStyle TenancyStyle { get; }
+
+    /// <summary>
+    ///     Point the nested session's identity map / version entries for this document type at the
+    ///     very same underlying dictionaries used by the parent session. Only ever called for
+    ///     tenancy-neutral documents living in the same database.
+    /// </summary>
+    void ShareTenantNeutralStateWith(IStorageSession parent, IStorageSession nested);
+}
+
+public abstract class IdentityMapDocumentStorage<T, TId>: DocumentStorage<T, TId>, ISharedTenantNeutralSessionState where T: notnull where TId: notnull
 {
     public IdentityMapDocumentStorage(DocumentMapping document): this(StorageStyle.IdentityMap, document)
     {
@@ -23,6 +40,66 @@ protected IdentityMapDocumentStorage(StorageStyle storageStyle, DocumentMapping
     {
     }
 
+    void ISharedTenantNeutralSessionState.ShareTenantNeutralStateWith(IStorageSession parent, IStorageSession nested)
+    {
+        if (ReferenceEquals(parent.ItemMap, nested.ItemMap))
+        {
+            return;
+        }
+
+        shareIdentityMap(parent, nested);
+
+        if (UseOptimisticConcurrency || UseNumericRevisions)
+        {
+            shareVersions(parent, nested);
+        }
+    }
+
+    private static void shareIdentityMap(IStorageSession parent, IStorageSession nested)
+    {
+        if (parent.ItemMap.TryGetValue(typeof(T), out var items))
+        {
+            // A mismatched key type is diagnosed (and thrown on) by the normal storage paths --
+            // don't propagate a broken entry into the nested session here.
+            if (items is not Dictionary<TId, T>)
+            {
+                return;
+            }
+        }
+        else
+        {
+            items = new Dictionary<TId, T>();
+            parent.ItemMap[typeof(T)] = items;
+        }
+
+        nested.ItemMap[typeof(T)] = items;
+    }
+
+    private void shareVersions(IStorageSession parent, IStorageSession nested)
+    {
+        if (parent.Versions is not VersionTracker parentVersions ||
+            nested.Versions is not VersionTracker nestedVersions)
+        {
+            return;
+        }
+
+        if (ReferenceEquals(parentVersions, nestedVersions))
+        {
+            return;
+        }
+
+        if (!parentVersions.ByType.TryGetValue(typeof(T), out var versions))
+        {
+            versions = UseNumericRevisions
+                ? new Dictionary<TId, long>()
+                : new Dictionary<TId, Guid>();
+
+            parentVersions.ByType[typeof(T)] = versions;
+        }
+
+        nestedVersions.ByType[typeof(T)] = versions;
+    }
+
     public sealed override void Eject(IStorageSession session, T document)
     {
         var id = Identity(document);

src/Marten/Internal/VersionTracker.cs

@@ -8,6 +8,12 @@ public class VersionTracker: IVersionTracker
 {
     private readonly Dictionary<Type, object> _byType = new();
 
+    /// <summary>
+    ///     #4947 -- exposed so that a nested ForTenant() session can alias the parent session's
+    ///     version state for a tenancy-neutral document type instead of tracking its own copy.
+    /// </summary>
+    internal Dictionary<Type, object> ByType => _byType;
+
     public Dictionary<TId, long> RevisionsFor<TDoc, TId>() where TId : notnull
     {
         if (_byType.TryGetValue(typeof(TDoc), out var item))