Add context to the partition offset table
MODIFY
src/github.com/matrix-org/dendrite/clientapi/auth/storage/accounts/storage.go
src/github.com/matrix-org/dendrite/clientapi/auth/storage/accounts/storage.go
@@ -30,7 +30,7 @@ import (
// Database represents an account database
type Database struct {
db *sql.DB
- partitions common.PartitionOffsetStatements
+ common.PartitionOffsetStatements
accounts accountsStatements
profiles profilesStatements
memberships membershipStatements
@@ -127,16 +127,6 @@ func (d *Database) CreateAccount(
return d.accounts.insertAccount(ctx, localpart, hash)
}
-// PartitionOffsets implements common.PartitionStorer
-func (d *Database) PartitionOffsets(topic string) ([]common.PartitionOffset, error) {
- return d.partitions.SelectPartitionOffsets(topic)
-}
-
-// SetPartitionOffset implements common.PartitionStorer
-func (d *Database) SetPartitionOffset(topic string, partition int32, offset int64) error {
- return d.partitions.UpsertPartitionOffset(topic, partition, offset)
-}
-
// SaveMembership saves the user matching a given localpart as a member of a given
// room. It also stores the ID of the membership event and a flag on whether the user
// is still in the room.

MODIFY
src/github.com/matrix-org/dendrite/common/consumers.go
src/github.com/matrix-org/dendrite/common/consumers.go
package common
import (
+ "context"
"fmt"
sarama "gopkg.in/Shopify/sarama.v1"
@@ -31,9 +32,9 @@ type PartitionOffset struct {
// A PartitionStorer has the storage APIs needed by the consumer.
type PartitionStorer interface {
// PartitionOffsets returns the offsets the consumer has reached for each partition.
- PartitionOffsets(topic string) ([]PartitionOffset, error)
+ PartitionOffsets(ctx context.Context, topic string) ([]PartitionOffset, error)
// SetPartitionOffset records where the consumer has reached for a partition.
- SetPartitionOffset(topic string, partition int32, offset int64) error
+ SetPartitionOffset(ctx context.Context, topic string, partition int32, offset int64) error
}
// A ContinualConsumer continually consumes logs even across restarts. It requires a PartitionStorer to
@@ -75,7 +76,7 @@ func (c *ContinualConsumer) Start() error {
offsets[partition] = sarama.OffsetOldest
}
- storedOffsets, err := c.PartitionStore.PartitionOffsets(c.Topic)
+ storedOffsets, err := c.PartitionStore.PartitionOffsets(context.TODO(), c.Topic)
if err != nil {
return err
}
@@ -110,7 +111,7 @@ func (c *ContinualConsumer) consumePartition(pc sarama.PartitionConsumer) {
for message := range pc.Messages() {
msgErr := c.ProcessMessage(message)
// Advance our position in the stream so that we will start at the right position after a restart.
- if err := c.PartitionStore.SetPartitionOffset(c.Topic, message.Partition, message.Offset); err != nil {
+ if err := c.PartitionStore.SetPartitionOffset(context.TODO(), c.Topic, message.Partition, message.Offset); err != nil {
panic(fmt.Errorf("the ContinualConsumer failed to SetPartitionOffset: %s", err))
}
// Shutdown if we were told to do so.

MODIFY
src/github.com/matrix-org/dendrite/common/partition_offset_table.go
src/github.com/matrix-org/dendrite/common/partition_offset_table.go
package common
-import "database/sql"
-import "strings"
+import (
+ "context"
+ "database/sql"
+ "strings"
+)
const partitionOffsetsSchema = `
-- The offsets that the server has processed up to.
@@ -66,9 +69,25 @@ func (s *PartitionOffsetStatements) Prepare(db *sql.DB, prefix string) (err erro
return
}
-// SelectPartitionOffsets returns all the partition offsets for the given topic.
-func (s *PartitionOffsetStatements) SelectPartitionOffsets(topic string) ([]PartitionOffset, error) {
- rows, err := s.selectPartitionOffsetsStmt.Query(topic)
+// PartitionOffsets implements PartitionStorer
+func (s *PartitionOffsetStatements) PartitionOffsets(
+ ctx context.Context, topic string,
+) ([]PartitionOffset, error) {
+ return s.selectPartitionOffsets(ctx, topic)
+}
+
+// SetPartitionOffset implements PartitionStorer
+func (s *PartitionOffsetStatements) SetPartitionOffset(
+ ctx context.Context, topic string, partition int32, offset int64,
+) error {
+ return s.upsertPartitionOffset(ctx, topic, partition, offset)
+}
+
+// selectPartitionOffsets returns all the partition offsets for the given topic.
+func (s *PartitionOffsetStatements) selectPartitionOffsets(
+ ctx context.Context, topic string,
+) ([]PartitionOffset, error) {
+ rows, err := s.selectPartitionOffsetsStmt.QueryContext(ctx, topic)
if err != nil {
return nil, err
}
@@ -85,7 +104,9 @@ func (s *PartitionOffsetStatements) SelectPartitionOffsets(topic string) ([]Part
}
// UpsertPartitionOffset updates or inserts the partition offset for the given topic.
-func (s *PartitionOffsetStatements) UpsertPartitionOffset(topic string, partition int32, offset int64) error {
- _, err := s.upsertPartitionOffsetStmt.Exec(topic, partition, offset)
+func (s *PartitionOffsetStatements) upsertPartitionOffset(
+ ctx context.Context, topic string, partition int32, offset int64,
+) error {
+ _, err := s.upsertPartitionOffsetStmt.ExecContext(ctx, topic, partition, offset)
return err
}

MODIFY
src/github.com/matrix-org/dendrite/federationsender/storage/storage.go
src/github.com/matrix-org/dendrite/federationsender/storage/storage.go
@@ -61,16 +61,6 @@ func (d *Database) prepare() error {
return nil
}
-// PartitionOffsets implements common.PartitionStorer
-func (d *Database) PartitionOffsets(topic string) ([]common.PartitionOffset, error) {
- return d.SelectPartitionOffsets(topic)
-}
-
-// SetPartitionOffset implements common.PartitionStorer
-func (d *Database) SetPartitionOffset(topic string, partition int32, offset int64) error {
- return d.UpsertPartitionOffset(topic, partition, offset)
-}
-
// UpdateRoom updates the joined hosts for a room and returns what the joined
// hosts were before the update.
func (d *Database) UpdateRoom(

MODIFY
src/github.com/matrix-org/dendrite/publicroomsapi/storage/storage.go
src/github.com/matrix-org/dendrite/publicroomsapi/storage/storage.go
@@ -28,7 +28,7 @@ import (
// PublicRoomsServerDatabase represents a public rooms server database.
type PublicRoomsServerDatabase struct {
db *sql.DB
- partitions common.PartitionOffsetStatements
+ common.PartitionOffsetStatements
statements publicRoomsStatements
}
@@ -52,16 +52,6 @@ func NewPublicRoomsServerDatabase(dataSourceName string) (*PublicRoomsServerData
return &PublicRoomsServerDatabase{db, partitions, statements}, nil
}
-// PartitionOffsets implements common.PartitionStorer
-func (d *PublicRoomsServerDatabase) PartitionOffsets(topic string) ([]common.PartitionOffset, error) {
- return d.partitions.SelectPartitionOffsets(topic)
-}
-
-// SetPartitionOffset implements common.PartitionStorer
-func (d *PublicRoomsServerDatabase) SetPartitionOffset(topic string, partition int32, offset int64) error {
- return d.partitions.UpsertPartitionOffset(topic, partition, offset)
-}
-
// GetRoomVisibility returns the room visibility as a boolean: true if the room
// is publicly visible, false if not.
// Returns an error if the retrieval failed.

MODIFY
src/github.com/matrix-org/dendrite/syncapi/storage/syncserver.go
src/github.com/matrix-org/dendrite/syncapi/storage/syncserver.go
@@ -43,7 +43,7 @@ type streamEvent struct {
// SyncServerDatabase represents a sync server database
type SyncServerDatabase struct {
db *sql.DB
- partitions common.PartitionOffsetStatements
+ common.PartitionOffsetStatements
accountData accountDataStatements
events outputRoomEventsStatements
roomstate currentRoomStateStatements
@@ -57,7 +57,7 @@ func NewSyncServerDatabase(dataSourceName string) (*SyncServerDatabase, error) {
if d.db, err = sql.Open("postgres", dataSourceName); err != nil {
return nil, err
}
- if err = d.partitions.Prepare(d.db, "syncapi"); err != nil {
+ if err = d.PartitionOffsetStatements.Prepare(d.db, "syncapi"); err != nil {
return nil, err
}
if err = d.accountData.prepare(d.db); err != nil {
@@ -162,16 +162,6 @@ func (d *SyncServerDatabase) GetStateEvent(
return d.roomstate.selectStateEvent(ctx, evType, roomID, stateKey)
}
-// PartitionOffsets implements common.PartitionStorer
-func (d *SyncServerDatabase) PartitionOffsets(topic string) ([]common.PartitionOffset, error) {
- return d.partitions.SelectPartitionOffsets(topic)
-}
-
-// SetPartitionOffset implements common.PartitionStorer
-func (d *SyncServerDatabase) SetPartitionOffset(topic string, partition int32, offset int64) error {
- return d.partitions.UpsertPartitionOffset(topic, partition, offset)
-}
-
// SyncStreamPosition returns the latest position in the sync stream. Returns 0 if there are no events yet.
func (d *SyncServerDatabase) SyncStreamPosition(ctx context.Context) (types.StreamPosition, error) {
return d.syncStreamPositionTx(ctx, nil)

Add context to the partition offset table (#249)