package com.atlassian.bitbucket.mesh.management;

import com.atlassian.bitbucket.mesh.AbstractStatusException;
import com.atlassian.bitbucket.mesh.annotation.RequiresRole;
import com.atlassian.bitbucket.mesh.auth.ApplicationRole;
import com.atlassian.bitbucket.mesh.auth.ExpirablePublicKey;
import com.atlassian.bitbucket.mesh.auth.InvalidPublicKeyException;
import com.atlassian.bitbucket.mesh.auth.UpdatePublicKeysParameters;
import com.atlassian.bitbucket.mesh.execution.GrpcAsyncExecutor;
import com.atlassian.bitbucket.mesh.git.exception.InvalidArgumentException;
import com.atlassian.bitbucket.mesh.grpc.BackoffStreamObserver;
import com.atlassian.bitbucket.mesh.io.MoreFiles;
import com.atlassian.bitbucket.mesh.network.*;
import com.atlassian.bitbucket.mesh.repair.RepairManager;
import com.atlassian.bitbucket.mesh.repository.*;
import com.atlassian.bitbucket.mesh.rpc.v1.*;
import com.atlassian.bitbucket.mesh.util.KeyUtils;
import com.google.common.collect.Collections2;
import com.google.common.collect.Lists;
import com.google.common.io.ByteStreams;
import com.google.protobuf.ByteString;
import com.google.protobuf.UnsafeByteOperations;
import io.grpc.Context;
import io.grpc.Deadline;
import io.grpc.Status;
import io.grpc.StatusRuntimeException;
import io.grpc.stub.StreamObserver;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.io.InputStream;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.security.PublicKey;
import java.time.Duration;
import java.util.*;
import java.util.concurrent.CancellationException;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionStage;
import java.util.function.BooleanSupplier;
import java.util.stream.Collectors;

import static com.atlassian.bitbucket.mesh.rpc.ManagedConfigurationProperties.NODE_ID;
import static com.atlassian.bitbucket.mesh.rpc.v1.RpcInvalidPublicKey.Reason.EXPIRY_NOT_ALLOWED;
import static com.atlassian.bitbucket.mesh.rpc.v1.RpcInvalidPublicKey.Reason.INVALID;
import static org.apache.commons.lang3.StringUtils.isBlank;
import static org.apache.commons.lang3.StringUtils.trimToNull;

/**
 * RPC service that is used by the control plane to manage the node
 */
public class GrpcManagementService extends ManagementServiceGrpc.ManagementServiceImplBase {

    private static final Logger log = LoggerFactory.getLogger(GrpcManagementService.class);

    private final GrpcAsyncExecutor asyncExecutor;
    private final RepositoryAttributeManager attributeManager;
    private final DataStoreManager dataStoreManager;
    private final MeshManager meshManager;
    private final NodeManager nodeManager;
    private final RepairManager repairManager;
    private final ReplicaStateRegistry replicaStateRegistry;
    private final RpcRepositoryAccessTracker repositoryAccessTracker;
    private final RepositoryManager repositoryManager;
    private final SidebandManager sidebandManager;

    public GrpcManagementService(GrpcAsyncExecutor asyncExecutor, RepositoryAttributeManager attributeManager,
                                 DataStoreManager dataStoreManager, MeshManager meshManager, NodeManager nodeManager,
                                 RepairManager repairManager, ReplicaStateRegistry replicaStateRegistry,
                                 RpcRepositoryAccessTracker repositoryAccessTracker,
                                 RepositoryManager repositoryManager,
                                 SidebandManager sidebandManager) {
        this.asyncExecutor = asyncExecutor;
        this.attributeManager = attributeManager;
        this.dataStoreManager = dataStoreManager;
        this.meshManager = meshManager;
        this.nodeManager = nodeManager;
        this.repairManager = repairManager;
        this.repositoryAccessTracker = repositoryAccessTracker;
        this.replicaStateRegistry = replicaStateRegistry;
        this.repositoryManager = repositoryManager;
        this.sidebandManager = sidebandManager;
    }

    public static Duration toDuration(com.google.protobuf.Duration duration) {
        return Duration.ofSeconds(duration.getSeconds(), duration.getNanos());
    }

