Examples with Creator org.apache.beam.runners.fnexecution.control.ReferenceCountingExecutableStageContextFactory.Creator used on opensource projects

Example 1 with Creator

use of org.apache.beam.runners.fnexecution.control.ReferenceCountingExecutableStageContextFactory.Creator in project beam by apache.

the class ReferenceCountingExecutableStageContextFactoryTest method testCreateReuseReleaseCreate.

@Test
public void testCreateReuseReleaseCreate() throws Exception {
    Creator creator = mock(Creator.class);
    ExecutableStageContext c1 = mock(ExecutableStageContext.class);
    ExecutableStageContext c2 = mock(ExecutableStageContext.class);
    ExecutableStageContext c3 = mock(ExecutableStageContext.class);
    ExecutableStageContext c4 = mock(ExecutableStageContext.class);
    when(creator.apply(any(JobInfo.class))).thenReturn(c1).thenReturn(c2).thenReturn(c3).thenReturn(c4);
    ReferenceCountingExecutableStageContextFactory factory = ReferenceCountingExecutableStageContextFactory.create(creator, (x) -> true);
    JobInfo jobA = mock(JobInfo.class);
    when(jobA.jobId()).thenReturn("jobA");
    JobInfo jobB = mock(JobInfo.class);
    when(jobB.jobId()).thenReturn("jobB");
    // 1 open jobA
    ExecutableStageContext ac1A = factory.get(jobA);
    // 1 open jobB
    ExecutableStageContext ac2B = factory.get(jobB);
    Assert.assertSame("Context should be cached and reused.", ac1A, // 2 open jobA
    factory.get(jobA));
    Assert.assertSame("Context should be cached and reused.", ac2B, // 2 open jobB
    factory.get(jobB));
    // 1 open jobA
    factory.release(ac1A);
    Assert.assertSame("Context should be cached and reused.", ac1A, // 2 open jobA
    factory.get(jobA));
    // 1 open jobA
    factory.release(ac1A);
    // 0 open jobA
    factory.release(ac1A);
    // 1 open jobA
    ExecutableStageContext ac3A = factory.get(jobA);
    Assert.assertNotSame("We should get a new instance.", ac1A, ac3A);
    Assert.assertSame("Context should be cached and reused.", ac3A, // 2 open jobA
    factory.get(jobA));
    // 1 open jobA
    factory.release(ac3A);
    // 0 open jobA
    factory.release(ac3A);
    Assert.assertSame("Context should be cached and reused.", ac2B, // 3 open jobB
    factory.get(jobB));
    // 2 open jobB
    factory.release(ac2B);
    // 1 open jobB
    factory.release(ac2B);
    // 0 open jobB
    factory.release(ac2B);
    // 1 open jobB
    ExecutableStageContext ac4B = factory.get(jobB);
    Assert.assertNotSame("We should get a new instance.", ac2B, ac4B);
    // 0 open jobB
    factory.release(ac4B);
}

Example 2 with Creator

use of org.apache.beam.runners.fnexecution.control.ReferenceCountingExecutableStageContextFactory.Creator in project beam by apache.

the class ReferenceCountingExecutableStageContextFactoryTest method testCatchThrowablesAndLogThem.

@Test
public void testCatchThrowablesAndLogThem() throws Exception {
    PrintStream oldErr = System.err;
    oldErr.flush();
    ByteArrayOutputStream baos = new ByteArrayOutputStream();
    PrintStream newErr = new PrintStream(baos);
    try {
        System.setErr(newErr);
        Creator creator = mock(Creator.class);
        ExecutableStageContext c1 = mock(ExecutableStageContext.class);
        when(creator.apply(any(JobInfo.class))).thenReturn(c1);
        // throw an Throwable and ensure that it is caught and logged.
        doThrow(new NoClassDefFoundError()).when(c1).close();
        ReferenceCountingExecutableStageContextFactory factory = ReferenceCountingExecutableStageContextFactory.create(creator, (x) -> true);
        JobInfo jobA = mock(JobInfo.class);
        when(jobA.jobId()).thenReturn("jobA");
        ExecutableStageContext ac1A = factory.get(jobA);
        factory.release(ac1A);
        newErr.flush();
        String output = new String(baos.toByteArray(), Charsets.UTF_8);
        // Ensure that the error is logged
        assertTrue(output.contains("Unable to close ExecutableStageContext"));
    } finally {
        newErr.flush();
        System.setErr(oldErr);
    }
}