Use memberships to determine whether to reset latest events/state on room join (#1047)

* Track local/remote memberships, re-scope some input stuff

* Check if we're in the room already before resetting latest events/state

* Fix postgres, fix lint

* Review comments
This commit is contained in:
Neil Alexander 2020-05-20 18:03:06 +01:00 committed by GitHub
parent 6091bf044f
commit f2c07437fe
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
12 changed files with 162 additions and 100 deletions

View File

@ -60,7 +60,7 @@ func (r *RoomserverInternalAPI) InputRoomEvents(
defer r.mutex.Unlock() defer r.mutex.Unlock()
for i := range request.InputInviteEvents { for i := range request.InputInviteEvents {
var loopback *api.InputRoomEvent var loopback *api.InputRoomEvent
if loopback, err = processInviteEvent(ctx, r.DB, r, request.InputInviteEvents[i]); err != nil { if loopback, err = r.processInviteEvent(ctx, r, request.InputInviteEvents[i]); err != nil {
return err return err
} }
// The processInviteEvent function can optionally return a // The processInviteEvent function can optionally return a
@ -71,7 +71,7 @@ func (r *RoomserverInternalAPI) InputRoomEvents(
} }
} }
for i := range request.InputRoomEvents { for i := range request.InputRoomEvents {
if response.EventID, err = processRoomEvent(ctx, r.DB, r, request.InputRoomEvents[i]); err != nil { if response.EventID, err = r.processRoomEvent(ctx, request.InputRoomEvents[i]); err != nil {
return err return err
} }
} }

View File

@ -31,21 +31,13 @@ import (
log "github.com/sirupsen/logrus" log "github.com/sirupsen/logrus"
) )
// OutputRoomEventWriter has the APIs needed to write an event to the output logs.
type OutputRoomEventWriter interface {
// Write a list of events for a room
WriteOutputEvents(roomID string, updates []api.OutputEvent) error
}
// processRoomEvent can only be called once at a time // processRoomEvent can only be called once at a time
// //
// TODO(#375): This should be rewritten to allow concurrent calls. The // TODO(#375): This should be rewritten to allow concurrent calls. The
// difficulty is in ensuring that we correctly annotate events with the correct // difficulty is in ensuring that we correctly annotate events with the correct
// state deltas when sending to kafka streams // state deltas when sending to kafka streams
func processRoomEvent( func (r *RoomserverInternalAPI) processRoomEvent(
ctx context.Context, ctx context.Context,
db storage.Database,
ow OutputRoomEventWriter,
input api.InputRoomEvent, input api.InputRoomEvent,
) (eventID string, err error) { ) (eventID string, err error) {
// Parse and validate the event JSON // Parse and validate the event JSON
@ -54,7 +46,7 @@ func processRoomEvent(
// Check that the event passes authentication checks and work out // Check that the event passes authentication checks and work out
// the numeric IDs for the auth events. // the numeric IDs for the auth events.
authEventNIDs, err := checkAuthEvents(ctx, db, headered, input.AuthEventIDs) authEventNIDs, err := checkAuthEvents(ctx, r.DB, headered, input.AuthEventIDs)
if err != nil { if err != nil {
logrus.WithError(err).WithField("event_id", event.EventID()).WithField("auth_event_ids", input.AuthEventIDs).Error("processRoomEvent.checkAuthEvents failed for event") logrus.WithError(err).WithField("event_id", event.EventID()).WithField("auth_event_ids", input.AuthEventIDs).Error("processRoomEvent.checkAuthEvents failed for event")
return return
@ -63,7 +55,7 @@ func processRoomEvent(
// If we don't have a transaction ID then get one. // If we don't have a transaction ID then get one.
if input.TransactionID != nil { if input.TransactionID != nil {
tdID := input.TransactionID tdID := input.TransactionID
eventID, err = db.GetTransactionEventID( eventID, err = r.DB.GetTransactionEventID(
ctx, tdID.TransactionID, tdID.SessionID, event.Sender(), ctx, tdID.TransactionID, tdID.SessionID, event.Sender(),
) )
// On error OR event with the transaction already processed/processesing // On error OR event with the transaction already processed/processesing
@ -73,7 +65,7 @@ func processRoomEvent(
} }
// Store the event. // Store the event.
roomNID, stateAtEvent, err := db.StoreEvent(ctx, event, input.TransactionID, authEventNIDs) roomNID, stateAtEvent, err := r.DB.StoreEvent(ctx, event, input.TransactionID, authEventNIDs)
if err != nil { if err != nil {
return return
} }
@ -93,16 +85,14 @@ func processRoomEvent(
if stateAtEvent.BeforeStateSnapshotNID == 0 { if stateAtEvent.BeforeStateSnapshotNID == 0 {
// We haven't calculated a state for this event yet. // We haven't calculated a state for this event yet.
// Lets calculate one. // Lets calculate one.
err = calculateAndSetState(ctx, db, input, roomNID, &stateAtEvent, event) err = r.calculateAndSetState(ctx, input, roomNID, &stateAtEvent, event)
if err != nil { if err != nil {
return return
} }
} }
if err = updateLatestEvents( if err = r.updateLatestEvents(
ctx, // context ctx, // context
db, // roomserver database
ow, // output event writer
roomNID, // room NID to update roomNID, // room NID to update
stateAtEvent, // state at event (below) stateAtEvent, // state at event (below)
event, // event event, // event
@ -116,29 +106,36 @@ func processRoomEvent(
return event.EventID(), nil return event.EventID(), nil
} }
func calculateAndSetState( func (r *RoomserverInternalAPI) calculateAndSetState(
ctx context.Context, ctx context.Context,
db storage.Database,
input api.InputRoomEvent, input api.InputRoomEvent,
roomNID types.RoomNID, roomNID types.RoomNID,
stateAtEvent *types.StateAtEvent, stateAtEvent *types.StateAtEvent,
event gomatrixserverlib.Event, event gomatrixserverlib.Event,
) error { ) error {
var err error var err error
roomState := state.NewStateResolution(db) roomState := state.NewStateResolution(r.DB)
if input.HasState { if input.HasState {
// TODO: Check here if we think we're in the room already. // Check here if we think we're in the room already.
stateAtEvent.Overwrite = true stateAtEvent.Overwrite = true
var joinEventNIDs []types.EventNID
// Request join memberships only for local users only.
if joinEventNIDs, err = r.DB.GetMembershipEventNIDsForRoom(ctx, roomNID, true, true); err == nil {
// If we have no local users that are joined to the room then any state about
// the room that we have is quite possibly out of date. Therefore in that case
// we should overwrite it rather than merge it.
stateAtEvent.Overwrite = len(joinEventNIDs) == 0
}
// We've been told what the state at the event is so we don't need to calculate it. // We've been told what the state at the event is so we don't need to calculate it.
// Check that those state events are in the database and store the state. // Check that those state events are in the database and store the state.
var entries []types.StateEntry var entries []types.StateEntry
if entries, err = db.StateEntriesForEventIDs(ctx, input.StateEventIDs); err != nil { if entries, err = r.DB.StateEntriesForEventIDs(ctx, input.StateEventIDs); err != nil {
return err return err
} }
if stateAtEvent.BeforeStateSnapshotNID, err = db.AddState(ctx, roomNID, nil, entries); err != nil { if stateAtEvent.BeforeStateSnapshotNID, err = r.DB.AddState(ctx, roomNID, nil, entries); err != nil {
return err return err
} }
} else { } else {
@ -149,12 +146,11 @@ func calculateAndSetState(
return err return err
} }
} }
return db.SetState(ctx, stateAtEvent.EventNID, stateAtEvent.BeforeStateSnapshotNID) return r.DB.SetState(ctx, stateAtEvent.EventNID, stateAtEvent.BeforeStateSnapshotNID)
} }
func processInviteEvent( func (r *RoomserverInternalAPI) processInviteEvent(
ctx context.Context, ctx context.Context,
db storage.Database,
ow *RoomserverInternalAPI, ow *RoomserverInternalAPI,
input api.InputInviteEvent, input api.InputInviteEvent,
) (*api.InputRoomEvent, error) { ) (*api.InputRoomEvent, error) {
@ -172,7 +168,10 @@ func processInviteEvent(
"target_user_id": targetUserID, "target_user_id": targetUserID,
}).Info("processing invite event") }).Info("processing invite event")
updater, err := db.MembershipUpdater(ctx, roomID, targetUserID, input.RoomVersion) _, domain, _ := gomatrixserverlib.SplitID('@', targetUserID)
isTargetLocalUser := domain == r.Cfg.Matrix.ServerName
updater, err := r.DB.MembershipUpdater(ctx, roomID, targetUserID, isTargetLocalUser, input.RoomVersion)
if err != nil { if err != nil {
return nil, err return nil, err
} }
@ -239,7 +238,7 @@ func processInviteEvent(
// up from local data (which is most likely to be if the event came // up from local data (which is most likely to be if the event came
// from the CS API). If we know about the room then we can insert // from the CS API). If we know about the room then we can insert
// the invite room state, if we don't then we just fail quietly. // the invite room state, if we don't then we just fail quietly.
if irs, ierr := buildInviteStrippedState(ctx, db, input); ierr == nil { if irs, ierr := buildInviteStrippedState(ctx, r.DB, input); ierr == nil {
if err = event.SetUnsignedField("invite_room_state", irs); err != nil { if err = event.SetUnsignedField("invite_room_state", irs); err != nil {
return nil, err return nil, err
} }

View File

@ -23,7 +23,6 @@ import (
"github.com/matrix-org/dendrite/common" "github.com/matrix-org/dendrite/common"
"github.com/matrix-org/dendrite/roomserver/api" "github.com/matrix-org/dendrite/roomserver/api"
"github.com/matrix-org/dendrite/roomserver/state" "github.com/matrix-org/dendrite/roomserver/state"
"github.com/matrix-org/dendrite/roomserver/storage"
"github.com/matrix-org/dendrite/roomserver/types" "github.com/matrix-org/dendrite/roomserver/types"
"github.com/matrix-org/gomatrixserverlib" "github.com/matrix-org/gomatrixserverlib"
"github.com/matrix-org/util" "github.com/matrix-org/util"
@ -46,17 +45,15 @@ import (
// 7 <----- latest // 7 <----- latest
// //
// Can only be called once at a time // Can only be called once at a time
func updateLatestEvents( func (r *RoomserverInternalAPI) updateLatestEvents(
ctx context.Context, ctx context.Context,
db storage.Database,
ow OutputRoomEventWriter,
roomNID types.RoomNID, roomNID types.RoomNID,
stateAtEvent types.StateAtEvent, stateAtEvent types.StateAtEvent,
event gomatrixserverlib.Event, event gomatrixserverlib.Event,
sendAsServer string, sendAsServer string,
transactionID *api.TransactionID, transactionID *api.TransactionID,
) (err error) { ) (err error) {
updater, err := db.GetLatestEventsForUpdate(ctx, roomNID) updater, err := r.DB.GetLatestEventsForUpdate(ctx, roomNID)
if err != nil { if err != nil {
return return
} }
@ -70,9 +67,8 @@ func updateLatestEvents(
u := latestEventsUpdater{ u := latestEventsUpdater{
ctx: ctx, ctx: ctx,
db: db, api: r,
updater: updater, updater: updater,
ow: ow,
roomNID: roomNID, roomNID: roomNID,
stateAtEvent: stateAtEvent, stateAtEvent: stateAtEvent,
event: event, event: event,
@ -94,9 +90,8 @@ func updateLatestEvents(
// when there are so many variables to pass around. // when there are so many variables to pass around.
type latestEventsUpdater struct { type latestEventsUpdater struct {
ctx context.Context ctx context.Context
db storage.Database api *RoomserverInternalAPI
updater types.RoomRecentEventsUpdater updater types.RoomRecentEventsUpdater
ow OutputRoomEventWriter
roomNID types.RoomNID roomNID types.RoomNID
stateAtEvent types.StateAtEvent stateAtEvent types.StateAtEvent
event gomatrixserverlib.Event event gomatrixserverlib.Event
@ -181,7 +176,7 @@ func (u *latestEventsUpdater) doUpdateLatestEvents() error {
// If we need to generate any output events then here's where we do it. // If we need to generate any output events then here's where we do it.
// TODO: Move this! // TODO: Move this!
updates, err := updateMemberships(u.ctx, u.db, u.updater, u.removed, u.added) updates, err := u.api.updateMemberships(u.ctx, u.updater, u.removed, u.added)
if err != nil { if err != nil {
return err return err
} }
@ -200,7 +195,7 @@ func (u *latestEventsUpdater) doUpdateLatestEvents() error {
// send the event asynchronously but we would need to ensure that 1) the events are written to the log in // send the event asynchronously but we would need to ensure that 1) the events are written to the log in
// the correct order, 2) that pending writes are resent across restarts. In order to avoid writing all the // the correct order, 2) that pending writes are resent across restarts. In order to avoid writing all the
// necessary bookkeeping we'll keep the event sending synchronous for now. // necessary bookkeeping we'll keep the event sending synchronous for now.
if err = u.ow.WriteOutputEvents(u.event.RoomID(), updates); err != nil { if err = u.api.WriteOutputEvents(u.event.RoomID(), updates); err != nil {
return err return err
} }
@ -213,7 +208,7 @@ func (u *latestEventsUpdater) doUpdateLatestEvents() error {
func (u *latestEventsUpdater) latestState() error { func (u *latestEventsUpdater) latestState() error {
var err error var err error
roomState := state.NewStateResolution(u.db) roomState := state.NewStateResolution(u.api.DB)
// Get a list of the current latest events. // Get a list of the current latest events.
latestStateAtEvents := make([]types.StateAtEvent, len(u.latest)) latestStateAtEvents := make([]types.StateAtEvent, len(u.latest))
@ -303,7 +298,7 @@ func (u *latestEventsUpdater) makeOutputNewRoomEvent() (*api.OutputEvent, error)
latestEventIDs[i] = u.latest[i].EventID latestEventIDs[i] = u.latest[i].EventID
} }
roomVersion, err := u.db.GetRoomVersionForRoom(u.ctx, u.event.RoomID()) roomVersion, err := u.api.DB.GetRoomVersionForRoom(u.ctx, u.event.RoomID())
if err != nil { if err != nil {
return nil, err return nil, err
} }
@ -329,7 +324,7 @@ func (u *latestEventsUpdater) makeOutputNewRoomEvent() (*api.OutputEvent, error)
stateEventNIDs = append(stateEventNIDs, entry.EventNID) stateEventNIDs = append(stateEventNIDs, entry.EventNID)
} }
stateEventNIDs = stateEventNIDs[:util.SortAndUnique(eventNIDSorter(stateEventNIDs))] stateEventNIDs = stateEventNIDs[:util.SortAndUnique(eventNIDSorter(stateEventNIDs))]
eventIDMap, err := u.db.EventIDs(u.ctx, stateEventNIDs) eventIDMap, err := u.api.DB.EventIDs(u.ctx, stateEventNIDs)
if err != nil { if err != nil {
return nil, err return nil, err
} }

View File

@ -19,7 +19,6 @@ import (
"fmt" "fmt"
"github.com/matrix-org/dendrite/roomserver/api" "github.com/matrix-org/dendrite/roomserver/api"
"github.com/matrix-org/dendrite/roomserver/storage"
"github.com/matrix-org/dendrite/roomserver/types" "github.com/matrix-org/dendrite/roomserver/types"
"github.com/matrix-org/gomatrixserverlib" "github.com/matrix-org/gomatrixserverlib"
) )
@ -28,9 +27,8 @@ import (
// user affected by a change in the current state of the room. // user affected by a change in the current state of the room.
// Returns a list of output events to write to the kafka log to inform the // Returns a list of output events to write to the kafka log to inform the
// consumers about the invites added or retired by the change in current state. // consumers about the invites added or retired by the change in current state.
func updateMemberships( func (r *RoomserverInternalAPI) updateMemberships(
ctx context.Context, ctx context.Context,
db storage.Database,
updater types.RoomRecentEventsUpdater, updater types.RoomRecentEventsUpdater,
removed, added []types.StateEntry, removed, added []types.StateEntry,
) ([]api.OutputEvent, error) { ) ([]api.OutputEvent, error) {
@ -48,7 +46,7 @@ func updateMemberships(
// Load the event JSON so we can look up the "membership" key. // Load the event JSON so we can look up the "membership" key.
// TODO: Maybe add a membership key to the events table so we can load that // TODO: Maybe add a membership key to the events table so we can load that
// key without having to load the entire event JSON? // key without having to load the entire event JSON?
events, err := db.Events(ctx, eventNIDs) events, err := r.DB.Events(ctx, eventNIDs)
if err != nil { if err != nil {
return nil, err return nil, err
} }
@ -71,15 +69,16 @@ func updateMemberships(
ae = &ev.Event ae = &ev.Event
} }
} }
if updates, err = updateMembership(updater, targetUserNID, re, ae, updates); err != nil { if updates, err = r.updateMembership(updater, targetUserNID, re, ae, updates); err != nil {
return nil, err return nil, err
} }
} }
return updates, nil return updates, nil
} }
func updateMembership( func (r *RoomserverInternalAPI) updateMembership(
updater types.RoomRecentEventsUpdater, targetUserNID types.EventStateKeyNID, updater types.RoomRecentEventsUpdater,
targetUserNID types.EventStateKeyNID,
remove, add *gomatrixserverlib.Event, remove, add *gomatrixserverlib.Event,
updates []api.OutputEvent, updates []api.OutputEvent,
) ([]api.OutputEvent, error) { ) ([]api.OutputEvent, error) {
@ -113,7 +112,7 @@ func updateMembership(
return updates, nil return updates, nil
} }
mu, err := updater.MembershipUpdater(targetUserNID) mu, err := updater.MembershipUpdater(targetUserNID, r.isLocalTarget(add))
if err != nil { if err != nil {
return nil, err return nil, err
} }
@ -132,6 +131,15 @@ func updateMembership(
} }
} }
func (r *RoomserverInternalAPI) isLocalTarget(event *gomatrixserverlib.Event) bool {
isTargetLocalUser := false
if statekey := event.StateKey(); statekey != nil {
_, domain, _ := gomatrixserverlib.SplitID('@', *statekey)
isTargetLocalUser = domain == r.Cfg.Matrix.ServerName
}
return isTargetLocalUser
}
func updateToInviteMembership( func updateToInviteMembership(
mu types.MembershipUpdater, add *gomatrixserverlib.Event, updates []api.OutputEvent, mu types.MembershipUpdater, add *gomatrixserverlib.Event, updates []api.OutputEvent,
roomVersion gomatrixserverlib.RoomVersion, roomVersion gomatrixserverlib.RoomVersion,

View File

@ -267,7 +267,7 @@ func (r *RoomserverInternalAPI) QueryMembershipsForRoom(
var stateEntries []types.StateEntry var stateEntries []types.StateEntry
if stillInRoom { if stillInRoom {
var eventNIDs []types.EventNID var eventNIDs []types.EventNID
eventNIDs, err = r.DB.GetMembershipEventNIDsForRoom(ctx, roomNID, request.JoinedOnly) eventNIDs, err = r.DB.GetMembershipEventNIDsForRoom(ctx, roomNID, request.JoinedOnly, false)
if err != nil { if err != nil {
return err return err
} }
@ -591,7 +591,7 @@ func (r *RoomserverInternalAPI) isServerCurrentlyInRoom(ctx context.Context, ser
return false, err return false, err
} }
eventNIDs, err := r.DB.GetMembershipEventNIDsForRoom(ctx, roomNID, true) eventNIDs, err := r.DB.GetMembershipEventNIDsForRoom(ctx, roomNID, true, false)
if err != nil { if err != nil {
return false, err return false, err
} }

View File

@ -297,7 +297,7 @@ func joinEventsFromHistoryVisibility(
if err != nil { if err != nil {
return nil, err return nil, err
} }
joinEventNIDs, err := db.GetMembershipEventNIDsForRoom(ctx, roomNID, true) joinEventNIDs, err := db.GetMembershipEventNIDsForRoom(ctx, roomNID, true, false)
if err != nil { if err != nil {
return nil, err return nil, err
} }

View File

@ -83,9 +83,9 @@ type Database interface {
GetAliasesForRoomID(ctx context.Context, roomID string) ([]string, error) GetAliasesForRoomID(ctx context.Context, roomID string) ([]string, error)
GetCreatorIDForAlias(ctx context.Context, alias string) (string, error) GetCreatorIDForAlias(ctx context.Context, alias string) (string, error)
RemoveRoomAlias(ctx context.Context, alias string) error RemoveRoomAlias(ctx context.Context, alias string) error
MembershipUpdater(ctx context.Context, roomID, targetUserID string, roomVersion gomatrixserverlib.RoomVersion) (types.MembershipUpdater, error) MembershipUpdater(ctx context.Context, roomID, targetUserID string, targetLocal bool, roomVersion gomatrixserverlib.RoomVersion) (types.MembershipUpdater, error)
GetMembership(ctx context.Context, roomNID types.RoomNID, requestSenderUserID string) (membershipEventNID types.EventNID, stillInRoom bool, err error) GetMembership(ctx context.Context, roomNID types.RoomNID, requestSenderUserID string) (membershipEventNID types.EventNID, stillInRoom bool, err error)
GetMembershipEventNIDsForRoom(ctx context.Context, roomNID types.RoomNID, joinOnly bool) ([]types.EventNID, error) GetMembershipEventNIDsForRoom(ctx context.Context, roomNID types.RoomNID, joinOnly bool, localOnly bool) ([]types.EventNID, error)
EventsFromIDs(ctx context.Context, eventIDs []string) ([]types.Event, error) EventsFromIDs(ctx context.Context, eventIDs []string) ([]types.Event, error)
GetRoomVersionForRoom(ctx context.Context, roomID string) (gomatrixserverlib.RoomVersion, error) GetRoomVersionForRoom(ctx context.Context, roomID string) (gomatrixserverlib.RoomVersion, error)
} }

View File

@ -59,6 +59,10 @@ CREATE TABLE IF NOT EXISTS roomserver_membership (
-- This NID is updated if the join event gets updated (e.g. profile update), -- This NID is updated if the join event gets updated (e.g. profile update),
-- or if the user leaves/joins the room. -- or if the user leaves/joins the room.
event_nid BIGINT NOT NULL DEFAULT 0, event_nid BIGINT NOT NULL DEFAULT 0,
-- Local target is true if the target_nid refers to a local user rather than
-- a federated one. This is an optimisation for resetting state on federated
-- room joins.
target_local BOOLEAN NOT NULL DEFAULT false,
UNIQUE (room_nid, target_nid) UNIQUE (room_nid, target_nid)
); );
` `
@ -66,8 +70,8 @@ CREATE TABLE IF NOT EXISTS roomserver_membership (
// Insert a row in to membership table so that it can be locked by the // Insert a row in to membership table so that it can be locked by the
// SELECT FOR UPDATE // SELECT FOR UPDATE
const insertMembershipSQL = "" + const insertMembershipSQL = "" +
"INSERT INTO roomserver_membership (room_nid, target_nid)" + "INSERT INTO roomserver_membership (room_nid, target_nid, target_local)" +
" VALUES ($1, $2)" + " VALUES ($1, $2, $3)" +
" ON CONFLICT DO NOTHING" " ON CONFLICT DO NOTHING"
const selectMembershipFromRoomAndTargetSQL = "" + const selectMembershipFromRoomAndTargetSQL = "" +
@ -78,10 +82,20 @@ const selectMembershipsFromRoomAndMembershipSQL = "" +
"SELECT event_nid FROM roomserver_membership" + "SELECT event_nid FROM roomserver_membership" +
" WHERE room_nid = $1 AND membership_nid = $2" " WHERE room_nid = $1 AND membership_nid = $2"
const selectLocalMembershipsFromRoomAndMembershipSQL = "" +
"SELECT event_nid FROM roomserver_membership" +
" WHERE room_nid = $1 AND membership_nid = $2" +
" AND target_local = true"
const selectMembershipsFromRoomSQL = "" + const selectMembershipsFromRoomSQL = "" +
"SELECT event_nid FROM roomserver_membership" + "SELECT event_nid FROM roomserver_membership" +
" WHERE room_nid = $1" " WHERE room_nid = $1"
const selectLocalMembershipsFromRoomSQL = "" +
"SELECT event_nid FROM roomserver_membership" +
" WHERE room_nid = $1" +
" AND target_local = true"
const selectMembershipForUpdateSQL = "" + const selectMembershipForUpdateSQL = "" +
"SELECT membership_nid FROM roomserver_membership" + "SELECT membership_nid FROM roomserver_membership" +
" WHERE room_nid = $1 AND target_nid = $2 FOR UPDATE" " WHERE room_nid = $1 AND target_nid = $2 FOR UPDATE"
@ -91,12 +105,14 @@ const updateMembershipSQL = "" +
" WHERE room_nid = $1 AND target_nid = $2" " WHERE room_nid = $1 AND target_nid = $2"
type membershipStatements struct { type membershipStatements struct {
insertMembershipStmt *sql.Stmt insertMembershipStmt *sql.Stmt
selectMembershipForUpdateStmt *sql.Stmt selectMembershipForUpdateStmt *sql.Stmt
selectMembershipFromRoomAndTargetStmt *sql.Stmt selectMembershipFromRoomAndTargetStmt *sql.Stmt
selectMembershipsFromRoomAndMembershipStmt *sql.Stmt selectMembershipsFromRoomAndMembershipStmt *sql.Stmt
selectMembershipsFromRoomStmt *sql.Stmt selectLocalMembershipsFromRoomAndMembershipStmt *sql.Stmt
updateMembershipStmt *sql.Stmt selectMembershipsFromRoomStmt *sql.Stmt
selectLocalMembershipsFromRoomStmt *sql.Stmt
updateMembershipStmt *sql.Stmt
} }
func (s *membershipStatements) prepare(db *sql.DB) (err error) { func (s *membershipStatements) prepare(db *sql.DB) (err error) {
@ -110,7 +126,9 @@ func (s *membershipStatements) prepare(db *sql.DB) (err error) {
{&s.selectMembershipForUpdateStmt, selectMembershipForUpdateSQL}, {&s.selectMembershipForUpdateStmt, selectMembershipForUpdateSQL},
{&s.selectMembershipFromRoomAndTargetStmt, selectMembershipFromRoomAndTargetSQL}, {&s.selectMembershipFromRoomAndTargetStmt, selectMembershipFromRoomAndTargetSQL},
{&s.selectMembershipsFromRoomAndMembershipStmt, selectMembershipsFromRoomAndMembershipSQL}, {&s.selectMembershipsFromRoomAndMembershipStmt, selectMembershipsFromRoomAndMembershipSQL},
{&s.selectLocalMembershipsFromRoomAndMembershipStmt, selectLocalMembershipsFromRoomAndMembershipSQL},
{&s.selectMembershipsFromRoomStmt, selectMembershipsFromRoomSQL}, {&s.selectMembershipsFromRoomStmt, selectMembershipsFromRoomSQL},
{&s.selectLocalMembershipsFromRoomStmt, selectLocalMembershipsFromRoomSQL},
{&s.updateMembershipStmt, updateMembershipSQL}, {&s.updateMembershipStmt, updateMembershipSQL},
}.prepare(db) }.prepare(db)
} }
@ -118,9 +136,10 @@ func (s *membershipStatements) prepare(db *sql.DB) (err error) {
func (s *membershipStatements) insertMembership( func (s *membershipStatements) insertMembership(
ctx context.Context, ctx context.Context,
txn *sql.Tx, roomNID types.RoomNID, targetUserNID types.EventStateKeyNID, txn *sql.Tx, roomNID types.RoomNID, targetUserNID types.EventStateKeyNID,
localTarget bool,
) error { ) error {
stmt := common.TxStmt(txn, s.insertMembershipStmt) stmt := common.TxStmt(txn, s.insertMembershipStmt)
_, err := stmt.ExecContext(ctx, roomNID, targetUserNID) _, err := stmt.ExecContext(ctx, roomNID, targetUserNID, localTarget)
return err return err
} }
@ -145,9 +164,15 @@ func (s *membershipStatements) selectMembershipFromRoomAndTarget(
} }
func (s *membershipStatements) selectMembershipsFromRoom( func (s *membershipStatements) selectMembershipsFromRoom(
ctx context.Context, roomNID types.RoomNID, ctx context.Context, roomNID types.RoomNID, localOnly bool,
) (eventNIDs []types.EventNID, err error) { ) (eventNIDs []types.EventNID, err error) {
rows, err := s.selectMembershipsFromRoomStmt.QueryContext(ctx, roomNID) var stmt *sql.Stmt
if localOnly {
stmt = s.selectLocalMembershipsFromRoomStmt
} else {
stmt = s.selectMembershipsFromRoomStmt
}
rows, err := stmt.QueryContext(ctx, roomNID)
if err != nil { if err != nil {
return return
} }
@ -165,10 +190,16 @@ func (s *membershipStatements) selectMembershipsFromRoom(
func (s *membershipStatements) selectMembershipsFromRoomAndMembership( func (s *membershipStatements) selectMembershipsFromRoomAndMembership(
ctx context.Context, ctx context.Context,
roomNID types.RoomNID, membership membershipState, roomNID types.RoomNID, membership membershipState, localOnly bool,
) (eventNIDs []types.EventNID, err error) { ) (eventNIDs []types.EventNID, err error) {
stmt := s.selectMembershipsFromRoomAndMembershipStmt var rows *sql.Rows
rows, err := stmt.QueryContext(ctx, roomNID, membership) var stmt *sql.Stmt
if localOnly {
stmt = s.selectLocalMembershipsFromRoomAndMembershipStmt
} else {
stmt = s.selectMembershipsFromRoomAndMembershipStmt
}
rows, err = stmt.QueryContext(ctx, roomNID, membership)
if err != nil { if err != nil {
return return
} }

View File

@ -459,8 +459,8 @@ func (u *roomRecentEventsUpdater) MarkEventAsSent(eventNID types.EventNID) error
return u.d.statements.updateEventSentToOutput(u.ctx, u.txn, eventNID) return u.d.statements.updateEventSentToOutput(u.ctx, u.txn, eventNID)
} }
func (u *roomRecentEventsUpdater) MembershipUpdater(targetUserNID types.EventStateKeyNID) (types.MembershipUpdater, error) { func (u *roomRecentEventsUpdater) MembershipUpdater(targetUserNID types.EventStateKeyNID, targetLocal bool) (types.MembershipUpdater, error) {
return u.d.membershipUpdaterTxn(u.ctx, u.txn, u.roomNID, targetUserNID) return u.d.membershipUpdaterTxn(u.ctx, u.txn, u.roomNID, targetUserNID, targetLocal)
} }
// RoomNID implements query.RoomserverQueryAPIDB // RoomNID implements query.RoomserverQueryAPIDB
@ -558,7 +558,7 @@ func (d *Database) StateEntriesForTuples(
// MembershipUpdater implements input.RoomEventDatabase // MembershipUpdater implements input.RoomEventDatabase
func (d *Database) MembershipUpdater( func (d *Database) MembershipUpdater(
ctx context.Context, roomID, targetUserID string, ctx context.Context, roomID, targetUserID string,
roomVersion gomatrixserverlib.RoomVersion, targetLocal bool, roomVersion gomatrixserverlib.RoomVersion,
) (types.MembershipUpdater, error) { ) (types.MembershipUpdater, error) {
txn, err := d.db.Begin() txn, err := d.db.Begin()
if err != nil { if err != nil {
@ -581,7 +581,7 @@ func (d *Database) MembershipUpdater(
return nil, err return nil, err
} }
updater, err := d.membershipUpdaterTxn(ctx, txn, roomNID, targetUserNID) updater, err := d.membershipUpdaterTxn(ctx, txn, roomNID, targetUserNID, targetLocal)
if err != nil { if err != nil {
return nil, err return nil, err
} }
@ -603,9 +603,10 @@ func (d *Database) membershipUpdaterTxn(
txn *sql.Tx, txn *sql.Tx,
roomNID types.RoomNID, roomNID types.RoomNID,
targetUserNID types.EventStateKeyNID, targetUserNID types.EventStateKeyNID,
targetLocal bool,
) (types.MembershipUpdater, error) { ) (types.MembershipUpdater, error) {
if err := d.statements.insertMembership(ctx, txn, roomNID, targetUserNID); err != nil { if err := d.statements.insertMembership(ctx, txn, roomNID, targetUserNID, targetLocal); err != nil {
return nil, err return nil, err
} }
@ -748,15 +749,15 @@ func (d *Database) GetMembership(
// GetMembershipEventNIDsForRoom implements query.RoomserverQueryAPIDB // GetMembershipEventNIDsForRoom implements query.RoomserverQueryAPIDB
func (d *Database) GetMembershipEventNIDsForRoom( func (d *Database) GetMembershipEventNIDsForRoom(
ctx context.Context, roomNID types.RoomNID, joinOnly bool, ctx context.Context, roomNID types.RoomNID, joinOnly bool, localOnly bool,
) ([]types.EventNID, error) { ) ([]types.EventNID, error) {
if joinOnly { if joinOnly {
return d.statements.selectMembershipsFromRoomAndMembership( return d.statements.selectMembershipsFromRoomAndMembership(
ctx, roomNID, membershipStateJoin, ctx, roomNID, membershipStateJoin, localOnly,
) )
} }
return d.statements.selectMembershipsFromRoom(ctx, roomNID) return d.statements.selectMembershipsFromRoom(ctx, roomNID, localOnly)
} }
// EventsFromIDs implements query.RoomserverQueryAPIEventDB // EventsFromIDs implements query.RoomserverQueryAPIEventDB

View File

@ -38,6 +38,7 @@ const membershipSchema = `
sender_nid INTEGER NOT NULL DEFAULT 0, sender_nid INTEGER NOT NULL DEFAULT 0,
membership_nid INTEGER NOT NULL DEFAULT 1, membership_nid INTEGER NOT NULL DEFAULT 1,
event_nid INTEGER NOT NULL DEFAULT 0, event_nid INTEGER NOT NULL DEFAULT 0,
target_local BOOLEAN NOT NULL DEFAULT false,
UNIQUE (room_nid, target_nid) UNIQUE (room_nid, target_nid)
); );
` `
@ -45,8 +46,8 @@ const membershipSchema = `
// Insert a row in to membership table so that it can be locked by the // Insert a row in to membership table so that it can be locked by the
// SELECT FOR UPDATE // SELECT FOR UPDATE
const insertMembershipSQL = "" + const insertMembershipSQL = "" +
"INSERT INTO roomserver_membership (room_nid, target_nid)" + "INSERT INTO roomserver_membership (room_nid, target_nid, target_local)" +
" VALUES ($1, $2)" + " VALUES ($1, $2, $3)" +
" ON CONFLICT DO NOTHING" " ON CONFLICT DO NOTHING"
const selectMembershipFromRoomAndTargetSQL = "" + const selectMembershipFromRoomAndTargetSQL = "" +
@ -57,10 +58,20 @@ const selectMembershipsFromRoomAndMembershipSQL = "" +
"SELECT event_nid FROM roomserver_membership" + "SELECT event_nid FROM roomserver_membership" +
" WHERE room_nid = $1 AND membership_nid = $2" " WHERE room_nid = $1 AND membership_nid = $2"
const selectLocalMembershipsFromRoomAndMembershipSQL = "" +
"SELECT event_nid FROM roomserver_membership" +
" WHERE room_nid = $1 AND membership_nid = $2" +
" AND target_local = true"
const selectMembershipsFromRoomSQL = "" + const selectMembershipsFromRoomSQL = "" +
"SELECT event_nid FROM roomserver_membership" + "SELECT event_nid FROM roomserver_membership" +
" WHERE room_nid = $1" " WHERE room_nid = $1"
const selectLocalMembershipsFromRoomSQL = "" +
"SELECT event_nid FROM roomserver_membership" +
" WHERE room_nid = $1" +
" AND target_local = true"
const selectMembershipForUpdateSQL = "" + const selectMembershipForUpdateSQL = "" +
"SELECT membership_nid FROM roomserver_membership" + "SELECT membership_nid FROM roomserver_membership" +
" WHERE room_nid = $1 AND target_nid = $2" " WHERE room_nid = $1 AND target_nid = $2"
@ -70,12 +81,14 @@ const updateMembershipSQL = "" +
" WHERE room_nid = $4 AND target_nid = $5" " WHERE room_nid = $4 AND target_nid = $5"
type membershipStatements struct { type membershipStatements struct {
insertMembershipStmt *sql.Stmt insertMembershipStmt *sql.Stmt
selectMembershipForUpdateStmt *sql.Stmt selectMembershipForUpdateStmt *sql.Stmt
selectMembershipFromRoomAndTargetStmt *sql.Stmt selectMembershipFromRoomAndTargetStmt *sql.Stmt
selectMembershipsFromRoomAndMembershipStmt *sql.Stmt selectMembershipsFromRoomAndMembershipStmt *sql.Stmt
selectMembershipsFromRoomStmt *sql.Stmt selectLocalMembershipsFromRoomAndMembershipStmt *sql.Stmt
updateMembershipStmt *sql.Stmt selectMembershipsFromRoomStmt *sql.Stmt
selectLocalMembershipsFromRoomStmt *sql.Stmt
updateMembershipStmt *sql.Stmt
} }
func (s *membershipStatements) prepare(db *sql.DB) (err error) { func (s *membershipStatements) prepare(db *sql.DB) (err error) {
@ -89,7 +102,9 @@ func (s *membershipStatements) prepare(db *sql.DB) (err error) {
{&s.selectMembershipForUpdateStmt, selectMembershipForUpdateSQL}, {&s.selectMembershipForUpdateStmt, selectMembershipForUpdateSQL},
{&s.selectMembershipFromRoomAndTargetStmt, selectMembershipFromRoomAndTargetSQL}, {&s.selectMembershipFromRoomAndTargetStmt, selectMembershipFromRoomAndTargetSQL},
{&s.selectMembershipsFromRoomAndMembershipStmt, selectMembershipsFromRoomAndMembershipSQL}, {&s.selectMembershipsFromRoomAndMembershipStmt, selectMembershipsFromRoomAndMembershipSQL},
{&s.selectLocalMembershipsFromRoomAndMembershipStmt, selectLocalMembershipsFromRoomAndMembershipSQL},
{&s.selectMembershipsFromRoomStmt, selectMembershipsFromRoomSQL}, {&s.selectMembershipsFromRoomStmt, selectMembershipsFromRoomSQL},
{&s.selectLocalMembershipsFromRoomStmt, selectLocalMembershipsFromRoomSQL},
{&s.updateMembershipStmt, updateMembershipSQL}, {&s.updateMembershipStmt, updateMembershipSQL},
}.prepare(db) }.prepare(db)
} }
@ -97,9 +112,10 @@ func (s *membershipStatements) prepare(db *sql.DB) (err error) {
func (s *membershipStatements) insertMembership( func (s *membershipStatements) insertMembership(
ctx context.Context, txn *sql.Tx, ctx context.Context, txn *sql.Tx,
roomNID types.RoomNID, targetUserNID types.EventStateKeyNID, roomNID types.RoomNID, targetUserNID types.EventStateKeyNID,
localTarget bool,
) error { ) error {
stmt := common.TxStmt(txn, s.insertMembershipStmt) stmt := common.TxStmt(txn, s.insertMembershipStmt)
_, err := stmt.ExecContext(ctx, roomNID, targetUserNID) _, err := stmt.ExecContext(ctx, roomNID, targetUserNID, localTarget)
return err return err
} }
@ -127,9 +143,14 @@ func (s *membershipStatements) selectMembershipFromRoomAndTarget(
func (s *membershipStatements) selectMembershipsFromRoom( func (s *membershipStatements) selectMembershipsFromRoom(
ctx context.Context, txn *sql.Tx, ctx context.Context, txn *sql.Tx,
roomNID types.RoomNID, roomNID types.RoomNID, localOnly bool,
) (eventNIDs []types.EventNID, err error) { ) (eventNIDs []types.EventNID, err error) {
selectStmt := common.TxStmt(txn, s.selectMembershipsFromRoomStmt) var selectStmt *sql.Stmt
if localOnly {
selectStmt = common.TxStmt(txn, s.selectLocalMembershipsFromRoomStmt)
} else {
selectStmt = common.TxStmt(txn, s.selectMembershipsFromRoomStmt)
}
rows, err := selectStmt.QueryContext(ctx, roomNID) rows, err := selectStmt.QueryContext(ctx, roomNID)
if err != nil { if err != nil {
return nil, err return nil, err
@ -145,11 +166,17 @@ func (s *membershipStatements) selectMembershipsFromRoom(
} }
return return
} }
func (s *membershipStatements) selectMembershipsFromRoomAndMembership( func (s *membershipStatements) selectMembershipsFromRoomAndMembership(
ctx context.Context, txn *sql.Tx, ctx context.Context, txn *sql.Tx,
roomNID types.RoomNID, membership membershipState, roomNID types.RoomNID, membership membershipState, localOnly bool,
) (eventNIDs []types.EventNID, err error) { ) (eventNIDs []types.EventNID, err error) {
stmt := common.TxStmt(txn, s.selectMembershipsFromRoomAndMembershipStmt) var stmt *sql.Stmt
if localOnly {
stmt = common.TxStmt(txn, s.selectLocalMembershipsFromRoomAndMembershipStmt)
} else {
stmt = common.TxStmt(txn, s.selectMembershipsFromRoomAndMembershipStmt)
}
rows, err := stmt.QueryContext(ctx, roomNID, membership) rows, err := stmt.QueryContext(ctx, roomNID, membership)
if err != nil { if err != nil {
return return

View File

@ -569,9 +569,9 @@ func (u *roomRecentEventsUpdater) MarkEventAsSent(eventNID types.EventNID) error
return err return err
} }
func (u *roomRecentEventsUpdater) MembershipUpdater(targetUserNID types.EventStateKeyNID) (mu types.MembershipUpdater, err error) { func (u *roomRecentEventsUpdater) MembershipUpdater(targetUserNID types.EventStateKeyNID, targetLocal bool) (mu types.MembershipUpdater, err error) {
err = common.WithTransaction(u.d.db, func(txn *sql.Tx) error { err = common.WithTransaction(u.d.db, func(txn *sql.Tx) error {
mu, err = u.d.membershipUpdaterTxn(u.ctx, txn, u.roomNID, targetUserNID) mu, err = u.d.membershipUpdaterTxn(u.ctx, txn, u.roomNID, targetUserNID, targetLocal)
return err return err
}) })
return return
@ -680,7 +680,7 @@ func (d *Database) StateEntriesForTuples(
// MembershipUpdater implements input.RoomEventDatabase // MembershipUpdater implements input.RoomEventDatabase
func (d *Database) MembershipUpdater( func (d *Database) MembershipUpdater(
ctx context.Context, roomID, targetUserID string, ctx context.Context, roomID, targetUserID string,
roomVersion gomatrixserverlib.RoomVersion, targetLocal bool, roomVersion gomatrixserverlib.RoomVersion,
) (updater types.MembershipUpdater, err error) { ) (updater types.MembershipUpdater, err error) {
var txn *sql.Tx var txn *sql.Tx
txn, err = d.db.Begin() txn, err = d.db.Begin()
@ -716,7 +716,7 @@ func (d *Database) MembershipUpdater(
return nil, err return nil, err
} }
updater, err = d.membershipUpdaterTxn(ctx, txn, roomNID, targetUserNID) updater, err = d.membershipUpdaterTxn(ctx, txn, roomNID, targetUserNID, targetLocal)
if err != nil { if err != nil {
return nil, err return nil, err
} }
@ -738,9 +738,10 @@ func (d *Database) membershipUpdaterTxn(
txn *sql.Tx, txn *sql.Tx,
roomNID types.RoomNID, roomNID types.RoomNID,
targetUserNID types.EventStateKeyNID, targetUserNID types.EventStateKeyNID,
targetLocal bool,
) (types.MembershipUpdater, error) { ) (types.MembershipUpdater, error) {
if err := d.statements.insertMembership(ctx, txn, roomNID, targetUserNID); err != nil { if err := d.statements.insertMembership(ctx, txn, roomNID, targetUserNID, targetLocal); err != nil {
return nil, err return nil, err
} }
@ -896,17 +897,17 @@ func (d *Database) GetMembership(
// GetMembershipEventNIDsForRoom implements query.RoomserverQueryAPIDB // GetMembershipEventNIDsForRoom implements query.RoomserverQueryAPIDB
func (d *Database) GetMembershipEventNIDsForRoom( func (d *Database) GetMembershipEventNIDsForRoom(
ctx context.Context, roomNID types.RoomNID, joinOnly bool, ctx context.Context, roomNID types.RoomNID, joinOnly bool, localOnly bool,
) (eventNIDs []types.EventNID, err error) { ) (eventNIDs []types.EventNID, err error) {
err = common.WithTransaction(d.db, func(txn *sql.Tx) error { err = common.WithTransaction(d.db, func(txn *sql.Tx) error {
if joinOnly { if joinOnly {
eventNIDs, err = d.statements.selectMembershipsFromRoomAndMembership( eventNIDs, err = d.statements.selectMembershipsFromRoomAndMembership(
ctx, txn, roomNID, membershipStateJoin, ctx, txn, roomNID, membershipStateJoin, localOnly,
) )
return nil return nil
} }
eventNIDs, err = d.statements.selectMembershipsFromRoom(ctx, txn, roomNID) eventNIDs, err = d.statements.selectMembershipsFromRoom(ctx, txn, roomNID, localOnly)
return nil return nil
}) })
return return

View File

@ -172,7 +172,7 @@ type RoomRecentEventsUpdater interface {
MarkEventAsSent(eventNID EventNID) error MarkEventAsSent(eventNID EventNID) error
// Build a membership updater for the target user in this room. // Build a membership updater for the target user in this room.
// It will share the same transaction as this updater. // It will share the same transaction as this updater.
MembershipUpdater(targetUserNID EventStateKeyNID) (MembershipUpdater, error) MembershipUpdater(targetUserNID EventStateKeyNID, isTargetLocalUser bool) (MembershipUpdater, error)
// Implements Transaction so it can be committed or rolledback // Implements Transaction so it can be committed or rolledback
common.Transaction common.Transaction
} }