mrvc/roominfotree/roominfotree.go

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 getMembersByHomeserverByRoomID(roomIDs []id.RoomID, client *mautrix.Client) map[id.RoomID](map[string]uint) {
	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 userID := range joinedMembers.Joined {
			membersByHomeserver[userID.Homeserver()]++
		}
		membersByHomeserverByRoomID[roomID] = membersByHomeserver
	}

	return membersByHomeserverByRoomID
}

// Gets the homeservers appearing in the given membersByHomeserverByRoomID map tree.
func getHomeservers(membersByHomeserverByRoomID map[id.RoomID]map[string]uint) []string {
	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)
	}

	return homeservers
}

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 {
	membersByHomeserverByRoomID := getMembersByHomeserverByRoomID(roomIDs, client)
	homeservers := getHomeservers(membersByHomeserverByRoomID)
	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
}