com.aerospike.vector.client.adminclient.AdminClient Maven / Gradle / Ivy

Go to download
package com.aerospike.vector.client.adminclient;

import com.aerospike.vector.client.*;
import com.aerospike.vector.client.internal.HostPort;
import com.aerospike.vector.client.internal.ChannelProvider;
import com.google.protobuf.Empty;
import io.grpc.ManagedChannel;
import io.grpc.StatusRuntimeException;

import javax.annotation.Nullable;
import java.lang.Boolean;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.*;

/**
 * Implementation of Vector DB admin client
 */
public class AdminClient implements IAdminClient {

    private final ExecutorService adminExecutorService;
    private final ChannelProvider channelProvider;

    /**
     * Constructor for creating a new VectorDbAdminClientGrpc with specified parameters for the cluster.
     * @param seeds The seeds to initialize the connection.
     * @param listenerName Name of the listener for gRPC connection.
     * @param isLoadBalancer Whether to enable load balancing.
     */
    public AdminClient(List seeds, String listenerName, boolean isLoadBalancer) {
        this.channelProvider = new ChannelProvider(seeds, listenerName, isLoadBalancer);
        this.adminExecutorService = Executors.newThreadPerTaskExecutor(Thread.ofVirtual().name("avs-admin-", 0L).factory());
    }

    /**
     * {@inheritDoc}
     * This implementation checks if the index already exists before attempting to create it. If the index exists,
     * no action is taken. This check helps prevent unnecessary creation attempts that might result in exceptions
     * or errors from the database.
     * Note: This method uses a high-concurrency ExecutorService to manage gRPC calls, improving performance
     * under load but requiring proper shutdown management.
     *
     * @param indexId {@inheritDoc} Unique identifier for the index; must not be null.
     * @param vectorBinName {@inheritDoc} Name of the bin storing vector data; must not be null.
     * @param dimensions {@inheritDoc} Must be a positive integer, throws IllegalArgumentException otherwise.
     * @param vectorDistanceMetric {@inheritDoc}
     * @param setFilter {@inheritDoc} Can be null, which means no set filter is applied.
     * @param indexParams {@inheritDoc} Optional parameters for fine-tuning the index.
     * @param storage {@inheritDoc} Storage settings; can be null if default storage options are used.
     * @param labels {@inheritDoc} Must not be empty; required for indexing metadata.
     * @throws RuntimeException if an error occurs during the index creation process, including timeout exceptions.
     */
    @Override
    public void indexCreate(IndexId indexId, String vectorBinName, int dimensions,
                            VectorDistanceMetric vectorDistanceMetric, @Nullable String setFilter,
                            @Nullable HnswParams indexParams, @Nullable IndexStorage storage,
                            Map labels,
                            long timeoutInMillis,
                            long waitTimeInMillis) {
        Objects.requireNonNull(indexId, "Index ID cannot be null.");
        Objects.requireNonNull(vectorBinName, "Vector bin name cannot be null.");
        Objects.requireNonNull(vectorDistanceMetric, "Vector distance metric cannot be null.");
        Objects.requireNonNull(labels, "Labels cannot be null.");
        if (dimensions <= 0) {
            throw new IllegalArgumentException("Dimensions must be a positive integer.");
        }
        ManagedChannel channel = (ManagedChannel) channelProvider.getChannel();
        IndexServiceGrpc.IndexServiceBlockingStub indexService = IndexServiceGrpc.newBlockingStub(channel).withExecutor(adminExecutorService);

        if (indexExists(indexService, indexId)) {
            return; // Index already exists, so we do nothing
        }

        IndexDefinition.Builder builder = IndexDefinition.newBuilder()
                .setField(vectorBinName)
                .setId(indexId)
                .setVectorDistanceMetric(vectorDistanceMetric)
                .setDimensions(dimensions);

        if (setFilter != null) {
            builder.setSetFilter(setFilter);
        }

        if (indexParams != null) {
            builder.setHnswParams(indexParams);
        }

        if (storage != null) {
            builder.setStorage(storage);
        }

        if (!labels.isEmpty()) {
            builder.putAllLabels(labels);
        }

        indexService.create(builder.build());

        try {
            waitForIndexCreation(indexService, indexId, timeoutInMillis, waitTimeInMillis);
        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        }
    }

