From 112f4ef920ebbb08641ad440d601220ec8107908 Mon Sep 17 00:00:00 2001
From: Brendan Abolivier <contact@brendanabolivier.com>
Date: Fri, 11 Aug 2017 17:59:14 +0100
Subject: [PATCH] Count number of joined user in rooms

---
 .../roomserver/publicroom/public_room.go      |  9 ++++-
 .../roomserver/storage/membership_table.go    | 36 ++++++++++++++++++-
 .../dendrite/roomserver/storage/storage.go    |  5 +++
 3 files changed, 48 insertions(+), 2 deletions(-)

diff --git a/src/github.com/matrix-org/dendrite/roomserver/publicroom/public_room.go b/src/github.com/matrix-org/dendrite/roomserver/publicroom/public_room.go
index 6fa68cd8a..a4da4fcde 100644
--- a/src/github.com/matrix-org/dendrite/roomserver/publicroom/public_room.go
+++ b/src/github.com/matrix-org/dendrite/roomserver/publicroom/public_room.go
@@ -54,6 +54,9 @@ type RoomserverPublicRoomAPIDatabase interface {
 	// ordered by room NID (ie map[roomNID] = []alias)
 	// Returns an error if the retrieval failed
 	GetAliasesFromRoomNIDs(roomNIDs []types.RoomNID) (map[types.RoomNID][]string, error)
+	// Returns the number of joined user in roms identified by given room numeric IDs.
+	// Returns an error if the retrieval failed
+	CountJoinedMembersInRooms(roomNIDs []types.RoomNID) (map[types.RoomNID]int64, error)
 }
 
 // RoomserverPublicRoomAPI is an implementation of api.RoomserverPublicRoomAPI
@@ -137,6 +140,10 @@ func (r *RoomserverPublicRoomAPI) GetPublicRooms(
 	if err != nil {
 		return err
 	}
+	nbMembers, err := r.DB.CountJoinedMembersInRooms(roomNIDs)
+	if err != nil {
+		return err
+	}
 
 	chunks := []api.PublicRoomsChunk{}
 	// Iterate over the array of aliases instead of the array of rooms, because
@@ -145,7 +152,7 @@ func (r *RoomserverPublicRoomAPI) GetPublicRooms(
 		chunk := api.PublicRoomsChunk{
 			RoomID:           roomIDs[roomNID],
 			Aliases:          as,
-			NumJoinedMembers: 0,
+			NumJoinedMembers: nbMembers[roomNID],
 			WorldReadable:    true,
 			GuestCanJoin:     true,
 		}
diff --git a/src/github.com/matrix-org/dendrite/roomserver/storage/membership_table.go b/src/github.com/matrix-org/dendrite/roomserver/storage/membership_table.go
index 725e5b8d9..f209c0d2c 100644
--- a/src/github.com/matrix-org/dendrite/roomserver/storage/membership_table.go
+++ b/src/github.com/matrix-org/dendrite/roomserver/storage/membership_table.go
@@ -17,6 +17,7 @@ package storage
 import (
 	"database/sql"
 
+	"github.com/lib/pq"
 	"github.com/matrix-org/dendrite/roomserver/types"
 )
 
@@ -33,7 +34,7 @@ const membershipSchema = `
 -- and the room state tables.
 -- This table is updated in one of 3 ways:
 --   1) The membership of a user changes within the current state of the room.
---   2) An invite is received outside of a room over federation. 
+--   2) An invite is received outside of a room over federation.
 --   3) An invite is rejected outside of a room over federation.
 CREATE TABLE IF NOT EXISTS roomserver_membership (
 	room_nid BIGINT NOT NULL,
@@ -61,6 +62,11 @@ const selectMembershipForUpdateSQL = "" +
 	"SELECT membership_nid FROM roomserver_membership" +
 	" WHERE room_nid = $1 AND target_nid = $2 FOR UPDATE"
 
+const countJoinedMembersInRoomsSQL = "" +
+	"SELECT room_nid, COUNT(*) FROM roomserver_membership" +
+	" WHERE room_nid = ANY($1) AND membership_nid = $2" +
+	" GROUP BY room_nid"
+
 const updateMembershipSQL = "" +
 	"UPDATE roomserver_membership SET sender_nid = $3, membership_nid = $4" +
 	" WHERE room_nid = $1 AND target_nid = $2"
@@ -68,6 +74,7 @@ const updateMembershipSQL = "" +
 type membershipStatements struct {
 	insertMembershipStmt          *sql.Stmt
 	selectMembershipForUpdateStmt *sql.Stmt
+	countJoinedMembersInRoomsStmt *sql.Stmt
 	updateMembershipStmt          *sql.Stmt
 }
 
@@ -80,6 +87,7 @@ func (s *membershipStatements) prepare(db *sql.DB) (err error) {
 	return statementList{
 		{&s.insertMembershipStmt, insertMembershipSQL},
 		{&s.selectMembershipForUpdateStmt, selectMembershipForUpdateSQL},
+		{&s.countJoinedMembersInRoomsStmt, countJoinedMembersInRoomsSQL},
 		{&s.updateMembershipStmt, updateMembershipSQL},
 	}.prepare(db)
 }
@@ -100,6 +108,32 @@ func (s *membershipStatements) selectMembershipForUpdate(
 	return
 }
 
+func (s *membershipStatements) countJoinedMembersInRooms(roomNIDs []types.RoomNID) (map[types.RoomNID]int64, error) {
+	nIDs := []int64{}
+	for _, roomNID := range roomNIDs {
+		nIDs = append(nIDs, int64(roomNID))
+	}
+
+	rows, err := s.countJoinedMembersInRoomsStmt.Query(pq.Int64Array(nIDs), membershipStateJoin)
+	if err != nil {
+		return nil, err
+	}
+
+	members := make(map[types.RoomNID]int64)
+	for rows.Next() {
+		var roomNID types.RoomNID
+		var count int64
+
+		if err = rows.Scan(&roomNID, &count); err != nil {
+			return members, err
+		}
+
+		members[roomNID] = count
+	}
+
+	return members, err
+}
+
 func (s *membershipStatements) updateMembership(
 	txn *sql.Tx, roomNID types.RoomNID, targetUserNID types.EventStateKeyNID,
 	senderUserNID types.EventStateKeyNID, membership membershipState,
diff --git a/src/github.com/matrix-org/dendrite/roomserver/storage/storage.go b/src/github.com/matrix-org/dendrite/roomserver/storage/storage.go
index 7704183b5..388f61cc2 100644
--- a/src/github.com/matrix-org/dendrite/roomserver/storage/storage.go
+++ b/src/github.com/matrix-org/dendrite/roomserver/storage/storage.go
@@ -410,6 +410,11 @@ func (d *Database) UpdateRoomVisibility(roomNID types.RoomNID, visibility bool)
 	return d.statements.updateVisibilityForRoomNID(roomNID, visibility)
 }
 
+// CountJoinedMembersInRooms implements publicroom.RoomserverPublicRoomAPIDB
+func (d *Database) CountJoinedMembersInRooms(roomNIDs []types.RoomNID) (map[types.RoomNID]int64, error) {
+	return d.statements.countJoinedMembersInRooms(roomNIDs)
+}
+
 type membershipUpdater struct {
 	transaction
 	d             *Database