Examples with UserCodeClassLoader org.apache.flink.util.UserCodeClassLoader used on opensource projects

Example 11 with UserCodeClassLoader

use of org.apache.flink.util.UserCodeClassLoader in project flink by apache.

the class BlobLibraryCacheManagerTest method releaseUserCodeClassLoader_willRunReleaseHooks.

@Test
public void releaseUserCodeClassLoader_willRunReleaseHooks() throws IOException, InterruptedException {
    final BlobLibraryCacheManager libraryCacheManager = new TestingBlobLibraryCacheManagerBuilder().build();
    final LibraryCacheManager.ClassLoaderLease classLoaderLease = libraryCacheManager.registerClassLoaderLease(new JobID());
    final UserCodeClassLoader userCodeClassLoader = classLoaderLease.getOrResolveClassLoader(Collections.emptyList(), Collections.emptyList());
    final OneShotLatch releaseHookLatch = new OneShotLatch();
    userCodeClassLoader.registerReleaseHookIfAbsent("test", releaseHookLatch::trigger);
    // this should trigger the release of the class loader
    classLoaderLease.release();
    releaseHookLatch.await();
}

Example 12 with UserCodeClassLoader

use of org.apache.flink.util.UserCodeClassLoader in project flink by apache.

the class BlobLibraryCacheManagerTest method testLibraryCacheManagerDifferentJobsCleanup.

/**
 * Tests that the {@link BlobLibraryCacheManager} cleans up after the class loader leases for
 * different jobs are closed.
 */
@Test
public void testLibraryCacheManagerDifferentJobsCleanup() throws Exception {
    JobID jobId1 = new JobID();
    JobID jobId2 = new JobID();
    List<PermanentBlobKey> keys1 = new ArrayList<>();
    List<PermanentBlobKey> keys2 = new ArrayList<>();
    BlobServer server = null;
    PermanentBlobCache cache = null;
    BlobLibraryCacheManager libCache = null;
    final byte[] buf = new byte[128];
    try {
        Configuration config = new Configuration();
        config.setLong(BlobServerOptions.CLEANUP_INTERVAL, 1L);
        server = new BlobServer(config, temporaryFolder.newFolder(), new VoidBlobStore());
        server.start();
        InetSocketAddress serverAddress = new InetSocketAddress("localhost", server.getPort());
        cache = new PermanentBlobCache(config, temporaryFolder.newFolder(), new VoidBlobStore(), serverAddress);
        keys1.add(server.putPermanent(jobId1, buf));
        buf[0] += 1;
        keys1.add(server.putPermanent(jobId1, buf));
        keys2.add(server.putPermanent(jobId2, buf));
        libCache = createBlobLibraryCacheManager(cache);
        cache.registerJob(jobId1);
        cache.registerJob(jobId2);
        assertEquals(0, libCache.getNumberOfManagedJobs());
        assertEquals(0, libCache.getNumberOfReferenceHolders(jobId1));
        checkFileCountForJob(2, jobId1, server);
        checkFileCountForJob(0, jobId1, cache);
        checkFileCountForJob(1, jobId2, server);
        checkFileCountForJob(0, jobId2, cache);
        final LibraryCacheManager.ClassLoaderLease classLoaderLeaseJob1 = libCache.registerClassLoaderLease(jobId1);
        final UserCodeClassLoader classLoader1 = classLoaderLeaseJob1.getOrResolveClassLoader(keys1, Collections.emptyList());
        assertEquals(1, libCache.getNumberOfManagedJobs());
        assertEquals(1, libCache.getNumberOfReferenceHolders(jobId1));
        assertEquals(0, libCache.getNumberOfReferenceHolders(jobId2));
        assertEquals(2, checkFilesExist(jobId1, keys1, cache, true));
        checkFileCountForJob(2, jobId1, server);
        checkFileCountForJob(2, jobId1, cache);
        assertEquals(0, checkFilesExist(jobId2, keys2, cache, false));
        checkFileCountForJob(1, jobId2, server);
        checkFileCountForJob(0, jobId2, cache);
        final LibraryCacheManager.ClassLoaderLease classLoaderLeaseJob2 = libCache.registerClassLoaderLease(jobId2);
        final UserCodeClassLoader classLoader2 = classLoaderLeaseJob2.getOrResolveClassLoader(keys2, Collections.emptyList());
        assertThat(classLoader1, not(sameInstance(classLoader2)));
        try {
            classLoaderLeaseJob2.getOrResolveClassLoader(keys1, Collections.<URL>emptyList());
            fail("Should fail with an IllegalStateException");
        } catch (IllegalStateException e) {
        // that's what we want
        }
        try {
            classLoaderLeaseJob2.getOrResolveClassLoader(keys2, Collections.singletonList(new URL("file:///tmp/does-not-exist")));
            fail("Should fail with an IllegalStateException");
        } catch (IllegalStateException e) {
        // that's what we want
        }
        assertEquals(2, libCache.getNumberOfManagedJobs());
        assertEquals(1, libCache.getNumberOfReferenceHolders(jobId1));
        assertEquals(1, libCache.getNumberOfReferenceHolders(jobId2));
        assertEquals(2, checkFilesExist(jobId1, keys1, cache, true));
        checkFileCountForJob(2, jobId1, server);
        checkFileCountForJob(2, jobId1, cache);
        assertEquals(1, checkFilesExist(jobId2, keys2, cache, true));
        checkFileCountForJob(1, jobId2, server);
        checkFileCountForJob(1, jobId2, cache);
        classLoaderLeaseJob1.release();
        assertEquals(1, libCache.getNumberOfManagedJobs());
        assertEquals(0, libCache.getNumberOfReferenceHolders(jobId1));
        assertEquals(1, libCache.getNumberOfReferenceHolders(jobId2));
        assertEquals(2, checkFilesExist(jobId1, keys1, cache, true));
        checkFileCountForJob(2, jobId1, server);
        checkFileCountForJob(2, jobId1, cache);
        assertEquals(1, checkFilesExist(jobId2, keys2, cache, true));
        checkFileCountForJob(1, jobId2, server);
        checkFileCountForJob(1, jobId2, cache);
        classLoaderLeaseJob2.release();
        assertEquals(0, libCache.getNumberOfManagedJobs());
        assertEquals(0, libCache.getNumberOfReferenceHolders(jobId1));
        assertEquals(0, libCache.getNumberOfReferenceHolders(jobId2));
        assertEquals(2, checkFilesExist(jobId1, keys1, cache, true));
        checkFileCountForJob(2, jobId1, server);
        checkFileCountForJob(2, jobId1, cache);
        assertEquals(1, checkFilesExist(jobId2, keys2, cache, true));
        checkFileCountForJob(1, jobId2, server);
        checkFileCountForJob(1, jobId2, cache);
    // only PermanentBlobCache#releaseJob() calls clean up files (tested in
    // BlobCacheCleanupTest etc.
    } finally {
        if (libCache != null) {
            libCache.shutdown();
        }
        // should have been closed by the libraryCacheManager, but just in case
        if (cache != null) {
            cache.close();
        }
        if (server != null) {
            server.close();
        }
    }
}