    private boolean indexExists(IndexServiceGrpc.IndexServiceBlockingStub indexService, IndexId indexId) {

        // Check if index already exists
        List existingIndices = indexService.list(Empty.getDefaultInstance()).getIndicesList();
        return existingIndices.stream().anyMatch(def -> def.getId().equals(indexId));
    }


    @Override
    public void indexDrop(IndexId indexId, long timeoutInMillis, long waitTimeInMillis ) {
        ManagedChannel channel = (ManagedChannel) channelProvider.getChannel();
        IndexServiceGrpc.IndexServiceBlockingStub indexService = IndexServiceGrpc.newBlockingStub(channel).withExecutor(adminExecutorService);
        indexService.drop(indexId);

        try {
            waitForIndexDeletion(indexService, indexId, timeoutInMillis, waitTimeInMillis);
        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * List all available vector index
     * @return list all index related information
     */
    @Override
    public List indexList() {
        ManagedChannel channel = (ManagedChannel) channelProvider.getChannel();
        IndexServiceGrpc.IndexServiceBlockingStub indexService = IndexServiceGrpc.newBlockingStub(channel).withExecutor(adminExecutorService);
        return indexService.list(Empty.getDefaultInstance()).getIndicesList();
    }

    /**
     * Number of unmerged index records in the given index
     * @param indexId Unique identifier for the index.
     * @return number of unmerged record in the index
     */
    @Override
    public IndexStatusResponse indexStatus(IndexId indexId) {
        ManagedChannel channel = (ManagedChannel) channelProvider.getChannel();
        IndexServiceGrpc.IndexServiceBlockingStub indexService = IndexServiceGrpc.newBlockingStub(channel).withExecutor(adminExecutorService);
        return indexService.getStatus(indexId);
    }

    @Override
    public IndexDefinition getIndex(IndexId indexId) {
        ManagedChannel channel = (ManagedChannel) channelProvider.getChannel();
        IndexServiceGrpc.IndexServiceBlockingStub indexService = IndexServiceGrpc.newBlockingStub(channel).withExecutor(adminExecutorService);
        return indexService.get(indexId);
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public void close() {
        try {
            if (channelProvider != null) {
                channelProvider.close();
            }
            if (adminExecutorService != null && !adminExecutorService.isShutdown()) {
                adminExecutorService.shutdown();
            }
        } catch (Exception e) {
            throw new RuntimeException("Failed to close resources properly", e);
        }
    }

    private void waitForIndexCreation(IndexServiceGrpc.IndexServiceBlockingStub indexService, IndexId indexId, long timeoutInMillis, long waitTimeInMillis) throws InterruptedException {
        Callable task = () -> {
            long startTime = System.currentTimeMillis();
            while (System.currentTimeMillis() - startTime < timeoutInMillis) {
                try {
                    if (getIndex(indexId) != null) {
                        return true; // Index is found
                    }
                } catch (StatusRuntimeException e) {
                    // Handle the case where the index is not found yet
                    Thread.sleep(waitTimeInMillis);
                }
            }
            return false;
        };

        Future future = adminExecutorService.submit(task);
        try {
            if (!future.get()) {
                throw new TimeoutException("Failed to verify index creation within the timeout period.");
            }
        } catch (InterruptedException | ExecutionException | TimeoutException e) {
            throw new RuntimeException("Error waiting for index to be created", e);
        }
        //Not shutting down executor because other tasks may be using it.
    }

    private void waitForIndexDeletion(IndexServiceGrpc.IndexServiceBlockingStub indexService,IndexId indexId, long timeoutInMillis, long waitTime) throws InterruptedException {

            long startTime = System.currentTimeMillis();
            while (System.currentTimeMillis() - startTime < timeoutInMillis) {
                //if index is not in the list then it was deleted
                if (!indexExists(indexService, indexId) ) {
                    return;
                }else {
                    Thread.sleep(waitTime);
                }
            }

            throw  new RuntimeException(String.format("Timed out in %s index deletion", indexId));
        }
    }