Java GraphQL数据加载器批处理的实现详解

{
  "user": {
    "name": "zhangsan",
    “friends”: [
      {
        "name": "lisi",
      },
      {
        "name": "wanmgwu",
      },
      {
        "name": "zhouliu",
      }
    ]
  }
}

List<User> loadUsersById(List<Long> userIds) {
   System.out.println("Api call to load users = " + userIds);
   return users.stream().filter(u -> userIds.contains(u.id())).toList();
}

BatchLoader<Long, User> userBatchLoader = new BatchLoader<>() {
   @Override
   public CompletionStage<List<User>> load(List<Long> userIds) {
      return CompletableFuture.supplyAsync(() -> {
         return loadUsersById(userIds);
      });
   }
};

var userLoader = DataLoaderFactory.newDataLoader(userBatchLoader);

var userDTO = new UserDTO();
userLoader.load(1L).thenAccept(user -> {
   userDTO.id = user.id();
   userDTO.name = user.name();
   user.friends().forEach(friendId -> {
      userLoader.load(friendId).thenAccept(friend -> {
         userDTO.friends.add(new FriendDTO(friend.id(), friend.name()));
      });
   });
});

userLoader.dispatchAndJoin();
System.out.println(userDTO);

static class UserLoader {
   BatchLoader<Long, User> userBatchLoader;

   record QueueEntry(long id, CompletableFuture<User> value) { }
   List<QueueEntry> loaderQueue = new ArrayList<>();

   UserLoader(BatchLoader<Long, User> userBatchLoader) {
      this.userBatchLoader = userBatchLoader;
   }

   CompletableFuture<User> load(long userId) {
      var future = new CompletableFuture<User>();
      loaderQueue.add(new QueueEntry(userId, future));
      return future;
   }

   List<User> dispatchAndJoin() {
      List<User> joinedResults = dispatch().join();
      List<User> results = new ArrayList<>(joinedResults);
      while (loaderQueue.size() > 0) {
         joinedResults = dispatch().join();
         results.addAll(joinedResults);
      }
      return results;
   }

   CompletableFuture<List<User>> dispatch() {
      var userIds = new ArrayList<Long>();
      final List<CompletableFuture<User>> queuedFutures = new ArrayList<>();

      loaderQueue.forEach(qe -> {
         userIds.add(qe.id());
         queuedFutures.add(qe.value());
      });

      loaderQueue.clear();

      var userFutures = userBatchLoader.load(userIds).toCompletableFuture();

      return userFutures.thenApply(users -> {
         for (int i = 0; i < queuedFutures.size(); i++) {
            var userId = userIds.get(i);
            var user = users.get(i);
            var future = queuedFutures.get(i);
            future.complete(user);
         }
         return users;
      });
   }
}