Examples with Registry org.apache.drill.exec.proto.UserBitShared.Registry used on opensource projects

Example 6 with Registry

use of org.apache.drill.exec.proto.UserBitShared.Registry in project drill by apache.

the class TestDynamicUDFSupport method testConcurrentRemoteRegistryUpdateForDifferentJars.

@Test
public void testConcurrentRemoteRegistryUpdateForDifferentJars() throws Exception {
    RemoteFunctionRegistry remoteFunctionRegistry = spyRemoteFunctionRegistry();
    final CountDownLatch latch1 = new CountDownLatch(1);
    final CountDownLatch latch2 = new CountDownLatch(2);
    doAnswer(new Answer<Void>() {

        @Override
        public Void answer(InvocationOnMock invocation) throws Throwable {
            latch2.countDown();
            latch1.await();
            invocation.callRealMethod();
            return null;
        }
    }).when(remoteFunctionRegistry).updateRegistry(any(Registry.class), any(DataChangeVersion.class));
    final String jarName1 = default_binary_name;
    final String jarName2 = "DrillUDF-2.0.jar";
    final String query = "create function using jar '%s'";
    copyDefaultJarsToStagingArea();
    copyJarsToStagingArea(jarName2, JarUtil.getSourceName(jarName2));
    Thread thread1 = new Thread(new TestBuilderRunner(testBuilder().sqlQuery(query, jarName1).unOrdered().baselineColumns("ok", "summary").baselineValues(true, String.format("The following UDFs in jar %s have been registered:\n" + "[custom_lower(VARCHAR-REQUIRED)]", jarName1))));
    Thread thread2 = new Thread(new TestBuilderRunner(testBuilder().sqlQuery(query, jarName2).unOrdered().baselineColumns("ok", "summary").baselineValues(true, String.format("The following UDFs in jar %s have been registered:\n" + "[custom_upper(VARCHAR-REQUIRED)]", jarName2))));
    thread1.start();
    thread2.start();
    latch2.await();
    latch1.countDown();
    thread1.join();
    thread2.join();
    DataChangeVersion version = new DataChangeVersion();
    Registry registry = remoteFunctionRegistry.getRegistry(version);
    assertEquals("Remote registry version should match", 2, version.getVersion());
    List<Jar> actualJars = registry.getJarList();
    List<String> expectedJars = Lists.newArrayList(jarName1, jarName2);
    assertEquals("Only one jar should be registered", 2, actualJars.size());
    for (Jar jar : actualJars) {
        assertTrue("Jar should be present in remote function registry", expectedJars.contains(jar.getName()));
    }
    verify(remoteFunctionRegistry, times(3)).updateRegistry(any(Registry.class), any(DataChangeVersion.class));
}

Example 7 with Registry

use of org.apache.drill.exec.proto.UserBitShared.Registry in project drill by axbaretto.

the class RemoteFunctionRegistry method prepareStores.

/**
 * Connects to three stores: REGISTRY, UNREGISTRATION, JARS.
 * Puts in REGISTRY store with default instance of remote function registry if store is initiated for the first time.
 * Registers unregistration listener in UNREGISTRATION store.
 */
private void prepareStores(PersistentStoreProvider storeProvider, ClusterCoordinator coordinator) {
    try {
        PersistentStoreConfig<Registry> registrationConfig = PersistentStoreConfig.newProtoBuilder(SchemaUserBitShared.Registry.WRITE, SchemaUserBitShared.Registry.MERGE).name("udf").persist().build();
        registry = storeProvider.getOrCreateVersionedStore(registrationConfig);
        registry.putIfAbsent(registry_path, Registry.getDefaultInstance());
    } catch (StoreException e) {
        throw new DrillRuntimeException("Failure while loading remote registry.", e);
    }
    TransientStoreConfig<String> unregistrationConfig = TransientStoreConfig.newJacksonBuilder(mapper, String.class).name("udf/unregister").build();
    unregistration = coordinator.getOrCreateTransientStore(unregistrationConfig);
    unregistration.addListener(unregistrationListener);
    TransientStoreConfig<String> jarsConfig = TransientStoreConfig.newJacksonBuilder(mapper, String.class).name("udf/jars").build();
    jars = coordinator.getOrCreateTransientStore(jarsConfig);
}

Example 8 with Registry

use of org.apache.drill.exec.proto.UserBitShared.Registry in project drill by axbaretto.

the class CreateFunctionHandler method initRemoteRegistration.

/**
 * Instantiates remote registration. First gets remote function registry with version.
 * Version is used to ensure that we update the same registry we validated against.
 * Then validates against list of remote jars.
 * If validation is successful, first copies jars to registry area and starts updating remote function registry.
 * If during update {@link VersionMismatchException} was detected,
 * attempts to repeat remote registration process till retry attempts exceeds the limit.
 * If retry attempts number hits 0, throws exception that failed to update remote function registry.
 * In case of any error, if jars have been already copied to registry area, they will be deleted.
 *
 * @param functions list of functions present in jar
 * @param jarManager helper class for copying jars to registry area
 * @param remoteRegistry remote function registry
 * @throws IOException in case of problems with copying jars to registry area
 */
