From c648c671a326f2d626cf34db52cbcc9999b95bba Mon Sep 17 00:00:00 2001
From: Till <2353100+S7evinK@users.noreply.github.com>
Date: Fri, 11 Nov 2022 10:52:43 +0100
Subject: [PATCH] Fix issue with missing user NIDs (#2874)

This should fix #2696 and possibly other related issues regarding
missing user NIDs.
(https://github.com/matrix-org/dendrite/issues/2094?)
---
 roomserver/storage/shared/storage.go | 17 +++++++++++++++--
 1 file changed, 15 insertions(+), 2 deletions(-)

diff --git a/roomserver/storage/shared/storage.go b/roomserver/storage/shared/storage.go
index 4455ec3bf..958787070 100644
--- a/roomserver/storage/shared/storage.go
+++ b/roomserver/storage/shared/storage.go
@@ -110,6 +110,19 @@ func (d *Database) eventStateKeyNIDs(
 	for eventStateKey, nid := range nids {
 		result[eventStateKey] = nid
 	}
+	// We received some nids, but are still missing some, work out which and create them
+	if len(eventStateKeys) < len(result) {
+		for _, eventStateKey := range eventStateKeys {
+			if _, ok := result[eventStateKey]; ok {
+				continue
+			}
+			nid, err := d.assignStateKeyNID(ctx, txn, eventStateKey)
+			if err != nil {
+				return result, err
+			}
+			result[eventStateKey] = nid
+		}
+	}
 	return result, nil
 }
 
@@ -1243,7 +1256,7 @@ func (d *Database) GetBulkStateContent(ctx context.Context, roomIDs []string, tu
 
 	}
 
-	eventStateKeyNIDMap, err := d.EventStateKeysTable.BulkSelectEventStateKeyNID(ctx, nil, eventStateKeys)
+	eventStateKeyNIDMap, err := d.eventStateKeyNIDs(ctx, nil, eventStateKeys)
 	if err != nil {
 		return nil, fmt.Errorf("GetBulkStateContent: failed to map state key nids: %w", err)
 	}
@@ -1309,7 +1322,7 @@ func (d *Database) JoinedUsersSetInRooms(ctx context.Context, roomIDs, userIDs [
 	if err != nil {
 		return nil, err
 	}
-	userNIDsMap, err := d.EventStateKeysTable.BulkSelectEventStateKeyNID(ctx, nil, userIDs)
+	userNIDsMap, err := d.eventStateKeyNIDs(ctx, nil, userIDs)
 	if err != nil {
 		return nil, err
 	}