refactor RoomInfoTree building into separate package for better org.

roominfotree/roominfotree.go

@@ -0,0 +1,169 @@
+package roominfotree
+
+import (
+	"context"
+	"log"
+	"sync"
+
+	"github.com/matrix-org/gomatrixserverlib/fclient"
+	"github.com/matrix-org/gomatrixserverlib/spec"
+	"maunium.net/go/mautrix"
+	"maunium.net/go/mautrix/id"
+)
+
+type homeserverServerVersionInfo struct {
+	homeserver        string
+	serverVersionInfo fclient.Version
+}
+
+type RoomInfoTree map[id.RoomID](*RoomInfo)
+
+type RoomInfo struct {
+	MemberCount     uint
+	MaxRoomVersions map[string](*MaxRoomVersionInfo)
+}
+
+type MaxRoomVersionInfo struct {
+	MemberCount uint
+	Servers     map[string](*ServerInfo)
+}
+
+type ServerInfo struct {
+	MemberCount uint
+	Versions    map[string](*VersionInfo)
+}
+
+type VersionInfo struct {
+	MemberCount uint
+	Homeservers map[string](*HomeserverInfo)
+}
+
+type HomeserverInfo struct {
+	MemberCount uint
+}
+
+var unknownServerVersionInfo = fclient.Version{
+	Server: struct {
+		Name    string `json:"name"`
+		Version string `json:"version"`
+	}{
+		Name:    "unknown",
+		Version: "unknown",
+	},
+}
+
+func getServerVersionInfoByHomeserver(homeservers []string, federationClient *fclient.Client) map[string](fclient.Version) {
+	homeserverChannel := make(chan string)
+	go func() {
+		for _, hs := range homeservers {
+			homeserverChannel <- hs
+		}
+		close(homeserverChannel)
+	}()
+
+	homeserverSVInfoChannel := make(chan homeserverServerVersionInfo)
+	var wg sync.WaitGroup
+	const numHomeserverSVInfoReceivers = 100
+	wg.Add(numHomeserverSVInfoReceivers)
+	for i := 0; i < numHomeserverSVInfoReceivers; i++ {
+		go func() {
+			for hs := range homeserverChannel {
+				serverVersionInfo, err := federationClient.GetVersion(context.Background(), spec.ServerName(hs))
+				if err != nil {
+					serverVersionInfo = unknownServerVersionInfo
+				}
+				homeserverSVInfoChannel <- homeserverServerVersionInfo{hs, serverVersionInfo}
+			}
+			wg.Done()
+		}()
+	}
+
+	go func() {
+		wg.Wait()
+		close(homeserverSVInfoChannel)
+	}()
+
+	serverVersionInfoByHomeserver := make(map[string](fclient.Version))
+	for homeserverSVInfo := range homeserverSVInfoChannel {
+		serverVersionInfoByHomeserver[homeserverSVInfo.homeserver] = homeserverSVInfo.serverVersionInfo
+	}
+	return serverVersionInfoByHomeserver
+}
+
+func Get(roomIDs []id.RoomID, client *mautrix.Client, federationClient *fclient.Client) RoomInfoTree {
+	joinedMembersByRoomID := make(map[id.RoomID]*mautrix.RespJoinedMembers)
+	for _, roomID := range roomIDs {
+		joinedMembers, err := client.JoinedMembers(context.Background(), roomID)
+		if err != nil {
+			log.Fatal(err)
+		}
+		joinedMembersByRoomID[roomID] = joinedMembers
+	}
+
+	membersByHomeserverByRoomID := make(map[id.RoomID](map[string]uint))
+	for roomID, joinedMembers := range joinedMembersByRoomID {
+		membersByHomeserver := make(map[string]uint)
+		for key := range joinedMembers.Joined {
+			membersByHomeserver[key.Homeserver()]++
+		}
+		membersByHomeserverByRoomID[roomID] = membersByHomeserver
+	}
+
+	homeserverSet := make(map[string]bool)
+	for _, membersByHomeserver := range membersByHomeserverByRoomID {
+		for hs := range membersByHomeserver {
+			homeserverSet[hs] = true
+		}
+	}
+	homeservers := make([]string, 0, len(homeserverSet))
+	for hs := range homeserverSet {
+		homeservers = append(homeservers, hs)
+	}
+
+	serverVersionInfoByHomeserver := getServerVersionInfoByHomeserver(homeservers, federationClient)
+
+	roomInfoTree := make(RoomInfoTree)
+
+	for roomID, membersByHomeserver := range membersByHomeserverByRoomID {
+		for hs, members := range membersByHomeserver {
+			serverVersionInfo := serverVersionInfoByHomeserver[hs]
+
+			maxRoomVersion := getMaxRoomVersion(serverVersionInfo)
+
+			// Sort into roomInfoTree and add to counters.
+			roomInfo, ok := roomInfoTree[roomID]
+			if !ok {
+				roomInfo = &RoomInfo{MaxRoomVersions: make(map[string]*MaxRoomVersionInfo)}
+				roomInfoTree[roomID] = roomInfo
+			}
+			maxRoomVersionInfo, ok := roomInfo.MaxRoomVersions[maxRoomVersion]
+			if !ok {
+				maxRoomVersionInfo = &MaxRoomVersionInfo{Servers: make(map[string]*ServerInfo)}
+				roomInfo.MaxRoomVersions[maxRoomVersion] = maxRoomVersionInfo
+			}
+			serverInfo, ok := maxRoomVersionInfo.Servers[serverVersionInfo.Server.Name]
+			if !ok {
+				serverInfo = &ServerInfo{Versions: make(map[string]*VersionInfo)}
+				maxRoomVersionInfo.Servers[serverVersionInfo.Server.Name] = serverInfo
+			}
+			versionInfo, ok := serverInfo.Versions[serverVersionInfo.Server.Version]
+			if !ok {
+				versionInfo = &VersionInfo{Homeservers: make(map[string]*HomeserverInfo)}
+				serverInfo.Versions[serverVersionInfo.Server.Version] = versionInfo
+			}
+			homeserverInfo, ok := versionInfo.Homeservers[hs]
+			if !ok {
+				homeserverInfo = &HomeserverInfo{}
+				versionInfo.Homeservers[hs] = homeserverInfo
+			}
+
+			homeserverInfo.MemberCount = members
+			versionInfo.MemberCount += members
+			serverInfo.MemberCount += members
+			maxRoomVersionInfo.MemberCount += members
+			roomInfo.MemberCount += members
+		}
+	}
+
+	return roomInfoTree
+}