private void initRemoteRegistration(List<String> functions, JarManager jarManager, RemoteFunctionRegistry remoteRegistry) throws IOException {
    int retryAttempts = remoteRegistry.getRetryAttempts();
    boolean copyJars = true;
    try {
        while (retryAttempts >= 0) {
            DataChangeVersion version = new DataChangeVersion();
            List<Jar> remoteJars = remoteRegistry.getRegistry(version).getJarList();
            validateAgainstRemoteRegistry(remoteJars, jarManager.getBinaryName(), functions);
            if (copyJars) {
                jarManager.copyToRegistryArea();
                copyJars = false;
            }
            List<Jar> jars = Lists.newArrayList(remoteJars);
            jars.add(Jar.newBuilder().setName(jarManager.getBinaryName()).addAllFunctionSignature(functions).build());
            Registry updatedRegistry = Registry.newBuilder().addAllJar(jars).build();
            try {
                remoteRegistry.updateRegistry(updatedRegistry, version);
                return;
            } catch (VersionMismatchException ex) {
                logger.debug("Failed to update function registry during registration, version mismatch was detected.", ex);
                retryAttempts--;
            }
        }
        throw new DrillRuntimeException("Failed to update remote function registry. Exceeded retry attempts limit.");
    } catch (Exception e) {
        if (!copyJars) {
            jarManager.deleteQuietlyFromRegistryArea();
        }
        throw e;
    }
}

Example 9 with Registry

use of org.apache.drill.exec.proto.UserBitShared.Registry in project drill by axbaretto.

the class DropFunctionHandler method unregister.

/**
 * Gets remote function registry with version.
 * Version is used to ensure that we update the same registry we removed jars from.
 * Looks for a jar to be deleted, if founds one,
 * attempts to update remote registry with list of jars, that excludes jar to be deleted.
 * If during update {@link VersionMismatchException} was detected,
 * attempts to repeat unregistration process till retry attempts exceeds the limit.
 * If retry attempts number hits 0, throws exception that failed to update remote function registry.
 *
 * @param jarName jar name
 * @param remoteFunctionRegistry remote function registry
 * @return jar that was unregistered, null otherwise
 */
private Jar unregister(String jarName, RemoteFunctionRegistry remoteFunctionRegistry) {
    int retryAttempts = remoteFunctionRegistry.getRetryAttempts();
    while (retryAttempts >= 0) {
        DataChangeVersion version = new DataChangeVersion();
        Registry registry = remoteFunctionRegistry.getRegistry(version);
        Jar jarToBeDeleted = null;
        List<Jar> jars = Lists.newArrayList();
        for (Jar j : registry.getJarList()) {
            if (j.getName().equals(jarName)) {
                jarToBeDeleted = j;
            } else {
                jars.add(j);
            }
        }
        if (jarToBeDeleted == null) {
            return null;
        }
        Registry updatedRegistry = Registry.newBuilder().addAllJar(jars).build();
        try {
            remoteFunctionRegistry.updateRegistry(updatedRegistry, version);
            return jarToBeDeleted;
        } catch (VersionMismatchException ex) {
            logger.debug("Failed to update function registry during unregistration, version mismatch was detected.", ex);
            retryAttempts--;
        }
    }
    throw new DrillRuntimeException("Failed to update remote function registry. Exceeded retry attempts limit.");
}

Example 10 with Registry

use of org.apache.drill.exec.proto.UserBitShared.Registry in project drill by axbaretto.

the class TestDynamicUDFSupport method testConcurrentRemoteRegistryUpdateWithDuplicates.

@Test
public void testConcurrentRemoteRegistryUpdateWithDuplicates() throws Exception {
    RemoteFunctionRegistry remoteFunctionRegistry = spyRemoteFunctionRegistry();
    final CountDownLatch latch1 = new CountDownLatch(1);
    final CountDownLatch latch2 = new CountDownLatch(1);
    final CountDownLatch latch3 = new CountDownLatch(1);
    doAnswer(new Answer<Void>() {

        @Override
        public Void answer(InvocationOnMock invocation) throws Throwable {
            latch3.countDown();
            latch1.await();
            invocation.callRealMethod();
            latch2.countDown();
            return null;
        }
    }).doAnswer(new Answer<Void>() {

        @Override
        public Void answer(InvocationOnMock invocation) throws Throwable {
            latch1.countDown();
            latch2.await();
            invocation.callRealMethod();
            return null;
        }
    }).when(remoteFunctionRegistry).updateRegistry(any(Registry.class), any(DataChangeVersion.class));
    final String jarName1 = default_binary_name;
    final String jarName2 = "DrillUDF_Copy-1.0.jar";
    final String query = "create function using jar '%s'";
    copyDefaultJarsToStagingArea();
    copyJarsToStagingArea(jarName2, JarUtil.getSourceName(jarName2));
    Thread thread1 = new Thread(new TestBuilderRunner(testBuilder().sqlQuery(query, jarName1).unOrdered().baselineColumns("ok", "summary").baselineValues(true, String.format("The following UDFs in jar %s have been registered:\n" + "[custom_lower(VARCHAR-REQUIRED)]", jarName1))));
    Thread thread2 = new Thread(new TestBuilderRunner(testBuilder().sqlQuery(query, jarName2).unOrdered().baselineColumns("ok", "summary").baselineValues(false, String.format("Found duplicated function in %s: custom_lower(VARCHAR-REQUIRED)", jarName1))));
    thread1.start();
    latch3.await();
    thread2.start();
    thread1.join();
    thread2.join();
    DataChangeVersion version = new DataChangeVersion();
    Registry registry = remoteFunctionRegistry.getRegistry(version);
    assertEquals("Remote registry version should match", 1, version.getVersion());
    List<Jar> jarList = registry.getJarList();
    assertEquals("Only one jar should be registered", 1, jarList.size());
    assertEquals("Jar name should match", jarName1, jarList.get(0).getName());
    verify(remoteFunctionRegistry, times(2)).updateRegistry(any(Registry.class), any(DataChangeVersion.class));
}