Add ability to limit size of grpc request for getTree

buildgrid/server/cas/instance.py

@@ -75,23 +75,35 @@ class ContentAddressableStorageInstance:
     def get_tree(self, request):
         storage = self._storage
 
-        # Create getTreeResponse message
-        # TODO: handle page_token, page_size
         response = re_pb2.GetTreeResponse()
+        page_size = request.page_size
+        next_page_token = ""
+        directories = []
 
-        # Set to MAX_REQUEST_COUNT, will use in the future to limit size of GRPC response
-        if not request.page_size:
-            request.page_size = 500
+        if not page_size:
+            # MAX_REQUEST_COUNT
+            page_size = 500
 
         def _get_tree(node_digest):
+            nonlocal directories, page_size, next_page_token
+            if next_page_token:
+                # next page token has been set unwind the stack.
+                return
+            if page_size <= 0:
+                # save the next digest hash in order to continue later
+                next_page_token = str(node_digest.hash)
+                return
             directory_from_digest = storage.get_message(node_digest, re_pb2.Directory)
-            directories = [directory_from_digest]
+            directories.append(directory_from_digest)
             for directory in directory_from_digest.directories:
-                directories.extend(_get_tree(directory.digest))
-            return directories
+                page_size -= 1
+                _get_tree(directory.digest)
+            return
 
-        response.directories.extend(_get_tree(request.root_digest))
-        yield response
+        _get_tree(request.root_digest)
+        response.directories.extend(directories)
+        response.next_page_token = next_page_token
+        return response
 
 
 class ByteStreamInstance:

buildgrid/server/cas/service.py

@@ -88,9 +88,17 @@ class ContentAddressableStorageService(remote_execution_pb2_grpc.ContentAddressa
 
         try:
             instance = self._get_instance(request.instance_name)
-            response = instance.get_tree(request)
-
-            return response
+            while True:
+                response = instance.get_tree(request)
+                if not response.next_page_token:
+                    # This is ugly, but handles the case in which there is only one iteration
+                    # of the while loop. We need to manually raise stopiteration from this generator.
+                    # If return is called, before a yield is ever called, we get an empty response.
+                    yield response
+                    raise StopIteration  # pylint: disable=stop-iteration-return
+
+                yield response
+                request.root_digest.hash = response.next_page_token
 
         except InvalidArgumentError as e:
             self.__logger.error(e)