    @RequiresRole(ApplicationRole.CONTROL_PLANE)
    @Override
    public void addDataStore(RpcAddDataStoreRequest request, StreamObserver<RpcAddDataStoreResponse> responseObserver) {
        RpcAddDataStoreResponse response = dataStoreManager.addShard(buildAddShardParameters(request)) ?
                RpcAddDataStoreResponse.newBuilder().setSuccess(true).build() :
                RpcAddDataStoreResponse.getDefaultInstance();

        responseObserver.onNext(response);
        responseObserver.onCompleted();
    }

    @RequiresRole(ApplicationRole.CONTROL_PLANE)
    @Override
    public void drain(RpcDrainRequest request, StreamObserver<RpcDrainResponseFragment> responseObserver) {
        switch (request.getResourceOneofCase()) {
            case ALL_REQUEST:
                drainAll(responseObserver);
                break;
            case HIERARCHIES_REQUEST:
                drainHierarchies(request.getHierarchiesRequest().getHierarchiesList(), responseObserver);
                break;
            case PARTITIONS_REQUEST:
                drainPartitions(request.getPartitionsRequest().getPartitionsList(), responseObserver);
                break;
            default:
                throw new InvalidArgumentException("Received [" + request.getResourceOneofCase() +
                        "] where [ALL_REQUEST, HIERARCHIES_REQUEST or PARTITIONS_REQUEST] was expected");

        }
    }

    @RequiresRole(ApplicationRole.CONTROL_PLANE)
    @Override
    public void getConfiguration(RpcGetConfigurationRequest request,
                                 StreamObserver<RpcGetConfigurationResponse> responseObserver) {
        RpcGetConfigurationResponse.Builder builder = RpcGetConfigurationResponse.newBuilder()
                //Add the current node configuration
                .setNodeConfiguration(RpcNodeConfiguration.newBuilder()
                        .putAllProperties(nodeManager.getConfigurationProperties()))
                //Include the topology
                .setTopology(RpcTopology.newBuilder()
                        .addAllNodes(Collections2.transform(meshManager.getNodes(), MeshNode::toNode))
                        .addAllPartitions(Collections2.transform(meshManager.getPartitions(), Partition::toPartition))
                        .addAllPartitionMigrations(Collections2.transform(meshManager.getPartitionMigrations(),
                                PartitionReplicaMigration::toPartitionReplicaMigration)));
        //Add currently inconsistent replicas
        replicaStateRegistry.streamInconsistentReplicas((repositoryId, observedVersion) ->
                builder.addInconsistentReplicas(RpcInconsistentReplica.newBuilder()
                        .setRepository(repositoryId)
                        .setVersion(observedVersion)));
        responseObserver.onNext(builder.build());
        responseObserver.onCompleted();
    }

    @RequiresRole(ApplicationRole.CONTROL_PLANE)
    @Override
    public void getDataStores(RpcGetDataStoresRequest request,
                              StreamObserver<RpcGetDataStoresResponse> responseObserver) {
        List<DataStoreInfo> dataStores = dataStoreManager.getInfo();
        RpcGetDataStoresResponse response = RpcGetDataStoresResponse.newBuilder()
                .addAllDataStores(Lists.transform(dataStores, DataStoreInfo::toDataStore))
                .build();
        responseObserver.onNext(response);
        responseObserver.onCompleted();
    }

    @Override
    public void getNodeInfo(RpcGetNodeInfoRequest ignored, StreamObserver<RpcGetNodeInfoResponse> responseObserver) {
        responseObserver.onNext(RpcGetNodeInfoResponse.newBuilder()
                .putAllProperties(nodeManager.getNodeInfo())
                .build());
        responseObserver.onCompleted();
    }

    @RequiresRole(ApplicationRole.CONTROL_PLANE)
    @Override
    public void getSupportInfo(RpcGetSupportInfoRequest ignored,
                               StreamObserver<RpcGetSupportInfoResponse> responseObserver) {
        responseObserver.onNext(nodeManager.getSupportInfo().toSupportInfoResponse());
        responseObserver.onCompleted();
    }

    @RequiresRole(ApplicationRole.CONTROL_PLANE)
    @Override
    public void getSupportZip(RpcGetSupportZipRequest request,
                              StreamObserver<RpcStreamChunk> responseObserver) {
        CreateSupportZipParameters parameters = new CreateSupportZipParameters(request.getConfig(), request.getLogs());
        BackoffStreamObserver<RpcStreamChunk> backoffObserver = new BackoffStreamObserver<>(responseObserver);
        asyncExecutor.execute(responseObserver, () -> createSupportZip(parameters, backoffObserver));
    }

