diff --git a/federationsender/queue/destinationqueue.go b/federationsender/queue/destinationqueue.go index 237e50829..76a237a47 100644 --- a/federationsender/queue/destinationqueue.go +++ b/federationsender/queue/destinationqueue.go @@ -33,11 +33,7 @@ import ( ) const ( - maxPDUsPerTransaction = 50 - maxEDUsPerTransaction = 50 - maxPDUsInMemory = 128 - maxEDUsInMemory = 128 - queueIdleTimeout = time.Second * 30 + queueIdleTimeout = time.Second * 30 ) // destinationQueue is a queue of events for a single destination. @@ -45,23 +41,24 @@ const ( // ensures that only one request is in flight to a given destination // at a time. type destinationQueue struct { - db storage.Database - signing *SigningInfo - rsAPI api.RoomserverInternalAPI - client *gomatrixserverlib.FederationClient // federation client - origin gomatrixserverlib.ServerName // origin of requests - destination gomatrixserverlib.ServerName // destination of requests - running atomic.Bool // is the queue worker running? - backingOff atomic.Bool // true if we're backing off - overflowed atomic.Bool // the queues exceed maxPDUsInMemory/maxEDUsInMemory, so we should consult the database for more - statistics *statistics.ServerStatistics // statistics about this remote server - transactionIDMutex sync.Mutex // protects transactionID - transactionID gomatrixserverlib.TransactionID // last transaction ID if retrying, or "" if last txn was successful - notify chan struct{} // interrupts idle wait pending PDUs/EDUs - pendingPDUs []*queuedPDU // PDUs waiting to be sent - pendingEDUs []*queuedEDU // EDUs waiting to be sent - pendingMutex sync.RWMutex // protects pendingPDUs and pendingEDUs - interruptBackoff chan bool // interrupts backoff + db storage.Database + signing *SigningInfo + rsAPI api.RoomserverInternalAPI + client *gomatrixserverlib.FederationClient // federation client + origin gomatrixserverlib.ServerName // origin of requests + destination gomatrixserverlib.ServerName // destination of requests + running atomic.Bool // is the queue worker running? + backingOff atomic.Bool // true if we're backing off + overflowed atomic.Bool // the queues exceed maxPDUsInMemory/maxEDUsInMemory, so we should consult the database for more + transactionID gomatrixserverlib.TransactionID // the ID to commit to the database with + transactionIDMutex sync.RWMutex // protects transactionID + transactionPDUCount atomic.Int32 // how many PDUs in database transaction? + transactionEDUCount atomic.Int32 // how many EDUs in database transaction? + statistics *statistics.ServerStatistics // statistics about this remote server + notify chan struct{} // interrupts idle wait pending PDUs/EDUs + pendingTransactions queuedTransactions // transactions waiting to be sent + pendingMutex sync.RWMutex // protects pendingTransactions + interruptBackoff chan bool // interrupts backoff } // Send event adds the event to the pending queue for the destination. @@ -72,39 +69,55 @@ func (oq *destinationQueue) sendEvent(event *gomatrixserverlib.HeaderedEvent, re log.Errorf("attempt to send nil PDU with destination %q", oq.destination) return } + + // Try to queue the PDU up in memory. If there was enough free + // space then we'll get a transaction ID back. + oq.pendingMutex.Lock() + transactionID := oq.pendingTransactions.queuePDUs(receipt, event) + oq.pendingMutex.Unlock() + + // Check if we got a transaction ID back. + if transactionID == "" { + // If we hit this point then we weren't able to fit the event + // into the memory cache, therefore we need to generate a new + // transaction ID to commit to the database. If we don't have + // a transaction ID for the database, or we've exceeded the + // number of PDUs we can fit in the last one, generate a new + // one. + oq.transactionIDMutex.Lock() + if oq.transactionID == "" || oq.transactionPDUCount.Load() > maxPDUsPerTransaction { + now := gomatrixserverlib.AsTimestamp(time.Now()) + oq.transactionID = gomatrixserverlib.TransactionID(fmt.Sprintf("%d-%d", now, oq.statistics.SuccessCount())) + oq.transactionPDUCount.Store(0) + oq.transactionEDUCount.Store(0) + transactionID = oq.transactionID + } + oq.transactionIDMutex.Unlock() + oq.overflowed.Store(true) + } + // Create a database entry that associates the given PDU NID with // this destination queue. We'll then be able to retrieve the PDU // later. if err := oq.db.AssociatePDUWithDestination( context.TODO(), - "", // TODO: remove this, as we don't need to persist the transaction ID + transactionID, // the transaction ID oq.destination, // the destination server name receipt, // NIDs from federationsender_queue_json table ); err != nil { log.WithError(err).Errorf("failed to associate PDU %q with destination %q", event.EventID(), oq.destination) return } - // Check if the destination is blacklisted. If it isn't then wake - // up the queue. - if !oq.statistics.Blacklisted() { - // If there's room in memory to hold the event then add it to the - // list. - oq.pendingMutex.Lock() - if len(oq.pendingPDUs) < maxPDUsInMemory { - oq.pendingPDUs = append(oq.pendingPDUs, &queuedPDU{ - pdu: event, - receipt: receipt, - }) - } else { - oq.overflowed.Store(true) - } - oq.pendingMutex.Unlock() - // Wake up the queue if it's asleep. - oq.wakeQueueIfNeeded() - select { - case oq.notify <- struct{}{}: - default: - } + + // We've successfully added a PDU to the transaction so increase + // the counter. + oq.transactionPDUCount.Add(1) + + // Wake up the queue. + oq.wakeQueueIfNeeded() + select { + case oq.notify <- struct{}{}: + default: } } @@ -116,38 +129,57 @@ func (oq *destinationQueue) sendEDU(event *gomatrixserverlib.EDU, receipt *share log.Errorf("attempt to send nil EDU with destination %q", oq.destination) return } - // Create a database entry that associates the given PDU NID with - // this destination queue. We'll then be able to retrieve the PDU + + // Try to queue the PDU up in memory. If there was enough free + // space then we'll get a transaction ID back. + oq.pendingMutex.Lock() + transactionID := oq.pendingTransactions.queueEDUs(receipt, event) + oq.pendingMutex.Unlock() + + // Check if we got a transaction ID back. + if transactionID == "" { + // If we hit this point then we weren't able to fit the event + // into the memory cache, therefore we need to generate a new + // transaction ID to commit to the database. If we don't have + // a transaction ID for the database, or we've exceeded the + // number of PDUs we can fit in the last one, generate a new + // one. + /* + oq.transactionIDMutex.Lock() + if oq.transactionID == "" || oq.transactionPDUCount.Load() > maxPDUsPerTransaction { + now := gomatrixserverlib.AsTimestamp(time.Now()) + oq.transactionID = gomatrixserverlib.TransactionID(fmt.Sprintf("%d-%d", now, oq.statistics.SuccessCount())) + oq.transactionPDUCount.Store(0) + oq.transactionEDUCount.Store(0) + transactionID = oq.transactionID + } + oq.transactionIDMutex.Unlock() + */ + oq.overflowed.Store(true) + } + + // Create a database entry that associates the given EDU NID with + // this destination queue. We'll then be able to retrieve the EDU // later. if err := oq.db.AssociateEDUWithDestination( context.TODO(), + //transactionID, // the transaction ID oq.destination, // the destination server name receipt, // NIDs from federationsender_queue_json table ); err != nil { log.WithError(err).Errorf("failed to associate EDU with destination %q", oq.destination) return } - // Check if the destination is blacklisted. If it isn't then wake - // up the queue. - if !oq.statistics.Blacklisted() { - // If there's room in memory to hold the event then add it to the - // list. - oq.pendingMutex.Lock() - if len(oq.pendingEDUs) < maxEDUsInMemory { - oq.pendingEDUs = append(oq.pendingEDUs, &queuedEDU{ - edu: event, - receipt: receipt, - }) - } else { - oq.overflowed.Store(true) - } - oq.pendingMutex.Unlock() - // Wake up the queue if it's asleep. - oq.wakeQueueIfNeeded() - select { - case oq.notify <- struct{}{}: - default: - } + + // We've successfully added a PDU to the transaction so increase + // the counter. + oq.transactionEDUCount.Add(1) + + // Wake up the queue. + oq.wakeQueueIfNeeded() + select { + case oq.notify <- struct{}{}: + default: } } @@ -167,66 +199,33 @@ func (oq *destinationQueue) wakeQueueIfNeeded() { } } -// getPendingFromDatabase will look at the database and see if -// there are any persisted events that haven't been sent to this +// getNextTransactionFromDatabase will look at the database and see +// if there are any persisted events that haven't been sent to this // destination yet. If so, they will be queued up. // nolint:gocyclo -func (oq *destinationQueue) getPendingFromDatabase() { +func (oq *destinationQueue) getNextTransactionFromDatabase() { // Check to see if there's anything to do for this server // in the database. - retrieved := false ctx := context.Background() oq.pendingMutex.Lock() defer oq.pendingMutex.Unlock() - // Take a note of all of the PDUs and EDUs that we already - // have cached. We will index them based on the receipt, - // which ultimately just contains the index of the PDU/EDU - // in the database. - gotPDUs := map[string]struct{}{} - gotEDUs := map[string]struct{}{} - for _, pdu := range oq.pendingPDUs { - gotPDUs[pdu.receipt.String()] = struct{}{} - } - for _, edu := range oq.pendingEDUs { - gotEDUs[edu.receipt.String()] = struct{}{} + transactionID, pdus, pduReceipt, err := oq.db.GetNextTransactionPDUs(ctx, oq.destination, maxPDUsPerTransaction) + if err != nil { + logrus.WithError(err).Errorf("Failed to get pending PDUs for %q", oq.destination) } - if pduCapacity := maxPDUsInMemory - len(oq.pendingPDUs); pduCapacity > 0 { - // We have room in memory for some PDUs - let's request no more than that. - if pdus, err := oq.db.GetPendingPDUs(ctx, oq.destination, pduCapacity); err == nil { - for receipt, pdu := range pdus { - if _, ok := gotPDUs[receipt.String()]; ok { - continue - } - oq.pendingPDUs = append(oq.pendingPDUs, &queuedPDU{receipt, pdu}) - retrieved = true - } - } else { - logrus.WithError(err).Errorf("Failed to get pending PDUs for %q", oq.destination) - } - } - if eduCapacity := maxEDUsInMemory - len(oq.pendingEDUs); eduCapacity > 0 { - // We have room in memory for some EDUs - let's request no more than that. - if edus, err := oq.db.GetPendingEDUs(ctx, oq.destination, eduCapacity); err == nil { - for receipt, edu := range edus { - if _, ok := gotEDUs[receipt.String()]; ok { - continue - } - oq.pendingEDUs = append(oq.pendingEDUs, &queuedEDU{receipt, edu}) - retrieved = true - } - } else { - logrus.WithError(err).Errorf("Failed to get pending EDUs for %q", oq.destination) - } - } - // If we've retrieved all of the events from the database with room to spare - // in memory then we'll no longer consider this queue to be overflowed. - if len(oq.pendingPDUs) < maxPDUsInMemory && len(oq.pendingEDUs) < maxEDUsInMemory { - oq.overflowed.Store(false) + edus, eduReceipt, err := oq.db.GetNextTransactionEDUs(ctx, oq.destination, maxEDUsPerTransaction) + if err != nil { + logrus.WithError(err).Errorf("Failed to get pending EDUs for %q", oq.destination) } + + oq.pendingTransactions.createNew(transactionID) + oq.pendingTransactions.queuePDUs(pduReceipt, pdus...) + oq.pendingTransactions.queueEDUs(eduReceipt, edus...) + // If we've retrieved some events then notify the destination queue goroutine. - if retrieved { + if len(pdus) > 0 || len(edus) > 0 { select { case oq.notify <- struct{}{}: default: @@ -252,7 +251,7 @@ func (oq *destinationQueue) backgroundSend() { // If we are overflowing memory and have sent things out to the // database then we can look up what those things are. if oq.overflowed.Load() { - oq.getPendingFromDatabase() + oq.getNextTransactionFromDatabase() } // If we have nothing to do then wait either for incoming events, or @@ -279,14 +278,10 @@ func (oq *destinationQueue) backgroundSend() { // buffers at this point. The PDU clean-up is already on a defer. log.Warnf("Blacklisting %q due to exceeding backoff threshold", oq.destination) oq.pendingMutex.Lock() - for i := range oq.pendingPDUs { - oq.pendingPDUs[i] = nil + for i := range oq.pendingTransactions.queue { + oq.pendingTransactions.queue[i] = nil } - for i := range oq.pendingEDUs { - oq.pendingEDUs[i] = nil - } - oq.pendingPDUs = nil - oq.pendingEDUs = nil + oq.pendingTransactions.queue = nil oq.pendingMutex.Lock() return } @@ -303,21 +298,15 @@ func (oq *destinationQueue) backgroundSend() { // Work out which PDUs/EDUs to include in the next transaction. oq.pendingMutex.RLock() - pduCount := len(oq.pendingPDUs) - eduCount := len(oq.pendingEDUs) - if pduCount > maxPDUsPerTransaction { - pduCount = maxPDUsPerTransaction + if len(oq.pendingTransactions.queue) == 0 { + continue } - if eduCount > maxEDUsPerTransaction { - eduCount = maxEDUsPerTransaction - } - toSendPDUs := oq.pendingPDUs[:pduCount] - toSendEDUs := oq.pendingEDUs[:eduCount] + next := oq.pendingTransactions.queue[0] oq.pendingMutex.RUnlock() // If we have pending PDUs or EDUs then construct a transaction. // Try sending the next transaction and see what happens. - transaction, pc, ec, terr := oq.nextTransaction(toSendPDUs, toSendEDUs) + transaction, _, _, terr := oq.nextTransaction(next) if terr != nil { // We failed to send the transaction. Mark it as a failure. oq.statistics.Failure() @@ -327,14 +316,13 @@ func (oq *destinationQueue) backgroundSend() { // the pending events and EDUs, and wipe our transaction ID. oq.statistics.Success() oq.pendingMutex.Lock() - for i := range oq.pendingPDUs[:pc] { - oq.pendingPDUs[i] = nil + for i := range next.pdus { + next.pdus[i] = nil } - for i := range oq.pendingEDUs[:ec] { - oq.pendingEDUs[i] = nil + for i := range next.edus { + next.edus[i] = nil } - oq.pendingPDUs = oq.pendingPDUs[pc:] - oq.pendingEDUs = oq.pendingEDUs[ec:] + oq.pendingTransactions.queue = oq.pendingTransactions.queue[1:] oq.pendingMutex.Unlock() } } @@ -344,10 +332,7 @@ func (oq *destinationQueue) backgroundSend() { // queue and sends it. Returns true if a transaction was sent or // false otherwise. // nolint:gocyclo -func (oq *destinationQueue) nextTransaction( - pdus []*queuedPDU, - edus []*queuedEDU, -) (bool, int, int, error) { +func (oq *destinationQueue) nextTransaction(transaction *queuedTransaction) (bool, int, int, error) { // If there's no projected transaction ID then generate one. If // the transaction succeeds then we'll set it back to "" so that // we generate a new one next time. If it fails, we'll preserve @@ -371,32 +356,27 @@ func (oq *destinationQueue) nextTransaction( // If we didn't get anything from the database and there are no // pending EDUs then there's nothing to do - stop here. - if len(pdus) == 0 && len(edus) == 0 { + if len(transaction.pdus) == 0 && len(transaction.edus) == 0 { return false, 0, 0, nil } - var pduReceipts []*shared.Receipt - var eduReceipts []*shared.Receipt - // Go through PDUs that we retrieved from the database, if any, // and add them into the transaction. - for _, pdu := range pdus { - if pdu == nil || pdu.pdu == nil { + for _, pdu := range transaction.pdus { + if pdu == nil { continue } // Append the JSON of the event, since this is a json.RawMessage type in the // gomatrixserverlib.Transaction struct - t.PDUs = append(t.PDUs, pdu.pdu.JSON()) - pduReceipts = append(pduReceipts, pdu.receipt) + t.PDUs = append(t.PDUs, pdu.JSON()) } // Do the same for pending EDUS in the queue. - for _, edu := range edus { - if edu == nil || edu.edu == nil { + for _, edu := range transaction.edus { + if edu == nil { continue } - t.EDUs = append(t.EDUs, *edu.edu) - eduReceipts = append(eduReceipts, edu.receipt) + t.EDUs = append(t.EDUs, *edu) } logrus.WithField("server_name", oq.destination).Debugf("Sending transaction %q containing %d PDUs, %d EDUs", t.TransactionID, len(t.PDUs), len(t.EDUs)) @@ -411,15 +391,15 @@ func (oq *destinationQueue) nextTransaction( switch err.(type) { case nil: // Clean up the transaction in the database. - if pduReceipts != nil { + for _, receipt := range transaction.pduReceipts { //logrus.Infof("Cleaning PDUs %q", pduReceipt.String()) - if err = oq.db.CleanPDUs(context.Background(), oq.destination, pduReceipts); err != nil { + if err = oq.db.CleanPDUs(context.Background(), oq.destination, receipt); err != nil { log.WithError(err).Errorf("Failed to clean PDUs for server %q", t.Destination) } } - if eduReceipts != nil { + for _, receipt := range transaction.eduReceipts { //logrus.Infof("Cleaning EDUs %q", eduReceipt.String()) - if err = oq.db.CleanEDUs(context.Background(), oq.destination, eduReceipts); err != nil { + if err = oq.db.CleanEDUs(context.Background(), oq.destination, receipt); err != nil { log.WithError(err).Errorf("Failed to clean EDUs for server %q", t.Destination) } } diff --git a/federationsender/queue/queue.go b/federationsender/queue/queue.go index da30e4de1..c59752452 100644 --- a/federationsender/queue/queue.go +++ b/federationsender/queue/queue.go @@ -24,7 +24,6 @@ import ( "github.com/matrix-org/dendrite/federationsender/statistics" "github.com/matrix-org/dendrite/federationsender/storage" - "github.com/matrix-org/dendrite/federationsender/storage/shared" "github.com/matrix-org/dendrite/roomserver/api" "github.com/matrix-org/gomatrixserverlib" log "github.com/sirupsen/logrus" @@ -101,16 +100,6 @@ type SigningInfo struct { PrivateKey ed25519.PrivateKey } -type queuedPDU struct { - receipt *shared.Receipt - pdu *gomatrixserverlib.HeaderedEvent -} - -type queuedEDU struct { - receipt *shared.Receipt - edu *gomatrixserverlib.EDU -} - func (oqs *OutgoingQueues) getQueue(destination gomatrixserverlib.ServerName) *destinationQueue { oqs.queuesMutex.Lock() defer oqs.queuesMutex.Unlock() @@ -127,6 +116,7 @@ func (oqs *OutgoingQueues) getQueue(destination gomatrixserverlib.ServerName) *d interruptBackoff: make(chan bool), signing: oqs.signing, } + oq.pendingTransactions.statistics = oq.statistics oqs.queues[destination] = oq } return oq diff --git a/federationsender/queue/transactions.go b/federationsender/queue/transactions.go new file mode 100644 index 000000000..a9275885f --- /dev/null +++ b/federationsender/queue/transactions.go @@ -0,0 +1,92 @@ +package queue + +import ( + "fmt" + "time" + + "github.com/matrix-org/gomatrixserverlib" + + "github.com/matrix-org/dendrite/federationsender/statistics" + "github.com/matrix-org/dendrite/federationsender/storage/shared" +) + +const ( + maxTransactionsInMemory = 4 + maxPDUsPerTransaction = 50 + maxEDUsPerTransaction = 50 +) + +type queuedTransaction struct { + id gomatrixserverlib.TransactionID + pdus []*gomatrixserverlib.HeaderedEvent + edus []*gomatrixserverlib.EDU + pduReceipts []*shared.Receipt + eduReceipts []*shared.Receipt +} + +type queuedTransactions struct { + statistics *statistics.ServerStatistics + queue []*queuedTransaction +} + +func (q *queuedTransactions) createNew(transactionID gomatrixserverlib.TransactionID) bool { + now := gomatrixserverlib.AsTimestamp(time.Now()) + if len(q.queue) == maxTransactionsInMemory { + return false + } + if transactionID == "" { + transactionID = gomatrixserverlib.TransactionID(fmt.Sprintf("%d-%d", now, q.statistics.SuccessCount())) + } + q.queue = append(q.queue, &queuedTransaction{ + id: transactionID, + }) + return true +} + +func (q *queuedTransactions) getNewestForPDU() *queuedTransaction { + if len(q.queue) == 0 || len(q.queue[len(q.queue)-1].pdus) == maxPDUsPerTransaction { + if len(q.queue) == maxTransactionsInMemory { + return nil + } + if !q.createNew("") { + return nil + } + } + return q.queue[len(q.queue)-1] +} + +func (q *queuedTransactions) getNewestForEDU() *queuedTransaction { + if len(q.queue) == 0 || len(q.queue[len(q.queue)-1].edus) == maxEDUsPerTransaction { + if len(q.queue) == maxTransactionsInMemory { + return nil + } + if !q.createNew("") { + return nil + } + } + return q.queue[len(q.queue)-1] +} + +func (q *queuedTransactions) queuePDUs(receipt *shared.Receipt, pdu ...*gomatrixserverlib.HeaderedEvent) gomatrixserverlib.TransactionID { + last := q.getNewestForPDU() + if last == nil { + return "" + } + last.pdus = append(last.pdus, pdu...) + if receipt != nil { + last.pduReceipts = append(last.pduReceipts, receipt) + } + return last.id +} + +func (q *queuedTransactions) queueEDUs(receipt *shared.Receipt, edu ...*gomatrixserverlib.EDU) gomatrixserverlib.TransactionID { + last := q.getNewestForEDU() + if last == nil { + return "" + } + last.edus = append(last.edus, edu...) + if receipt != nil { + last.eduReceipts = append(last.eduReceipts, receipt) + } + return last.id +} diff --git a/federationsender/storage/interface.go b/federationsender/storage/interface.go index 03d616f1b..a3f5073f9 100644 --- a/federationsender/storage/interface.go +++ b/federationsender/storage/interface.go @@ -36,14 +36,14 @@ type Database interface { StoreJSON(ctx context.Context, js string) (*shared.Receipt, error) - GetPendingPDUs(ctx context.Context, serverName gomatrixserverlib.ServerName, limit int) (pdus map[*shared.Receipt]*gomatrixserverlib.HeaderedEvent, err error) - GetPendingEDUs(ctx context.Context, serverName gomatrixserverlib.ServerName, limit int) (edus map[*shared.Receipt]*gomatrixserverlib.EDU, err error) - AssociatePDUWithDestination(ctx context.Context, transactionID gomatrixserverlib.TransactionID, serverName gomatrixserverlib.ServerName, receipt *shared.Receipt) error AssociateEDUWithDestination(ctx context.Context, serverName gomatrixserverlib.ServerName, receipt *shared.Receipt) error - CleanPDUs(ctx context.Context, serverName gomatrixserverlib.ServerName, receipts []*shared.Receipt) error - CleanEDUs(ctx context.Context, serverName gomatrixserverlib.ServerName, receipts []*shared.Receipt) error + GetNextTransactionPDUs(ctx context.Context, serverName gomatrixserverlib.ServerName, limit int) (gomatrixserverlib.TransactionID, []*gomatrixserverlib.HeaderedEvent, *shared.Receipt, error) + GetNextTransactionEDUs(ctx context.Context, serverName gomatrixserverlib.ServerName, limit int) ([]*gomatrixserverlib.EDU, *shared.Receipt, error) + + CleanPDUs(ctx context.Context, serverName gomatrixserverlib.ServerName, receipt *shared.Receipt) error + CleanEDUs(ctx context.Context, serverName gomatrixserverlib.ServerName, receipt *shared.Receipt) error GetPendingPDUCount(ctx context.Context, serverName gomatrixserverlib.ServerName) (int64, error) GetPendingEDUCount(ctx context.Context, serverName gomatrixserverlib.ServerName) (int64, error) diff --git a/federationsender/storage/postgres/queue_pdus_table.go b/federationsender/storage/postgres/queue_pdus_table.go index f9a477483..95a3b9eee 100644 --- a/federationsender/storage/postgres/queue_pdus_table.go +++ b/federationsender/storage/postgres/queue_pdus_table.go @@ -45,10 +45,16 @@ const insertQueuePDUSQL = "" + const deleteQueuePDUSQL = "" + "DELETE FROM federationsender_queue_pdus WHERE server_name = $1 AND json_nid = ANY($2)" -const selectQueuePDUsSQL = "" + - "SELECT json_nid FROM federationsender_queue_pdus" + +const selectQueuePDUNextTransactionIDSQL = "" + + "SELECT transaction_id FROM federationsender_queue_pdus" + " WHERE server_name = $1" + - " LIMIT $2" + " ORDER BY transaction_id ASC" + + " LIMIT 1" + +const selectQueuePDUsByTransactionSQL = "" + + "SELECT json_nid FROM federationsender_queue_pdus" + + " WHERE server_name = $1 AND transaction_id = $2" + + " LIMIT $3" const selectQueuePDUReferenceJSONCountSQL = "" + "SELECT COUNT(*) FROM federationsender_queue_pdus" + @@ -65,7 +71,8 @@ type queuePDUsStatements struct { db *sql.DB insertQueuePDUStmt *sql.Stmt deleteQueuePDUsStmt *sql.Stmt - selectQueuePDUsStmt *sql.Stmt + selectQueuePDUNextTransactionIDStmt *sql.Stmt + selectQueuePDUsByTransactionStmt *sql.Stmt selectQueuePDUReferenceJSONCountStmt *sql.Stmt selectQueuePDUsCountStmt *sql.Stmt selectQueuePDUServerNamesStmt *sql.Stmt @@ -85,7 +92,10 @@ func NewPostgresQueuePDUsTable(db *sql.DB) (s *queuePDUsStatements, err error) { if s.deleteQueuePDUsStmt, err = s.db.Prepare(deleteQueuePDUSQL); err != nil { return } - if s.selectQueuePDUsStmt, err = s.db.Prepare(selectQueuePDUsSQL); err != nil { + if s.selectQueuePDUNextTransactionIDStmt, err = s.db.Prepare(selectQueuePDUNextTransactionIDSQL); err != nil { + return + } + if s.selectQueuePDUsByTransactionStmt, err = s.db.Prepare(selectQueuePDUsByTransactionSQL); err != nil { return } if s.selectQueuePDUReferenceJSONCountStmt, err = s.db.Prepare(selectQueuePDUReferenceJSONCountSQL); err != nil { @@ -127,6 +137,18 @@ func (s *queuePDUsStatements) DeleteQueuePDUs( return err } +func (s *queuePDUsStatements) SelectQueuePDUNextTransactionID( + ctx context.Context, txn *sql.Tx, serverName gomatrixserverlib.ServerName, +) (gomatrixserverlib.TransactionID, error) { + var transactionID gomatrixserverlib.TransactionID + stmt := sqlutil.TxStmt(txn, s.selectQueuePDUNextTransactionIDStmt) + err := stmt.QueryRowContext(ctx, serverName).Scan(&transactionID) + if err == sql.ErrNoRows { + return "", nil + } + return transactionID, err +} + func (s *queuePDUsStatements) SelectQueuePDUReferenceJSONCount( ctx context.Context, txn *sql.Tx, jsonNID int64, ) (int64, error) { @@ -160,10 +182,11 @@ func (s *queuePDUsStatements) SelectQueuePDUCount( func (s *queuePDUsStatements) SelectQueuePDUs( ctx context.Context, txn *sql.Tx, serverName gomatrixserverlib.ServerName, + transactionID gomatrixserverlib.TransactionID, limit int, ) ([]int64, error) { - stmt := sqlutil.TxStmt(txn, s.selectQueuePDUsStmt) - rows, err := stmt.QueryContext(ctx, serverName, limit) + stmt := sqlutil.TxStmt(txn, s.selectQueuePDUsByTransactionStmt) + rows, err := stmt.QueryContext(ctx, serverName, transactionID, limit) if err != nil { return nil, err } diff --git a/federationsender/storage/shared/storage.go b/federationsender/storage/shared/storage.go index fbf84c705..af9d0d6a3 100644 --- a/federationsender/storage/shared/storage.go +++ b/federationsender/storage/shared/storage.go @@ -17,6 +17,7 @@ package shared import ( "context" "database/sql" + "encoding/json" "fmt" "github.com/matrix-org/dendrite/federationsender/storage/tables" @@ -43,11 +44,16 @@ type Database struct { // to pass them back so that we can clean up if the transaction sends // successfully. type Receipt struct { - nid int64 + nids []int64 } -func (r *Receipt) String() string { - return fmt.Sprintf("%d", r.nid) +func (e *Receipt) Empty() bool { + return len(e.nids) == 0 +} + +func (e *Receipt) String() string { + j, _ := json.Marshal(e.nids) + return string(j) } // UpdateRoom updates the joined hosts for a room and returns what the joined @@ -140,7 +146,7 @@ func (d *Database) StoreJSON( return nil, fmt.Errorf("d.insertQueueJSON: %w", err) } return &Receipt{ - nid: nid, + nids: []int64{nid}, }, nil } diff --git a/federationsender/storage/shared/storage_edus.go b/federationsender/storage/shared/storage_edus.go index 86fee1a37..ae1d15118 100644 --- a/federationsender/storage/shared/storage_edus.go +++ b/federationsender/storage/shared/storage_edus.go @@ -33,14 +33,16 @@ func (d *Database) AssociateEDUWithDestination( receipt *Receipt, ) error { return d.Writer.Do(d.DB, nil, func(txn *sql.Tx) error { - if err := d.FederationSenderQueueEDUs.InsertQueueEDU( - ctx, // context - txn, // SQL transaction - "", // TODO: EDU type for coalescing - serverName, // destination server name - receipt.nid, // NID from the federationsender_queue_json table - ); err != nil { - return fmt.Errorf("InsertQueueEDU: %w", err) + for _, nid := range receipt.nids { + if err := d.FederationSenderQueueEDUs.InsertQueueEDU( + ctx, // context + txn, // SQL transaction + "", // TODO: EDU type for coalescing + serverName, // destination server name + nid, // NID from the federationsender_queue_json table + ); err != nil { + return fmt.Errorf("InsertQueueEDU: %w", err) + } } return nil }) @@ -48,25 +50,29 @@ func (d *Database) AssociateEDUWithDestination( // GetNextTransactionEDUs retrieves events from the database for // the next pending transaction, up to the limit specified. -func (d *Database) GetPendingEDUs( +func (d *Database) GetNextTransactionEDUs( ctx context.Context, serverName gomatrixserverlib.ServerName, limit int, ) ( - edus map[*Receipt]*gomatrixserverlib.EDU, + edus []*gomatrixserverlib.EDU, + receipt *Receipt, err error, ) { - edus = make(map[*Receipt]*gomatrixserverlib.EDU) err = d.Writer.Do(d.DB, nil, func(txn *sql.Tx) error { nids, err := d.FederationSenderQueueEDUs.SelectQueueEDUs(ctx, txn, serverName, limit) if err != nil { return fmt.Errorf("SelectQueueEDUs: %w", err) } + receipt = &Receipt{ + nids: nids, + } + retrieve := make([]int64, 0, len(nids)) for _, nid := range nids { if edu, ok := d.Cache.GetFederationSenderQueuedEDU(nid); ok { - edus[&Receipt{nid}] = edu + edus = append(edus, edu) } else { retrieve = append(retrieve, nid) } @@ -77,12 +83,12 @@ func (d *Database) GetPendingEDUs( return fmt.Errorf("SelectQueueJSON: %w", err) } - for nid, blob := range blobs { + for _, blob := range blobs { var event gomatrixserverlib.EDU if err := json.Unmarshal(blob, &event); err != nil { return fmt.Errorf("json.Unmarshal: %w", err) } - edus[&Receipt{nid}] = &event + edus = append(edus, &event) } return nil @@ -95,24 +101,19 @@ func (d *Database) GetPendingEDUs( func (d *Database) CleanEDUs( ctx context.Context, serverName gomatrixserverlib.ServerName, - receipts []*Receipt, + receipt *Receipt, ) error { - if len(receipts) == 0 { + if receipt == nil { return errors.New("expected receipt") } - nids := make([]int64, len(receipts)) - for i := range receipts { - nids[i] = receipts[i].nid - } - return d.Writer.Do(d.DB, nil, func(txn *sql.Tx) error { - if err := d.FederationSenderQueueEDUs.DeleteQueueEDUs(ctx, txn, serverName, nids); err != nil { + if err := d.FederationSenderQueueEDUs.DeleteQueueEDUs(ctx, txn, serverName, receipt.nids); err != nil { return err } var deleteNIDs []int64 - for _, nid := range nids { + for _, nid := range receipt.nids { count, err := d.FederationSenderQueueEDUs.SelectQueueEDUReferenceJSONCount(ctx, txn, nid) if err != nil { return fmt.Errorf("SelectQueueEDUReferenceJSONCount: %w", err) diff --git a/federationsender/storage/shared/storage_pdus.go b/federationsender/storage/shared/storage_pdus.go index bc298a905..09235a5ec 100644 --- a/federationsender/storage/shared/storage_pdus.go +++ b/federationsender/storage/shared/storage_pdus.go @@ -34,14 +34,16 @@ func (d *Database) AssociatePDUWithDestination( receipt *Receipt, ) error { return d.Writer.Do(d.DB, nil, func(txn *sql.Tx) error { - if err := d.FederationSenderQueuePDUs.InsertQueuePDU( - ctx, // context - txn, // SQL transaction - transactionID, // transaction ID - serverName, // destination server name - receipt.nid, // NID from the federationsender_queue_json table - ); err != nil { - return fmt.Errorf("InsertQueuePDU: %w", err) + for _, nid := range receipt.nids { + if err := d.FederationSenderQueuePDUs.InsertQueuePDU( + ctx, // context + txn, // SQL transaction + transactionID, // transaction ID + serverName, // destination server name + nid, // NID from the federationsender_queue_json table + ); err != nil { + return fmt.Errorf("InsertQueuePDU: %w", err) + } } return nil }) @@ -49,12 +51,14 @@ func (d *Database) AssociatePDUWithDestination( // GetNextTransactionPDUs retrieves events from the database for // the next pending transaction, up to the limit specified. -func (d *Database) GetPendingPDUs( +func (d *Database) GetNextTransactionPDUs( ctx context.Context, serverName gomatrixserverlib.ServerName, limit int, ) ( - events map[*Receipt]*gomatrixserverlib.HeaderedEvent, + transactionID gomatrixserverlib.TransactionID, + events []*gomatrixserverlib.HeaderedEvent, + receipt *Receipt, err error, ) { // Strictly speaking this doesn't need to be using the writer @@ -62,17 +66,29 @@ func (d *Database) GetPendingPDUs( // a guarantee of transactional isolation, it's actually useful // to know in SQLite mode that nothing else is trying to modify // the database. - events = make(map[*Receipt]*gomatrixserverlib.HeaderedEvent) err = d.Writer.Do(d.DB, nil, func(txn *sql.Tx) error { - nids, err := d.FederationSenderQueuePDUs.SelectQueuePDUs(ctx, txn, serverName, limit) + transactionID, err = d.FederationSenderQueuePDUs.SelectQueuePDUNextTransactionID(ctx, txn, serverName) + if err != nil { + return fmt.Errorf("SelectQueuePDUNextTransactionID: %w", err) + } + + if transactionID == "" { + return nil + } + + nids, err := d.FederationSenderQueuePDUs.SelectQueuePDUs(ctx, txn, serverName, transactionID, limit) if err != nil { return fmt.Errorf("SelectQueuePDUs: %w", err) } + receipt = &Receipt{ + nids: nids, + } + retrieve := make([]int64, 0, len(nids)) for _, nid := range nids { if event, ok := d.Cache.GetFederationSenderQueuedPDU(nid); ok { - events[&Receipt{nid}] = event + events = append(events, event) } else { retrieve = append(retrieve, nid) } @@ -88,7 +104,7 @@ func (d *Database) GetPendingPDUs( if err := json.Unmarshal(blob, &event); err != nil { return fmt.Errorf("json.Unmarshal: %w", err) } - events[&Receipt{nid}] = &event + events = append(events, &event) d.Cache.StoreFederationSenderQueuedPDU(nid, &event) } @@ -103,24 +119,19 @@ func (d *Database) GetPendingPDUs( func (d *Database) CleanPDUs( ctx context.Context, serverName gomatrixserverlib.ServerName, - receipts []*Receipt, + receipt *Receipt, ) error { - if len(receipts) == 0 { + if receipt == nil { return errors.New("expected receipt") } - nids := make([]int64, len(receipts)) - for i := range receipts { - nids[i] = receipts[i].nid - } - return d.Writer.Do(d.DB, nil, func(txn *sql.Tx) error { - if err := d.FederationSenderQueuePDUs.DeleteQueuePDUs(ctx, txn, serverName, nids); err != nil { + if err := d.FederationSenderQueuePDUs.DeleteQueuePDUs(ctx, txn, serverName, receipt.nids); err != nil { return err } var deleteNIDs []int64 - for _, nid := range nids { + for _, nid := range receipt.nids { count, err := d.FederationSenderQueuePDUs.SelectQueuePDUReferenceJSONCount(ctx, txn, nid) if err != nil { return fmt.Errorf("SelectQueuePDUReferenceJSONCount: %w", err) diff --git a/federationsender/storage/sqlite3/queue_pdus_table.go b/federationsender/storage/sqlite3/queue_pdus_table.go index e0fdbda5f..70519c9ef 100644 --- a/federationsender/storage/sqlite3/queue_pdus_table.go +++ b/federationsender/storage/sqlite3/queue_pdus_table.go @@ -53,10 +53,10 @@ const selectQueueNextTransactionIDSQL = "" + " ORDER BY transaction_id ASC" + " LIMIT 1" -const selectQueuePDUsSQL = "" + +const selectQueuePDUsByTransactionSQL = "" + "SELECT json_nid FROM federationsender_queue_pdus" + - " WHERE server_name = $1" + - " LIMIT $2" + " WHERE server_name = $1 AND transaction_id = $2" + + " LIMIT $3" const selectQueuePDUsReferenceJSONCountSQL = "" + "SELECT COUNT(*) FROM federationsender_queue_pdus" + @@ -73,7 +73,7 @@ type queuePDUsStatements struct { db *sql.DB insertQueuePDUStmt *sql.Stmt selectQueueNextTransactionIDStmt *sql.Stmt - selectQueuePDUsStmt *sql.Stmt + selectQueuePDUsByTransactionStmt *sql.Stmt selectQueueReferenceJSONCountStmt *sql.Stmt selectQueuePDUsCountStmt *sql.Stmt selectQueueServerNamesStmt *sql.Stmt @@ -97,7 +97,7 @@ func NewSQLiteQueuePDUsTable(db *sql.DB) (s *queuePDUsStatements, err error) { if s.selectQueueNextTransactionIDStmt, err = db.Prepare(selectQueueNextTransactionIDSQL); err != nil { return } - if s.selectQueuePDUsStmt, err = db.Prepare(selectQueuePDUsSQL); err != nil { + if s.selectQueuePDUsByTransactionStmt, err = db.Prepare(selectQueuePDUsByTransactionSQL); err != nil { return } if s.selectQueueReferenceJSONCountStmt, err = db.Prepare(selectQueuePDUsReferenceJSONCountSQL); err != nil { @@ -193,10 +193,11 @@ func (s *queuePDUsStatements) SelectQueuePDUCount( func (s *queuePDUsStatements) SelectQueuePDUs( ctx context.Context, txn *sql.Tx, serverName gomatrixserverlib.ServerName, + transactionID gomatrixserverlib.TransactionID, limit int, ) ([]int64, error) { - stmt := sqlutil.TxStmt(txn, s.selectQueuePDUsStmt) - rows, err := stmt.QueryContext(ctx, serverName, limit) + stmt := sqlutil.TxStmt(txn, s.selectQueuePDUsByTransactionStmt) + rows, err := stmt.QueryContext(ctx, serverName, transactionID, limit) if err != nil { return nil, err } diff --git a/federationsender/storage/tables/interface.go b/federationsender/storage/tables/interface.go index 69e952de2..1167a212a 100644 --- a/federationsender/storage/tables/interface.go +++ b/federationsender/storage/tables/interface.go @@ -25,9 +25,10 @@ import ( type FederationSenderQueuePDUs interface { InsertQueuePDU(ctx context.Context, txn *sql.Tx, transactionID gomatrixserverlib.TransactionID, serverName gomatrixserverlib.ServerName, nid int64) error DeleteQueuePDUs(ctx context.Context, txn *sql.Tx, serverName gomatrixserverlib.ServerName, jsonNIDs []int64) error + SelectQueuePDUNextTransactionID(ctx context.Context, txn *sql.Tx, serverName gomatrixserverlib.ServerName) (gomatrixserverlib.TransactionID, error) SelectQueuePDUReferenceJSONCount(ctx context.Context, txn *sql.Tx, jsonNID int64) (int64, error) SelectQueuePDUCount(ctx context.Context, txn *sql.Tx, serverName gomatrixserverlib.ServerName) (int64, error) - SelectQueuePDUs(ctx context.Context, txn *sql.Tx, serverName gomatrixserverlib.ServerName, limit int) ([]int64, error) + SelectQueuePDUs(ctx context.Context, txn *sql.Tx, serverName gomatrixserverlib.ServerName, transactionID gomatrixserverlib.TransactionID, limit int) ([]int64, error) SelectQueuePDUServerNames(ctx context.Context, txn *sql.Tx) ([]gomatrixserverlib.ServerName, error) }