Example 13 with UserCodeClassLoader

use of org.apache.flink.util.UserCodeClassLoader in project flink by apache.

the class DataSourceTask method initInputFormat.

/**
 * Initializes the InputFormat implementation and configuration.
 *
 * @throws RuntimeException Throws if instance of InputFormat implementation can not be
 *     obtained.
 */
private void initInputFormat() {
    ClassLoader userCodeClassLoader = getUserCodeClassLoader();
    // obtain task configuration (including stub parameters)
    Configuration taskConf = getTaskConfiguration();
    this.config = new TaskConfig(taskConf);
    final Pair<OperatorID, InputFormat<OT, InputSplit>> operatorIdAndInputFormat;
    InputOutputFormatContainer formatContainer = new InputOutputFormatContainer(config, userCodeClassLoader);
    try {
        operatorIdAndInputFormat = formatContainer.getUniqueInputFormat();
        this.format = operatorIdAndInputFormat.getValue();
        // check if the class is a subclass, if the check is required
        if (!InputFormat.class.isAssignableFrom(this.format.getClass())) {
            throw new RuntimeException("The class '" + this.format.getClass().getName() + "' is not a subclass of '" + InputFormat.class.getName() + "' as is required.");
        }
    } catch (ClassCastException ccex) {
        throw new RuntimeException("The stub class is not a proper subclass of " + InputFormat.class.getName(), ccex);
    }
    Thread thread = Thread.currentThread();
    ClassLoader original = thread.getContextClassLoader();
    // user code
    try {
        thread.setContextClassLoader(userCodeClassLoader);
        this.format.configure(formatContainer.getParameters(operatorIdAndInputFormat.getKey()));
    } catch (Throwable t) {
        throw new RuntimeException("The user defined 'configure()' method caused an error: " + t.getMessage(), t);
    } finally {
        thread.setContextClassLoader(original);
    }
    // get the factory for the type serializer
    this.serializerFactory = this.config.getOutputSerializer(userCodeClassLoader);
}

Example 14 with UserCodeClassLoader

use of org.apache.flink.util.UserCodeClassLoader in project flink by apache.

the class BatchTask method initStub.

protected S initStub(Class<? super S> stubSuperClass) throws Exception {
    try {
        ClassLoader userCodeClassLoader = getUserCodeClassLoader();
        S stub = config.<S>getStubWrapper(userCodeClassLoader).getUserCodeObject(stubSuperClass, userCodeClassLoader);
        // check if the class is a subclass, if the check is required
        if (stubSuperClass != null && !stubSuperClass.isAssignableFrom(stub.getClass())) {
            throw new RuntimeException("The class '" + stub.getClass().getName() + "' is not a subclass of '" + stubSuperClass.getName() + "' as is required.");
        }
        FunctionUtils.setFunctionRuntimeContext(stub, this.runtimeUdfContext);
        return stub;
    } catch (ClassCastException ccex) {
        throw new Exception("The stub class is not a proper subclass of " + stubSuperClass.getName(), ccex);
    }
}