    @Override
    public void heartbeat(RpcHeartbeat request, StreamObserver<RpcHeartbeat> responseObserver) {
        responseObserver.onNext(RpcHeartbeat.getDefaultInstance());
        responseObserver.onCompleted();
    }

    @RequiresRole(ApplicationRole.CONTROL_PLANE)
    @Override
    public void markRepositoriesInconsistent(RpcMarkRepositoriesInconsistentRequest request,
                                             StreamObserver<RpcMarkRepositoriesInconsistentResponse> responseObserver) {
        replicaStateRegistry.markInconsistent(new MarkInconsistentParameters.Builder()
                .complete(false)
                .fromControlPlane(true)
                .inconsistentReplicas(request.getInconsistentReplicasList())
                .build());

        responseObserver.onNext(RpcMarkRepositoriesInconsistentResponse.getDefaultInstance());
        responseObserver.onCompleted();
    }

    @Override
    public void provisionNode(RpcProvisionNodeRequest request,
                              StreamObserver<RpcProvisionNodeResponse> responseObserver) {
        PublicKey publicKey = nodeManager.provision(buildProvisionParameters(request));
        RpcProvisionNodeResponse response = publicKey == null ?
                RpcProvisionNodeResponse.getDefaultInstance() :
                RpcProvisionNodeResponse.newBuilder()
                        .setPublicKey(toPublicKey(publicKey))
                        .build();

        responseObserver.onNext(response);
        responseObserver.onCompleted();
    }

    @RequiresRole(ApplicationRole.CONTROL_PLANE)
    @Override
    public void repairRepository(RpcRepairRepositoryRequest request,
                                 StreamObserver<RpcRepairRepositoryResponse> responseObserver) {
        Collection<MeshNode> sourceNodes;
        MeshNode targetNode;
        if (request.hasSources()) {
            targetNode = meshManager.getLocalNode();
            sourceNodes = request.getSources().getNodesList().stream()
                    .map(meshManager::getNodeById)
                    .collect(Collectors.toList());
        } else {
            targetNode = meshManager.getNodeById(request.getTarget());
            sourceNodes = Collections.singleton(meshManager.getLocalNode());
        }
        try {
            String targetRepository = request.getRepository();
            String sourceRepository = request.hasSourceRepository() ? request.getSourceRepository() : targetRepository;
            ReplicaState result = repairManager.repair(sourceRepository, sourceNodes, targetRepository, targetNode);

            responseObserver.onNext(RpcRepairRepositoryResponse.newBuilder()
                    .setState(result.toRpc())
                    .build());
            responseObserver.onCompleted();
        } catch (IllegalArgumentException e) {
            responseObserver.onError(Status.FAILED_PRECONDITION
                    .withCause(e)
                    .withDescription(e.getMessage())
                    .asRuntimeException());
        }
    }

    @RequiresRole(ApplicationRole.CONTROL_PLANE)
    @Override
    public StreamObserver<RpcReplaceReadOnlyRepositoriesFragment> replaceReadOnlyRepositories(
            StreamObserver<RpcReplaceReadOnlyRepositoriesResponse> responseObserver) {
        return new ReplaceReadOnlyRepositoriesRequestObserver(responseObserver, attributeManager);
    }

    @RequiresRole(ApplicationRole.CONTROL_PLANE)
    @Override
    public void resetHierarchy(RpcResetHierarchyRequest request,
                               StreamObserver<RpcResetHierarchyResponse> responseObserver) {
        int count = repositoryManager.resetHierarchy(request.getPartition(), request.getHierarchy());
        responseObserver.onNext(RpcResetHierarchyResponse.newBuilder()
                .setResetCount(count)
                .build());
        responseObserver.onCompleted();
    }

    @RequiresRole(ApplicationRole.CONTROL_PLANE)
    @Override
    public void setConfiguration(RpcSetConfigurationRequest request,
                                 StreamObserver<RpcSetConfigurationResponse> responseObserver) {
        nodeManager.updateConfiguration(buildUpdateConfigurationParameters(request));
        responseObserver.onNext(RpcSetConfigurationResponse.getDefaultInstance());
        responseObserver.onCompleted();
    }

