Add state key storage

src/github.com/matrix-org/dendrite/roomserver/storage/sql.go

@@ -10,6 +10,8 @@ type statements struct {
 	upsertPartitionOffsetStmt  *sql.Stmt
 	insertEventTypeNIDStmt     *sql.Stmt
 	selectEventTypeNIDStmt     *sql.Stmt
+	insertEventStateKeyNIDStmt *sql.Stmt
+	selectEventStateKeyNIDStmt *sql.Stmt
 }
 
 func (s *statements) prepare(db *sql.DB) error {
@@ -23,6 +25,10 @@ func (s *statements) prepare(db *sql.DB) error {
 		return err
 	}
 
+	if err = s.prepareEventStateKeys(db); err != nil {
+		return err
+	}
+
 	return nil
 }
 
@@ -160,3 +166,63 @@ func (s *statements) selectEventTypeNID(eventType string) (eventTypeNID int64, e
 	}
 	return
 }
+
+func (s *statements) prepareEventStateKeys(db *sql.DB) (err error) {
+	_, err = db.Exec(eventStateKeysSchema)
+	if err != nil {
+		return
+	}
+	if s.insertEventStateKeyNIDStmt, err = db.Prepare(insertEventStateKeyNIDSQL); err != nil {
+		return
+	}
+	if s.selectEventStateKeyNIDStmt, err = db.Prepare(selectEventStateKeyNIDSQL); err != nil {
+		return
+	}
+	return
+}
+
+const eventStateKeysSchema = `
+-- Numeric versions of the event "state_key"s. State keys tend to be reused so
+-- assigning each string a numeric ID should reduce the amount of data that
+-- needs to be stored and fetched from the database.
+-- It also means that many operations can work with int64 arrays rather than
+-- string arrays which may help reduce GC pressure.
+-- Well known state keys are pre-assigned numeric IDs:
+--   1 -> "" (the empty string)
+CREATE SEQUENCE IF NOT EXISTS event_state_key_nid_seq START 65536;
+CREATE TABLE event_state_keys (
+    -- Local numeric ID for the state key.
+    event_state_key_nid BIGINT PRIMARY KEY DEFAULT nextval('event_state_key_nid_seq'),
+    event_state_key TEXT NOT NULL CONSTRAINT event_state_key_unique UNIQUE
+);
+INSERT INTO event_state_keys (event_state_key_nid, event_state_key) VALUES (
+    (1, '')
+);
+`
+
+const insertEventStateKeyNIDSQL = "" +
+	"INSERT INTO event_state_keys (event_state_key) VALUES ($1)" +
+	" ON CONFLICT ON CONSTRAINT event_state_key_unique" +
+	" DO UPDATE SET event_state_key = $1" +
+	" RETURNING (event_state_key_nid)"
+
+const selectEventStateKeyNIDSQL = "" +
+	"SELECT event_state_key_nid FROM event_state_keys WHERE event_state_key = $1"
+
+func (s *statements) insertEventStateKeyNID(eventStateKey string) (eventStateKeyNID int64, err error) {
+	err = s.insertEventStateKeyNIDStmt.QueryRow(eventStateKey).Scan(&eventStateKeyNID)
+	if err == sql.ErrNoRows {
+		eventStateKeyNID = 0
+		err = nil
+	}
+	return
+}
+
+func (s *statements) selectEventStateKeyNID(eventStateKey string) (eventStateKeyNID int64, err error) {
+	err = s.selectEventStateKeyNIDStmt.QueryRow(eventStateKey).Scan(&eventStateKeyNID)
+	if err == sql.ErrNoRows {
+		eventStateKeyNID = 0
+		err = nil
+	}
+	return
+}