Refactor room hierarchy walker

clientapi/routing/room_hierarchy.go

@@ -30,24 +30,28 @@ import (
 )
 
 type RoomHierarchyPaginationCache struct {
-	cache map[string]roomserverAPI.CachedRoomHierarchyWalker
+	cache map[string]roomserverAPI.RoomHierarchyWalker
 	mu    sync.Mutex
 }
 
 func NewRoomHierarchyPaginationCache() RoomHierarchyPaginationCache {
 	return RoomHierarchyPaginationCache{
-		cache: map[string]roomserverAPI.CachedRoomHierarchyWalker{},
+		cache: map[string]roomserverAPI.RoomHierarchyWalker{},
 	}
 }
 
-func (c *RoomHierarchyPaginationCache) Get(token string) roomserverAPI.CachedRoomHierarchyWalker {
+func (c *RoomHierarchyPaginationCache) Get(token string) *roomserverAPI.RoomHierarchyWalker {
 	c.mu.Lock()
 	defer c.mu.Unlock()
-	line := c.cache[token]
-	return line
+	line, ok := c.cache[token]
+	if ok {
+		return &line
+	} else {
+		return nil
+	}
 }
 
-func (c *RoomHierarchyPaginationCache) AddLine(line roomserverAPI.CachedRoomHierarchyWalker) string {
+func (c *RoomHierarchyPaginationCache) AddLine(line roomserverAPI.RoomHierarchyWalker) string {
 	c.mu.Lock()
 	defer c.mu.Unlock()
 	token := uuid.NewString()
@@ -114,21 +118,21 @@ func QueryRoomHierarchy(req *http.Request, device *userapi.Device, roomIDStr str
 
 	var walker roomserverAPI.RoomHierarchyWalker
 	if from == "" { // No pagination token provided, so start new hierarchy walker
-		walker = rsAPI.QueryRoomHierarchy(req.Context(), types.NewDeviceNotServerName(*device), roomID, suggestedOnly, maxDepth)
+		walker = roomserverAPI.NewRoomHierarchyWalker(types.NewDeviceNotServerName(*device), roomID, suggestedOnly, maxDepth)
 	} else { // Attempt to resume cached walker
 		cachedWalker := paginationCache.Get(from)
 
-		if cachedWalker == nil || !cachedWalker.ValidateParams(suggestedOnly, maxDepth) {
+		if cachedWalker == nil || cachedWalker.SuggestedOnly != suggestedOnly || cachedWalker.MaxDepth != maxDepth {
 			return util.JSONResponse{
 				Code: http.StatusBadRequest,
 				JSON: spec.InvalidParam("pagination not found for provided token ('from') with given 'max_depth', 'suggested_only' and room ID"),
 			}
 		}
 
-		walker = cachedWalker.GetWalker()
+		walker = *cachedWalker
 	}
 
-	discoveredRooms, err := walker.NextPage(limit)
+	discoveredRooms, nextWalker, err := rsAPI.QueryNextRoomHierarchyPage(req.Context(), walker, limit)
 
 	if err != nil {
 		switch err.(type) {
@@ -147,9 +151,9 @@ func QueryRoomHierarchy(req *http.Request, device *userapi.Device, roomIDStr str
 	}
 
 	nextBatch := ""
-	if !walker.Done() {
-		cacheLine := walker.GetCached()
-		nextBatch = paginationCache.AddLine(cacheLine)
+	// nextWalker will be nil if there's no more rooms left to walk
+	if nextWalker != nil {
+		nextBatch = paginationCache.AddLine(*nextWalker)
 	}
 
 	return util.JSONResponse{

federationapi/routing/query.go

@@ -145,9 +145,8 @@ func QueryRoomHierarchy(httpReq *http.Request, request *fclient.FederationReques
 		}
 	}
 
-	walker := rsAPI.QueryRoomHierarchy(httpReq.Context(), types.NewServerNameNotDevice(request.Origin()), roomID, suggestedOnly, 1)
-
-	discoveredRooms, err := walker.NextPage(-1)
+	walker := roomserverAPI.NewRoomHierarchyWalker(types.NewServerNameNotDevice(request.Origin()), roomID, suggestedOnly, 1)
+	discoveredRooms, _, err := rsAPI.QueryNextRoomHierarchyPage(httpReq.Context(), walker, -1)
 
 	if err != nil {
 		switch err.(type) {

roomserver/api/api.go

@@ -5,6 +5,7 @@ import (
 	"crypto/ed25519"
 
 	"github.com/matrix-org/gomatrixserverlib"
+	"github.com/matrix-org/gomatrixserverlib/fclient"
 	"github.com/matrix-org/gomatrixserverlib/spec"
 	"github.com/matrix-org/util"
 
@@ -123,7 +124,7 @@ type QueryEventsAPI interface {
 }
 
 type QueryRoomHierarchyAPI interface {
-	QueryRoomHierarchy(ctx context.Context, caller types.DeviceOrServerName, roomID spec.RoomID, suggestedOnly bool, maxDepth int) RoomHierarchyWalker
+	QueryNextRoomHierarchyPage(ctx context.Context, walker RoomHierarchyWalker, limit int) ([]fclient.MSC2946Room, *RoomHierarchyWalker, error)
 }
 
 // API functions required by the syncapi

roomserver/api/query.go

@@ -23,7 +23,6 @@ import (
 	"strings"
 
 	"github.com/matrix-org/gomatrixserverlib"
-	"github.com/matrix-org/gomatrixserverlib/fclient"
 	"github.com/matrix-org/gomatrixserverlib/spec"
 	"github.com/matrix-org/util"
 
@@ -510,24 +509,63 @@ type QueryRoomHierarchyRequest struct {
 	From          int  `json:"json"`
 }
 
-// An iterator-like interface for walking a room/space hierarchy, returning each rooms information.
+// A struct storing the intermediate state of a room hierarchy query for pagination purposes.
 //
 // Used for implementing space summaries / room hierarchies
-type RoomHierarchyWalker interface {
-	// Walk the room hierarchy to retrieve room information until either
-	// no room left, or provided limit reached. If limit provided is -1, then this is
-	// treated as no limit.
-	NextPage(limit int) ([]fclient.MSC2946Room, error)
-	// Returns true if there are no more rooms left to walk
-	Done() bool
-	// Returns a stripped down version of the hiearchy walker suitable for pagination caching
-	GetCached() CachedRoomHierarchyWalker
-}
-
-// Stripped down version of RoomHierarchyWalker suitable for caching (for pagination)
-type CachedRoomHierarchyWalker interface {
-	// Converts this cached walker back into an actual walker, to resume walking from.
-	GetWalker() RoomHierarchyWalker
-	// Validates that the given parameters match those stored in the cache
-	ValidateParams(suggestedOnly bool, maxDepth int) bool
+//
+// Use NewRoomHierarchyWalker on the roomserver API to construct this.
+type RoomHierarchyWalker struct {
+	RootRoomID    spec.RoomID
+	Caller        types.DeviceOrServerName
+	SuggestedOnly bool
+	MaxDepth      int
+	Processed     RoomSet
+	Unvisited     []RoomHierarchyWalkerQueuedRoom
+}
+
+type RoomHierarchyWalkerQueuedRoom struct {
+	RoomID       spec.RoomID
+	ParentRoomID *spec.RoomID
+	Depth        int
+	Vias         []string // vias to query this room by
+}
+
+func NewRoomHierarchyWalker(caller types.DeviceOrServerName, roomID spec.RoomID, suggestedOnly bool, maxDepth int) RoomHierarchyWalker {
+	walker := RoomHierarchyWalker{
+		RootRoomID:    roomID,
+		Caller:        caller,
+		SuggestedOnly: suggestedOnly,
+		MaxDepth:      maxDepth,
+		Unvisited: []RoomHierarchyWalkerQueuedRoom{{
+			RoomID:       roomID,
+			ParentRoomID: nil,
+			Depth:        0,
+		}},
+		Processed: NewRoomSet(),
+	}
+
+	return walker
+}
+
+type RoomSet map[spec.RoomID]struct{}
+
+func NewRoomSet() RoomSet {
+	return RoomSet{}
+}
+
+func (s RoomSet) Contains(val spec.RoomID) bool {
+	_, ok := s[val]
+	return ok
+}
+
+func (s RoomSet) Add(val spec.RoomID) {
+	s[val] = struct{}{}
+}
+
+func (s RoomSet) Copy() RoomSet {
+	copied := make(RoomSet, len(s))
+	for k, _ := range s {
+		copied.Add(k)
+	}
+	return copied
 }