    @RequiresRole(ApplicationRole.CONTROL_PLANE)
    @Override
    public void setReplicaState(RpcSetReplicaStateRequest request,
                                StreamObserver<RpcSetReplicaStateResponse> responseObserver) {
        long version;
        if (request.hasVersion()) {
            version = request.getVersion();
        } else {
            Repository repository = repositoryManager.findById(request.getRepository());
            if (repository instanceof ManagedRepository) {
                version = ((ManagedRepository) repository).getVersion();
            } else {
                //Other repository types have no versions, so treat it as zero
                version = 0L;
            }
        }
        ReplicaState state = ReplicaState.fromRpc(request.getReplicaState());
        replicaStateRegistry.setState(request.getRepository(), state, version, request.getForce());
        responseObserver.onNext(RpcSetReplicaStateResponse.getDefaultInstance());
        responseObserver.onCompleted();
    }

    @RequiresRole(ApplicationRole.CONTROL_PLANE)
    @Override
    public void setRepositoryReadOnly(RpcSetRepositoryReadOnlyRequest request,
                                      StreamObserver<RpcSetRepositoryReadOnlyResponse> responseObserver) {
        Repository repository = repositoryManager.getById(request.getRepository());
        repository.setReadOnly(request.getReadOnly());

        responseObserver.onNext(RpcSetRepositoryReadOnlyResponse.getDefaultInstance());
        responseObserver.onCompleted();
    }

    @RequiresRole(ApplicationRole.CONTROL_PLANE)
    @Override
    public StreamObserver<RpcSidebandLeaderFragment> sideband(
            StreamObserver<RpcSidebandFollowerFragment> responseObserver) {
        SidebandSender sender = new ObserverSidebandSender(responseObserver);
        SidebandReceiver receiver = sidebandManager.open(new OpenSidebandParameters.Builder(sender).build());

        return new SidebandRequestObserver(receiver);
    }

    private static AddShardParameters buildAddShardParameters(RpcAddDataStoreRequest request) {
        return new AddShardParameters.Builder(request.getId(), Paths.get(request.getPath()))
                .build();
    }

    private static ProvisionParameters buildProvisionParameters(RpcProvisionNodeRequest request) {
        if (request.hasPublicKey()) {
            RpcPublicKey publicKey = request.getPublicKey();
            if (publicKey.hasExpiryInterval()) {
                log.warn("Rejecting provisioned public key; an expiry interval is set");
                throw new InvalidPublicKeyException(EXPIRY_NOT_ALLOWED, publicKey);
            }

            return new ProvisionParameters.Builder(toPublicKey(publicKey))
                    .nodeId(trimToNull(request.getNodeId()))
                    .build();
        }

        throw new InvalidArgumentException("A public key must be provided when provisioning a node");
    }

    private static UpdateTopologyParameters buildUpdateTopologyParameters(String localNodeId, RpcTopology topology) {
        //First construct all nodes
        Map<String, MeshNode> nodes = new HashMap<>();
        for (RpcNode node : topology.getNodesList()) {
            nodes.put(node.getId(), new MeshNode.Builder(node.getId())
                    .local(node.getId().equals(localNodeId))
                    .name(node.getName())
                    .rpcUrl(node.getRpcUrl())
                    .build());
        }
        UpdateTopologyParameters.Builder builder = new UpdateTopologyParameters.Builder()
                .nodes(nodes.values());

        //Add all the partitions
        for (RpcPartition partition : topology.getPartitionsList()) {
            String partitionId = partition.getId();
            Partition.Builder partitionBuilder = new Partition.Builder(partitionId);
            for (RpcPartitionReplica replica : partition.getReplicasList()) {
                MeshNode node = nodes.get(replica.getNodeId());
                if (node == null) {
                    log.warn("Topology references node {} for replica {} of partition {} but node is not defined. " +
                            "Skipping partition replica", replica.getNodeId(), replica.getId(), partitionId);
                } else {
                    partitionBuilder.replicas(new PartitionReplica.Builder(replica.getId(), node).build());
                }
            }
            builder.partitions(partitionBuilder.build());
        }

        //Add all migrations relevant to this node
        for (RpcPartitionReplicaMigration migration : topology.getPartitionMigrationsList()) {
            MeshNode sourceNode = nodes.get(migration.getSourceNodeId());
            if (sourceNode == null) {
                log.warn("Partition migration for partition {} references node {} which is not defined. " +
                        "Ignoring partition migration", migration.getPartition(), migration.getSourceNodeId());
                continue;
            }
            MeshNode targetNode = nodes.get(migration.getTargetNodeId());
            if (targetNode == null) {
                log.warn("Partition migration for partition {} references node {} which is not defined. " +
                        "Ignoring partition migration", migration.getPartition(), migration.getTargetNodeId());
                continue;
            }
            builder.migrations(new PartitionReplicaMigration.Builder()
                    .partition(migration.getPartition())
                    .sourceNode(sourceNode)
                    .targetNode(targetNode)
                    .build());
        }
        return builder.build();
    }

