Initial work on simplified state storage

roomserver/storage/postgres/events_table.go

@@ -88,6 +88,12 @@ const bulkSelectStateEventByIDSQL = "" +
 	" WHERE event_id = ANY($1)" +
 	" ORDER BY event_type_nid, event_state_key_nid ASC"
 
+// Bulk lookup of events by numeric ID.
+const bulkSelectStateEventByNIDSQL = "" +
+	"SELECT event_type_nid, event_state_key_nid, event_nid FROM roomserver_events" +
+	" WHERE event_nid = ANY($1)" +
+	" ORDER BY event_type_nid, event_state_key_nid ASC"
+
 const bulkSelectStateAtEventByIDSQL = "" +
 	"SELECT event_type_nid, event_state_key_nid, event_nid, state_snapshot_nid, is_rejected FROM roomserver_events" +
 	" WHERE event_id = ANY($1)"
@@ -127,6 +133,7 @@ type eventStatements struct {
 	insertEventStmt                        *sql.Stmt
 	selectEventStmt                        *sql.Stmt
 	bulkSelectStateEventByIDStmt           *sql.Stmt
+	bulkSelectStateEventByNIDStmt          *sql.Stmt
 	bulkSelectStateAtEventByIDStmt         *sql.Stmt
 	updateEventStateStmt                   *sql.Stmt
 	selectEventSentToOutputStmt            *sql.Stmt
@@ -151,6 +158,7 @@ func NewPostgresEventsTable(db *sql.DB) (tables.Events, error) {
 		{&s.insertEventStmt, insertEventSQL},
 		{&s.selectEventStmt, selectEventSQL},
 		{&s.bulkSelectStateEventByIDStmt, bulkSelectStateEventByIDSQL},
+		{&s.bulkSelectStateEventByNIDStmt, bulkSelectStateEventByNIDSQL},
 		{&s.bulkSelectStateAtEventByIDStmt, bulkSelectStateAtEventByIDSQL},
 		{&s.updateEventStateStmt, updateEventStateSQL},
 		{&s.updateEventSentToOutputStmt, updateEventSentToOutputSQL},
@@ -238,6 +246,48 @@ func (s *eventStatements) BulkSelectStateEventByID(
 	return results, nil
 }
 
+// bulkSelectStateEventByNID lookups a list of state events by internal numeric ID.
+// If any of the requested events are missing from the database it returns a types.MissingEventError
+func (s *eventStatements) BulkSelectStateEventByNID(
+	ctx context.Context, eventNIDs []types.EventNID,
+) ([]types.StateEntry, error) {
+	rows, err := s.bulkSelectStateEventByNIDStmt.QueryContext(ctx, eventNIDsAsArray(eventNIDs))
+	if err != nil {
+		return nil, err
+	}
+	defer internal.CloseAndLogIfError(ctx, rows, "bulkSelectStateEventByID: rows.close() failed")
+	// We know that we will only get as many results as event NIDs
+	// because of the unique constraint on event NIDs.
+	// So we can allocate an array of the correct size now.
+	// We might get fewer results than NIDs so we adjust the length of the slice before returning it.
+	results := make([]types.StateEntry, len(eventNIDs))
+	i := 0
+	for ; rows.Next(); i++ {
+		result := &results[i]
+		if err = rows.Scan(
+			&result.EventTypeNID,
+			&result.EventStateKeyNID,
+			&result.EventNID,
+		); err != nil {
+			return nil, err
+		}
+	}
+	if err = rows.Err(); err != nil {
+		return nil, err
+	}
+	if i != len(eventNIDs) {
+		// If there are fewer rows returned than IDs then we were asked to lookup event IDs we don't have.
+		// We don't know which ones were missing because we don't return the string IDs in the query.
+		// However it should be possible debug this by replaying queries or entries from the input kafka logs.
+		// If this turns out to be impossible and we do need the debug information here, it would be better
+		// to do it as a separate query rather than slowing down/complicating the internal case.
+		return nil, types.MissingEventError(
+			fmt.Sprintf("storage: state event IDs missing from the database (%d != %d)", i, len(eventNIDs)),
+		)
+	}
+	return results, nil
+}
+
 // bulkSelectStateAtEventByID lookups the state at a list of events by event ID.
 // If any of the requested events are missing from the database it returns a types.MissingEventError.
 // If we do not have the state for any of the requested events it returns a types.MissingEventError.

roomserver/storage/postgres/state_block_table.go

@@ -1,292 +0,0 @@
-// Copyright 2017-2018 New Vector Ltd
-// Copyright 2019-2020 The Matrix.org Foundation C.I.C.
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//     http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-package postgres
-
-import (
-	"context"
-	"database/sql"
-	"fmt"
-	"sort"
-
-	"github.com/lib/pq"
-	"github.com/matrix-org/dendrite/internal"
-	"github.com/matrix-org/dendrite/roomserver/storage/shared"
-	"github.com/matrix-org/dendrite/roomserver/storage/tables"
-	"github.com/matrix-org/dendrite/roomserver/types"
-	"github.com/matrix-org/util"
-)
-
-const stateDataSchema = `
--- The state data map.
--- Designed to give enough information to run the state resolution algorithm
--- without hitting the database in the internal case.
--- TODO: Is it worth replacing the unique btree index with a covering index so
--- that postgres could lookup the state using an index-only scan?
--- The type and state_key are included in the index to make it easier to
--- lookup a specific (type, state_key) pair for an event. It also makes it easy
--- to read the state for a given state_block_nid ordered by (type, state_key)
--- which in turn makes it easier to merge state data blocks.
-CREATE SEQUENCE IF NOT EXISTS roomserver_state_block_nid_seq;
-CREATE TABLE IF NOT EXISTS roomserver_state_block (
-    -- Local numeric ID for this state data.
-    state_block_nid bigint NOT NULL,
-    event_type_nid bigint NOT NULL,
-    event_state_key_nid bigint NOT NULL,
-    event_nid bigint NOT NULL,
-    UNIQUE (state_block_nid, event_type_nid, event_state_key_nid)
-);
-`
-
-const insertStateDataSQL = "" +
-	"INSERT INTO roomserver_state_block (state_block_nid, event_type_nid, event_state_key_nid, event_nid)" +
-	" VALUES ($1, $2, $3, $4)"
-
-const selectNextStateBlockNIDSQL = "" +
-	"SELECT nextval('roomserver_state_block_nid_seq')"
-
-// Bulk state lookup by numeric state block ID.
-// Sort by the state_block_nid, event_type_nid, event_state_key_nid
-// This means that all the entries for a given state_block_nid will appear
-// together in the list and those entries will sorted by event_type_nid
-// and event_state_key_nid. This property makes it easier to merge two
-// state data blocks together.
-const bulkSelectStateBlockEntriesSQL = "" +
-	"SELECT state_block_nid, event_type_nid, event_state_key_nid, event_nid" +
-	" FROM roomserver_state_block WHERE state_block_nid = ANY($1)" +
-	" ORDER BY state_block_nid, event_type_nid, event_state_key_nid"
-
-// Bulk state lookup by numeric state block ID.
-// Filters the rows in each block to the requested types and state keys.
-// We would like to restrict to particular type state key pairs but we are
-// restricted by the query language to pull the cross product of a list
-// of types and a list state_keys. So we have to filter the result in the
-// application to restrict it to the list of event types and state keys we
-// actually wanted.
-const bulkSelectFilteredStateBlockEntriesSQL = "" +
-	"SELECT state_block_nid, event_type_nid, event_state_key_nid, event_nid" +
-	" FROM roomserver_state_block WHERE state_block_nid = ANY($1)" +
-	" AND event_type_nid = ANY($2) AND event_state_key_nid = ANY($3)" +
-	" ORDER BY state_block_nid, event_type_nid, event_state_key_nid"
-
-type stateBlockStatements struct {
-	insertStateDataStmt                     *sql.Stmt
-	selectNextStateBlockNIDStmt             *sql.Stmt
-	bulkSelectStateBlockEntriesStmt         *sql.Stmt
-	bulkSelectFilteredStateBlockEntriesStmt *sql.Stmt
-}
-
-func NewPostgresStateBlockTable(db *sql.DB) (tables.StateBlock, error) {
-	s := &stateBlockStatements{}
-	_, err := db.Exec(stateDataSchema)
-	if err != nil {
-		return nil, err
-	}
-
-	return s, shared.StatementList{
-		{&s.insertStateDataStmt, insertStateDataSQL},
-		{&s.selectNextStateBlockNIDStmt, selectNextStateBlockNIDSQL},
-		{&s.bulkSelectStateBlockEntriesStmt, bulkSelectStateBlockEntriesSQL},
-		{&s.bulkSelectFilteredStateBlockEntriesStmt, bulkSelectFilteredStateBlockEntriesSQL},
-	}.Prepare(db)
-}
-
-func (s *stateBlockStatements) BulkInsertStateData(
-	ctx context.Context,
-	txn *sql.Tx,
-	entries []types.StateEntry,
-) (types.StateBlockNID, error) {
-	stateBlockNID, err := s.selectNextStateBlockNID(ctx)
-	if err != nil {
-		return 0, err
-	}
-	for _, entry := range entries {
-		_, err := s.insertStateDataStmt.ExecContext(
-			ctx,
-			int64(stateBlockNID),
-			int64(entry.EventTypeNID),
-			int64(entry.EventStateKeyNID),
-			int64(entry.EventNID),
-		)
-		if err != nil {
-			return 0, err
-		}
-	}
-	return stateBlockNID, nil
-}
-
-func (s *stateBlockStatements) selectNextStateBlockNID(
-	ctx context.Context,
-) (types.StateBlockNID, error) {
-	var stateBlockNID int64
-	err := s.selectNextStateBlockNIDStmt.QueryRowContext(ctx).Scan(&stateBlockNID)
-	return types.StateBlockNID(stateBlockNID), err
-}
-
-func (s *stateBlockStatements) BulkSelectStateBlockEntries(
-	ctx context.Context, stateBlockNIDs []types.StateBlockNID,
-) ([]types.StateEntryList, error) {
-	nids := make([]int64, len(stateBlockNIDs))
-	for i := range stateBlockNIDs {
-		nids[i] = int64(stateBlockNIDs[i])
-	}
-	rows, err := s.bulkSelectStateBlockEntriesStmt.QueryContext(ctx, pq.Int64Array(nids))
-	if err != nil {
-		return nil, err
-	}
-	defer internal.CloseAndLogIfError(ctx, rows, "bulkSelectStateBlockEntries: rows.close() failed")
-
-	results := make([]types.StateEntryList, len(stateBlockNIDs))
-	// current is a pointer to the StateEntryList to append the state entries to.
-	var current *types.StateEntryList
-	i := 0
-	for rows.Next() {
-		var (
-			stateBlockNID    int64
-			eventTypeNID     int64
-			eventStateKeyNID int64
-			eventNID         int64
-			entry            types.StateEntry
-		)
-		if err = rows.Scan(
-			&stateBlockNID, &eventTypeNID, &eventStateKeyNID, &eventNID,
-		); err != nil {
-			return nil, err
-		}
-		entry.EventTypeNID = types.EventTypeNID(eventTypeNID)
-		entry.EventStateKeyNID = types.EventStateKeyNID(eventStateKeyNID)
-		entry.EventNID = types.EventNID(eventNID)
-		if current == nil || types.StateBlockNID(stateBlockNID) != current.StateBlockNID {
-			// The state entry row is for a different state data block to the current one.
-			// So we start appending to the next entry in the list.
-			current = &results[i]
-			current.StateBlockNID = types.StateBlockNID(stateBlockNID)
-			i++
-		}
-		current.StateEntries = append(current.StateEntries, entry)
-	}
-	if err = rows.Err(); err != nil {
-		return nil, err
-	}
-	if i != len(stateBlockNIDs) {
-		return nil, fmt.Errorf("storage: state data NIDs missing from the database (%d != %d)", i, len(stateBlockNIDs))
-	}
-	return results, err
-}
-
-func (s *stateBlockStatements) BulkSelectFilteredStateBlockEntries(
-	ctx context.Context,
-	stateBlockNIDs []types.StateBlockNID,
-	stateKeyTuples []types.StateKeyTuple,
-) ([]types.StateEntryList, error) {
-	tuples := stateKeyTupleSorter(stateKeyTuples)
-	// Sort the tuples so that we can run binary search against them as we filter the rows returned by the db.
-	sort.Sort(tuples)
-
-	eventTypeNIDArray, eventStateKeyNIDArray := tuples.typesAndStateKeysAsArrays()
-	rows, err := s.bulkSelectFilteredStateBlockEntriesStmt.QueryContext(
-		ctx,
-		stateBlockNIDsAsArray(stateBlockNIDs),
-		eventTypeNIDArray,
-		eventStateKeyNIDArray,
-	)
-	if err != nil {
-		return nil, err
-	}
-	defer internal.CloseAndLogIfError(ctx, rows, "bulkSelectFilteredStateBlockEntries: rows.close() failed")
-
-	var results []types.StateEntryList
-	var current types.StateEntryList
-	for rows.Next() {
-		var (
-			stateBlockNID    int64
-			eventTypeNID     int64
-			eventStateKeyNID int64
-			eventNID         int64
-			entry            types.StateEntry
-		)
-		if err := rows.Scan(
-			&stateBlockNID, &eventTypeNID, &eventStateKeyNID, &eventNID,
-		); err != nil {
-			return nil, err
-		}
-		entry.EventTypeNID = types.EventTypeNID(eventTypeNID)
-		entry.EventStateKeyNID = types.EventStateKeyNID(eventStateKeyNID)
-		entry.EventNID = types.EventNID(eventNID)
-
-		// We can use binary search here because we sorted the tuples earlier
-		if !tuples.contains(entry.StateKeyTuple) {
-			// The select will return the cross product of types and state keys.
-			// So we need to check if type of the entry is in the list.
-			continue
-		}
-
-		if types.StateBlockNID(stateBlockNID) != current.StateBlockNID {
-			// The state entry row is for a different state data block to the current one.
-			// So we append the current entry to the results and start adding to a new one.
-			// The first time through the loop current will be empty.
-			if current.StateEntries != nil {
-				results = append(results, current)
-			}
-			current = types.StateEntryList{StateBlockNID: types.StateBlockNID(stateBlockNID)}
-		}
-		current.StateEntries = append(current.StateEntries, entry)
-	}
-	// Add the last entry to the list if it is not empty.
-	if current.StateEntries != nil {
-		results = append(results, current)
-	}
-	return results, rows.Err()
-}
-
-func stateBlockNIDsAsArray(stateBlockNIDs []types.StateBlockNID) pq.Int64Array {
-	nids := make([]int64, len(stateBlockNIDs))
-	for i := range stateBlockNIDs {
-		nids[i] = int64(stateBlockNIDs[i])
-	}
-	return pq.Int64Array(nids)
-}
-
-type stateKeyTupleSorter []types.StateKeyTuple
-
-func (s stateKeyTupleSorter) Len() int           { return len(s) }
-func (s stateKeyTupleSorter) Less(i, j int) bool { return s[i].LessThan(s[j]) }
-func (s stateKeyTupleSorter) Swap(i, j int)      { s[i], s[j] = s[j], s[i] }
-
-// Check whether a tuple is in the list. Assumes that the list is sorted.
-func (s stateKeyTupleSorter) contains(value types.StateKeyTuple) bool {
-	i := sort.Search(len(s), func(i int) bool { return !s[i].LessThan(value) })
-	return i < len(s) && s[i] == value
-}
-
-// List the unique eventTypeNIDs and eventStateKeyNIDs.
-// Assumes that the list is sorted.
-func (s stateKeyTupleSorter) typesAndStateKeysAsArrays() (eventTypeNIDs pq.Int64Array, eventStateKeyNIDs pq.Int64Array) {
-	eventTypeNIDs = make(pq.Int64Array, len(s))
-	eventStateKeyNIDs = make(pq.Int64Array, len(s))
-	for i := range s {
-		eventTypeNIDs[i] = int64(s[i].EventTypeNID)
-		eventStateKeyNIDs[i] = int64(s[i].EventStateKeyNID)
-	}
-	eventTypeNIDs = eventTypeNIDs[:util.SortAndUnique(int64Sorter(eventTypeNIDs))]
-	eventStateKeyNIDs = eventStateKeyNIDs[:util.SortAndUnique(int64Sorter(eventStateKeyNIDs))]
-	return
-}
-
-type int64Sorter []int64
-
-func (s int64Sorter) Len() int           { return len(s) }
-func (s int64Sorter) Less(i, j int) bool { return s[i] < s[j] }
-func (s int64Sorter) Swap(i, j int)      { s[i], s[j] = s[j], s[i] }

roomserver/storage/postgres/state_block_table_test.go

@@ -1,86 +0,0 @@
-// Copyright 2017-2018 New Vector Ltd
-// Copyright 2019-2020 The Matrix.org Foundation C.I.C.
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//     http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-package postgres
-
-import (
-	"sort"
-	"testing"
-
-	"github.com/matrix-org/dendrite/roomserver/types"
-)
-
-func TestStateKeyTupleSorter(t *testing.T) {
-	input := stateKeyTupleSorter{
-		{EventTypeNID: 1, EventStateKeyNID: 2},
-		{EventTypeNID: 1, EventStateKeyNID: 4},
-		{EventTypeNID: 2, EventStateKeyNID: 2},
-		{EventTypeNID: 1, EventStateKeyNID: 1},
-	}
-	want := []types.StateKeyTuple{
-		{EventTypeNID: 1, EventStateKeyNID: 1},
-		{EventTypeNID: 1, EventStateKeyNID: 2},
-		{EventTypeNID: 1, EventStateKeyNID: 4},
-		{EventTypeNID: 2, EventStateKeyNID: 2},
-	}
-	doNotWant := []types.StateKeyTuple{
-		{EventTypeNID: 0, EventStateKeyNID: 0},
-		{EventTypeNID: 1, EventStateKeyNID: 3},
-		{EventTypeNID: 2, EventStateKeyNID: 1},
-		{EventTypeNID: 3, EventStateKeyNID: 1},
-	}
-	wantTypeNIDs := []int64{1, 2}
-	wantStateKeyNIDs := []int64{1, 2, 4}
-
-	// Sort the input and check it's in the right order.
-	sort.Sort(input)
-	gotTypeNIDs, gotStateKeyNIDs := input.typesAndStateKeysAsArrays()
-
-	for i := range want {
-		if input[i] != want[i] {
-			t.Errorf("Wanted %#v at index %d got %#v", want[i], i, input[i])
-		}
-
-		if !input.contains(want[i]) {
-			t.Errorf("Wanted %#v.contains(%#v) to be true but got false", input, want[i])
-		}
-	}
-
-	for i := range doNotWant {
-		if input.contains(doNotWant[i]) {
-			t.Errorf("Wanted %#v.contains(%#v) to be false but got true", input, doNotWant[i])
-		}
-	}
-
-	if len(wantTypeNIDs) != len(gotTypeNIDs) {
-		t.Fatalf("Wanted type NIDs %#v got %#v", wantTypeNIDs, gotTypeNIDs)
-	}
-
-	for i := range wantTypeNIDs {
-		if wantTypeNIDs[i] != gotTypeNIDs[i] {
-			t.Fatalf("Wanted type NIDs %#v got %#v", wantTypeNIDs, gotTypeNIDs)
-		}
-	}
-
-	if len(wantStateKeyNIDs) != len(gotStateKeyNIDs) {
-		t.Fatalf("Wanted state key NIDs %#v got %#v", wantStateKeyNIDs, gotStateKeyNIDs)
-	}
-
-	for i := range wantStateKeyNIDs {
-		if wantStateKeyNIDs[i] != gotStateKeyNIDs[i] {
-			t.Fatalf("Wanted type NIDs %#v got %#v", wantTypeNIDs, gotTypeNIDs)
-		}
-	}
-}

roomserver/storage/postgres/state_snapshot_table.go

@@ -1,126 +0,0 @@
-// Copyright 2017-2018 New Vector Ltd
-// Copyright 2019-2020 The Matrix.org Foundation C.I.C.
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//     http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-package postgres
-
-import (
-	"context"
-	"database/sql"
-	"fmt"
-
-	"github.com/lib/pq"
-	"github.com/matrix-org/dendrite/internal/sqlutil"
-	"github.com/matrix-org/dendrite/roomserver/storage/shared"
-	"github.com/matrix-org/dendrite/roomserver/storage/tables"
-	"github.com/matrix-org/dendrite/roomserver/types"
-)
-
-const stateSnapshotSchema = `
--- The state of a room before an event.
--- Stored as a list of state_block entries stored in a separate table.
--- The actual state is constructed by combining all the state_block entries
--- referenced by state_block_nids together. If the same state key tuple appears
--- multiple times then the entry from the later state_block clobbers the earlier
--- entries.
--- This encoding format allows us to implement a delta encoding which is useful
--- because room state tends to accumulate small changes over time. Although if
--- the list of deltas becomes too long it becomes more efficient to encode
--- the full state under single state_block_nid.
-CREATE SEQUENCE IF NOT EXISTS roomserver_state_snapshot_nid_seq;
-CREATE TABLE IF NOT EXISTS roomserver_state_snapshots (
-    -- Local numeric ID for the state.
-    state_snapshot_nid bigint PRIMARY KEY DEFAULT nextval('roomserver_state_snapshot_nid_seq'),
-    -- Local numeric ID of the room this state is for.
-    -- Unused in normal operation, but useful for background work or ad-hoc debugging.
-    room_nid bigint NOT NULL,
-    -- List of state_block_nids, stored sorted by state_block_nid.
-    state_block_nids bigint[] NOT NULL
-);
-`
-
-const insertStateSQL = "" +
-	"INSERT INTO roomserver_state_snapshots (room_nid, state_block_nids)" +
-	" VALUES ($1, $2)" +
-	" RETURNING state_snapshot_nid"
-
-// Bulk state data NID lookup.
-// Sorting by state_snapshot_nid means we can use binary search over the result
-// to lookup the state data NIDs for a state snapshot NID.
-const bulkSelectStateBlockNIDsSQL = "" +
-	"SELECT state_snapshot_nid, state_block_nids FROM roomserver_state_snapshots" +
-	" WHERE state_snapshot_nid = ANY($1) ORDER BY state_snapshot_nid ASC"
-
-type stateSnapshotStatements struct {
-	insertStateStmt              *sql.Stmt
-	bulkSelectStateBlockNIDsStmt *sql.Stmt
-}
-
-func NewPostgresStateSnapshotTable(db *sql.DB) (tables.StateSnapshot, error) {
-	s := &stateSnapshotStatements{}
-	_, err := db.Exec(stateSnapshotSchema)
-	if err != nil {
-		return nil, err
-	}
-
-	return s, shared.StatementList{
-		{&s.insertStateStmt, insertStateSQL},
-		{&s.bulkSelectStateBlockNIDsStmt, bulkSelectStateBlockNIDsSQL},
-	}.Prepare(db)
-}
-
-func (s *stateSnapshotStatements) InsertState(
-	ctx context.Context, txn *sql.Tx, roomNID types.RoomNID, stateBlockNIDs []types.StateBlockNID,
-) (stateNID types.StateSnapshotNID, err error) {
-	nids := make([]int64, len(stateBlockNIDs))
-	for i := range stateBlockNIDs {
-		nids[i] = int64(stateBlockNIDs[i])
-	}
-	err = sqlutil.TxStmt(txn, s.insertStateStmt).QueryRowContext(ctx, int64(roomNID), pq.Int64Array(nids)).Scan(&stateNID)
-	return
-}
-
-func (s *stateSnapshotStatements) BulkSelectStateBlockNIDs(
-	ctx context.Context, stateNIDs []types.StateSnapshotNID,
-) ([]types.StateBlockNIDList, error) {
-	nids := make([]int64, len(stateNIDs))
-	for i := range stateNIDs {
-		nids[i] = int64(stateNIDs[i])
-	}
-	rows, err := s.bulkSelectStateBlockNIDsStmt.QueryContext(ctx, pq.Int64Array(nids))
-	if err != nil {
-		return nil, err
-	}
-	defer rows.Close() // nolint: errcheck
-	results := make([]types.StateBlockNIDList, len(stateNIDs))
-	i := 0
-	for ; rows.Next(); i++ {
-		result := &results[i]
-		var stateBlockNIDs pq.Int64Array
-		if err = rows.Scan(&result.StateSnapshotNID, &stateBlockNIDs); err != nil {
-			return nil, err
-		}
-		result.StateBlockNIDs = make([]types.StateBlockNID, len(stateBlockNIDs))
-		for k := range stateBlockNIDs {
-			result.StateBlockNIDs[k] = types.StateBlockNID(stateBlockNIDs[k])
-		}
-	}
-	if err = rows.Err(); err != nil {
-		return nil, err
-	}
-	if i != len(stateNIDs) {
-		return nil, fmt.Errorf("storage: state NIDs missing from the database (%d != %d)", i, len(stateNIDs))
-	}
-	return results, nil
-}

roomserver/storage/postgres/state_table.go

@@ -0,0 +1,124 @@
+// Copyright 2017-2018 New Vector Ltd
+// Copyright 2019-2020 The Matrix.org Foundation C.I.C.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package postgres
+
+import (
+	"context"
+	"database/sql"
+	"fmt"
+
+	"github.com/lib/pq"
+	"github.com/matrix-org/dendrite/internal"
+	"github.com/matrix-org/dendrite/internal/sqlutil"
+	"github.com/matrix-org/dendrite/roomserver/storage/shared"
+	"github.com/matrix-org/dendrite/roomserver/storage/tables"
+	"github.com/matrix-org/dendrite/roomserver/types"
+)
+
+const stateSchema = `
+CREATE TABLE IF NOT EXISTS roomserver_state (
+    state_nid BIGSERIAL PRIMARY KEY,
+	room_nid bigint NOT NULL,
+    event_nids bigint[] NOT NULL,
+    UNIQUE (room_nid, event_nids),
+	CONSTRAINT fk_room_id FOREIGN KEY(room_nid) REFERENCES roomserver_rooms(room_nid)
+);
+`
+
+const insertNewStateSnapshotSQL = "" +
+	"INSERT INTO roomserver_state (room_nid, event_nids)" +
+	" VALUES ($1, $2)" +
+	" ON CONFLICT (room_nid, event_nids) DO UPDATE SET room_nid = $1" +
+	" RETURNING state_nid"
+
+const bulkSelectNewStateSnapshotSQL = "" +
+	"SELECT state_nid, event_nids" +
+	" FROM roomserver_state WHERE state_nid = ANY($1)" +
+	" ORDER BY state_nid"
+
+type stateStatements struct {
+	insertStateStmt     *sql.Stmt
+	bulkSelectStateStmt *sql.Stmt
+}
+
+func NewPostgresStateTable(db *sql.DB) (tables.State, error) {
+	s := &stateStatements{}
+	_, err := db.Exec(stateSchema)
+	if err != nil {
+		return nil, err
+	}
+
+	return s, shared.StatementList{
+		{&s.insertStateStmt, insertNewStateSnapshotSQL},
+		{&s.bulkSelectStateStmt, bulkSelectNewStateSnapshotSQL},
+	}.Prepare(db)
+}
+
+func (s *stateStatements) InsertState(
+	ctx context.Context,
+	txn *sql.Tx,
+	roomNID types.RoomNID,
+	eventNIDs []types.EventNID,
+) (types.StateSnapshotNID, error) {
+	stmt := sqlutil.TxStmt(txn, s.insertStateStmt)
+	var id int64
+	var err error
+	eventNIDs = types.DeduplicateEventNIDs(eventNIDs)
+	if err = stmt.QueryRowContext(ctx, int64(roomNID), eventNIDsAsArray(eventNIDs)).Scan(&id); err != nil {
+		return 0, fmt.Errorf("stmt.ExecContext: %w", err)
+	}
+	return types.StateSnapshotNID(id), err
+}
+
+func (s *stateStatements) BulkSelectState(
+	ctx context.Context, stateNIDs []types.StateSnapshotNID,
+) (map[types.StateSnapshotNID][]types.EventNID, error) {
+	rows, err := s.bulkSelectStateStmt.QueryContext(ctx, stateSnapshotNIDsAsArray(stateNIDs))
+	if err != nil {
+		return nil, err
+	}
+	defer internal.CloseAndLogIfError(ctx, rows, "bulkSelectStateBlockEntries: rows.close() failed")
+
+	results := map[types.StateSnapshotNID][]types.EventNID{}
+	for rows.Next() {
+		var stateNID int64
+		var eventNIDs pq.Int64Array
+		if err = rows.Scan(&stateNID, &eventNIDs); err != nil {
+			return nil, err
+		}
+		for _, id := range eventNIDs {
+			results[types.StateSnapshotNID(stateNID)] = append(
+				results[types.StateSnapshotNID(stateNID)],
+				types.EventNID(id),
+			)
+		}
+	}
+	if err = rows.Err(); err != nil {
+		return nil, fmt.Errorf("rows.Err: %w", err)
+	}
+	if len(results) != len(stateNIDs) {
+		return nil, fmt.Errorf("storage: state data NIDs missing from the database (%d != %d)", len(results), len(stateNIDs))
+	}
+	return results, err
+}
+
+func stateSnapshotNIDsAsArray(stateSnapshotNIDs []types.StateSnapshotNID) pq.Int64Array {
+	nids := make([]int64, len(stateSnapshotNIDs))
+	for i := range stateSnapshotNIDs {
+		nids[i] = int64(stateSnapshotNIDs[i])
+	}
+	return nids
+}

roomserver/storage/postgres/storage.go

@@ -86,11 +86,17 @@ func (d *Database) prepare(db *sql.DB, cache caching.RoomServerCaches) (err erro
 	if err != nil {
 		return err
 	}
-	stateBlock, err := NewPostgresStateBlockTable(db)
-	if err != nil {
-		return err
-	}
-	stateSnapshot, err := NewPostgresStateSnapshotTable(db)
+	/*
+		stateBlock, err := NewPostgresStateBlockTable(db)
+		if err != nil {
+			return err
+		}
+		stateSnapshot, err := NewPostgresStateSnapshotTable(db)
+		if err != nil {
+			return err
+		}
+	*/
+	state, err := NewPostgresStateTable(db)
 	if err != nil {
 		return err
 	}
@@ -128,14 +134,15 @@ func (d *Database) prepare(db *sql.DB, cache caching.RoomServerCaches) (err erro
 		EventsTable:         events,
 		RoomsTable:          rooms,
 		TransactionsTable:   transactions,
-		StateBlockTable:     stateBlock,
-		StateSnapshotTable:  stateSnapshot,
-		PrevEventsTable:     prevEvents,
-		RoomAliasesTable:    roomAliases,
-		InvitesTable:        invites,
-		MembershipTable:     membership,
-		PublishedTable:      published,
-		RedactionsTable:     redactions,
+		StateTable:          state,
+		//	StateBlockTable:     stateBlock,
+		//	StateSnapshotTable:  stateSnapshot,
+		PrevEventsTable:  prevEvents,
+		RoomAliasesTable: roomAliases,
+		InvitesTable:     invites,
+		MembershipTable:  membership,
+		PublishedTable:   published,
+		RedactionsTable:  redactions,
 	}
 	return nil
 }