    private static ExpirablePublicKey toExpirablePublicKey(RpcPublicKey key) {
        return new ExpirablePublicKey.Builder(toPublicKey(key))
                .expiryInterval(key.hasExpiryInterval() ? toDuration(key.getExpiryInterval()) : null)
                .build();
    }

    private static PublicKey toPublicKey(RpcPublicKey publicKey) {
        try {
            return publicKey.hasDer() ?
                    KeyUtils.parsePublicKeyFromDer(publicKey.getDer().toByteArray()) :
                    KeyUtils.parsePublicKeyFromPkcs8(publicKey.getPkcs8());
        } catch (IllegalArgumentException e) {
            throw new InvalidPublicKeyException(INVALID, publicKey, e);
        }
    }

    private static RpcPublicKey toPublicKey(PublicKey publicKey) {
        return RpcPublicKey.newBuilder()
                .setDer(ByteString.copyFrom(publicKey.getEncoded()))
                .build();
    }

    private static StatusRuntimeException toStatusException(Throwable t) {
        if (t instanceof StatusRuntimeException) {
            return (StatusRuntimeException) t;
        }
        if (t instanceof AbstractStatusException) {
            return ((AbstractStatusException) t).toStatusException();
        }
        if (t instanceof CancellationException) {
            return Status.CANCELLED.withCause(t).asRuntimeException();
        }
        return Status.INTERNAL.withCause(t).asRuntimeException();
    }

    private UpdateConfigurationParameters buildUpdateConfigurationParameters(RpcSetConfigurationRequest request) {
        UpdateConfigurationParameters.Builder builder = new UpdateConfigurationParameters.Builder()
                .fromControlPlane(true);
        request.getInconsistentReplicasList().forEach(inconsistentReplica ->
                builder.inconsistentReplica(inconsistentReplica.getRepository(), inconsistentReplica.getVersion()));

        String nodeId = null;
        if (request.hasNodeConfiguration()) {
            Map<String, String> properties = request.getNodeConfiguration().getPropertiesMap();
            builder.properties(properties);

            nodeId = properties.get(NODE_ID);
        }
        if (request.hasPublicKeyConfiguration()) {
            builder.publicKeys(buildUpdatePublicKeysParameters(request.getPublicKeyConfiguration()));
        }
        if (request.hasTopology()) {
            RpcTopology topology = request.getTopology();
            //When constructing the topology, we need to create the correct node as the local node.
            //If the configuration request also changes the node's ID (which it will on the first time the
            //node is configured after being registered), we need to use the new ID.
            builder.topology(buildUpdateTopologyParameters(isBlank(nodeId) ? nodeManager.getId() : nodeId, topology));
        }

        return builder.build();
    }

    private UpdatePublicKeysParameters buildUpdatePublicKeysParameters(RpcPublicKeyConfiguration configuration) {
        UpdatePublicKeysParameters.Builder builder = new UpdatePublicKeysParameters.Builder()
                .complete(configuration.getComplete());
        if (configuration.hasControlPlaneKeys()) {
            for (RpcPublicKey publicKey : configuration.getControlPlaneKeys().getPublicKeysList()) {
                builder.controlPlaneKey(toExpirablePublicKey(publicKey));
            }
        }
        if (configuration.hasSidecarKeys()) {
            for (RpcPublicKey publicKey : configuration.getSidecarKeys().getPublicKeysList()) {
                builder.sidecarKey(toExpirablePublicKey(publicKey));
            }
        }
        for (Map.Entry<String, RpcPublicKeyList> entry : configuration.getNodePublicKeysMap().entrySet()) {
            String nodeId = entry.getKey();
            for (RpcPublicKey publicKey : entry.getValue().getPublicKeysList()) {
                builder.nodeKey(nodeId, toExpirablePublicKey(publicKey));
            }
        }

        return builder.build();
    }

    private void createSupportZip(CreateSupportZipParameters parameters,
                                  StreamObserver<RpcStreamChunk> responseObserver) {
        Path supportZip = null;
        try {
            log.debug("Starting to create the support zip");
            supportZip = nodeManager.createSupportZip(parameters);
            log.debug("Created the support zip and starting the transfer");

            try (InputStream is = Files.newInputStream(supportZip)) {
                boolean eof = false;
                long start = System.currentTimeMillis();
                byte[] zipBuffer = new byte[65536];
                long zipSize = 0L;

                RpcStreamChunk.Builder builder = RpcStreamChunk.newBuilder();
                Deadline deadline = Context.current().getDeadline();
                BooleanSupplier timedOut = deadline == null ? () -> false : deadline::isExpired;
                while (!eof && !timedOut.getAsBoolean()) {
                    builder.clear();
                    int bytesRead = ByteStreams.read(is, zipBuffer, 0, zipBuffer.length);
                    zipSize += bytesRead;
                    eof = bytesRead < zipBuffer.length;
                    builder.setClosed(eof)
                            .setData(UnsafeByteOperations.unsafeWrap(zipBuffer, 0, bytesRead));
                    responseObserver.onNext(builder.build());
                }

                responseObserver.onCompleted();
                log.debug("Transferred {} byte support zip in {} ms", zipSize, (System.currentTimeMillis() - start));
            } catch (IOException e) {
                log.warn("Could not (completely) transfer the support zip", e);
                throw new SupportZipFailedException("Read failure while transferring support zip", e);
            }
        } finally {
            if (supportZip != null) {
                MoreFiles.deleteQuietly(supportZip);
            }
        }
    }

    private void drainAll(StreamObserver<RpcDrainResponseFragment> observer) {
        //Start async draining and forward drain
        repositoryAccessTracker.drainAll()
                .handle((result, exception) -> {
                    if (exception == null) {
                        observer.onNext(RpcDrainResponseFragment.newBuilder()
                                .setDrainedResource(RpcDrainedResource.getDefaultInstance())
                                .build());
                        observer.onCompleted();
                    } else {
                        observer.onError(toStatusException(exception));
                    }
                    return null;
                });
    }

    private void drainHierarchies(List<RpcHierarchyRequest> hierarchies,
                                  StreamObserver<RpcDrainResponseFragment> observer) {
        CompletionStage<Void> future = CompletableFuture.completedFuture(null);
        for (RpcHierarchyRequest request : hierarchies) {
            CompletionStage<Void> drainStage =
                    repositoryAccessTracker.drainHierarchy(request.getPartition(), request.getHierarchy())
                            .handle((result, exception) -> {
                                if (exception == null) {
                                    observer.onNext(RpcDrainResponseFragment.newBuilder()
                                            .setDrainedResource(RpcDrainedResource.newBuilder()
                                                    .setPartition(request.getPartition())
                                                    .setHierarchy(request.getHierarchy()))
                                            .build());
                                } else {
                                    observer.onError(toStatusException(exception));
                                }
                                return null;
                            });
            future = future.thenCompose(result -> drainStage);
        }
        //When all hierarchies have drained, call onCompleted
        future.thenRun(observer::onCompleted);
    }

    private void drainPartitions(List<String> partitions, StreamObserver<RpcDrainResponseFragment> observer) {
        CompletionStage<Void> future = CompletableFuture.completedFuture(null);
        for (String partition : partitions) {
            CompletionStage<Void> drainStage = repositoryAccessTracker.drainPartition(partition)
                    .handle((result, exception) -> {
                        if (exception == null) {
                            observer.onNext(RpcDrainResponseFragment.newBuilder()
                                    .setDrainedResource(RpcDrainedResource.newBuilder()
                                            .setPartition(partition))
                                    .build());
                        } else {
                            observer.onError(toStatusException(exception));
                        }
                        return null;
                    });
            future = future.thenCompose(result -> drainStage);
        }
        //When all partitions have drained, call onCompleted
        future.thenRun(observer::onCompleted);
    }
}
