org.apache.nifi.controller.ControllerService Java Examples

/**
 * Uses the {@link ExtensionManager} to discover any {@link ControllerService} implementations that implement a specific
 * ControllerService API.
 *
 * @param parent the controller service API
 * @return a list of controller services that implement the controller service API
 */
private List<Class<? extends ControllerService>> lookupControllerServiceImpls(
        final Class<? extends ControllerService> parent) {

    final List<Class<? extends ControllerService>> implementations = new ArrayList<>();

    // first get all ControllerService implementations
    final Set<Class> controllerServices = extensionManager.getExtensions(ControllerService.class);

    // then iterate over all controller services looking for any that is a child of the parent
    // ControllerService API that was passed in as a parameter
    for (final Class<? extends ControllerService> controllerServiceClass : controllerServices) {
        if (parent.isAssignableFrom(controllerServiceClass)) {
            implementations.add(controllerServiceClass);
        }
    }

    return implementations;
}
 

@Override
public void disableControllerService(final ControllerService service) {
    final ControllerServiceConfiguration configuration = context.getConfiguration(service.getIdentifier());
    if (configuration == null) {
        throw new IllegalArgumentException("Controller Service " + service + " is not known");
    }

    if (!configuration.isEnabled()) {
        throw new IllegalStateException("Controller service " + service + " cannot be disabled because it is not enabled");
    }

    try {
        ReflectionUtils.invokeMethodsWithAnnotation(OnDisabled.class, service);
    } catch (final Exception e) {
        e.printStackTrace();
        Assert.fail("Failed to disable Controller Service " + service + " due to " + e);
    }

    configuration.setEnabled(false);
}
 

private static boolean checkControllerServiceReferenceEligibility(final ConfigurableComponent component, final ClassLoader classLoader) {
    // if the extension does not require instance classloading, its eligible
    final boolean requiresInstanceClassLoading = component.getClass().isAnnotationPresent(RequiresInstanceClassLoading.class);

    final Set<Class> cobundledApis = new HashSet<>();
    try (final NarCloseable closeable = NarCloseable.withComponentNarLoader(component.getClass().getClassLoader())) {
        final List<PropertyDescriptor> descriptors = component.getPropertyDescriptors();
        if (descriptors != null && !descriptors.isEmpty()) {
            for (final PropertyDescriptor descriptor : descriptors) {
                final Class<? extends ControllerService> serviceApi = descriptor.getControllerServiceDefinition();
                if (serviceApi != null && classLoader.equals(serviceApi.getClassLoader())) {
                    cobundledApis.add(serviceApi);
                }
            }
        }
    }

    if (!cobundledApis.isEmpty()) {
        logger.warn(String.format(
                "Component %s is bundled with its referenced Controller Service APIs %s. The service APIs should not be bundled with component implementations that reference it.",
                component.getClass().getName(), StringUtils.join(cobundledApis.stream().map(cls -> cls.getName()).collect(Collectors.toSet()), ", ")));
    }

    // the component is eligible when it does not require instance classloading or when the supporting APIs are bundled in a parent NAR
    return requiresInstanceClassLoading == false || cobundledApis.isEmpty();
}
 

@Test
public void testLookupShouldThrowExceptionWhenAttributeMapHasNoLookupAttribute() {
    // GIVEN
    String mappedCreatedServiceID = "mappedCreatedServiceID";
    ControllerService mappedCreatedService = mock(SERVICE_TYPE);
    MockControllerServiceInitializationContext serviceLookup = new MockControllerServiceInitializationContext(mappedCreatedService, mappedCreatedServiceID);

    // WHEN
    testSubject.onEnabled(new MockConfigurationContext(properties, serviceLookup));
    try {
        testSubject.lookupService(new HashMap<>());
        fail();
    } catch (ProcessException e) {
        assertEquals("Attributes must contain an attribute name '" + LOOKUP_ATTRIBUTE + "'", e.getMessage());
    }
}
 

/**
 * Set up the test environment and mock behaviour. This includes registering {@link #reporterServiceStub} in the
 * different contexts, overriding {@link MetricsReportingTask#currentStatusReference} and instantiating the test
 * subject.
 */
@Before
public void setUp() throws Exception {
    Map<String, ControllerService> services = new HashMap<>();
    services.put(REPORTER_SERVICE_IDENTIFIER, reporterServiceStub);
    testedReportingTask = new MetricsReportingTask();
    reportingContextStub = new MockReportingContext(
            services, new MockStateManager(testedReportingTask), new MockVariableRegistry());

    rootGroupStatus = new ProcessGroupStatus();
    innerGroupStatus = new ProcessGroupStatus();
    when(reporterServiceStub.createReporter(any())).thenReturn(reporterMock);
    reportingContextStub.setProperty(MetricsReportingTask.REPORTER_SERVICE.getName(), REPORTER_SERVICE_IDENTIFIER);
    reportingContextStub.addControllerService(reporterServiceStub, REPORTER_SERVICE_IDENTIFIER);

    configurationContextStub = new MockConfigurationContext(reportingContextStub.getProperties(),
            reportingContextStub.getControllerServiceLookup());
    reportingInitContextStub = new MockReportingInitializationContext(
            TEST_INIT_CONTEXT_ID,
            TEST_INIT_CONTEXT_NAME,
            new MockComponentLog(TEST_TASK_ID, testedReportingTask));
}
 

@Override
public ValidationResult setProperty(final ControllerService service, final PropertyDescriptor property, final String value) {
    final MockStateManager serviceStateManager = controllerServiceStateManagers.get(service.getIdentifier());
    if (serviceStateManager == null) {
        throw new IllegalStateException("Controller service " + service + " has not been added to this TestRunner via the #addControllerService method");
    }

    final ControllerServiceConfiguration configuration = getConfigToUpdate(service);
    final Map<PropertyDescriptor, String> curProps = configuration.getProperties();
    final Map<PropertyDescriptor, String> updatedProps = new HashMap<>(curProps);

    final ValidationContext validationContext = new MockValidationContext(context, serviceStateManager, variableRegistry).getControllerServiceValidationContext(service);
    final ValidationResult validationResult = property.validate(value, validationContext);

    final String oldValue = updatedProps.get(property);
    updatedProps.put(property, value);
    configuration.setProperties(updatedProps);

    if ((value == null && oldValue != null) || (value != null && !value.equals(oldValue))) {
        service.onPropertyModified(property, oldValue, value);
    }

    return validationResult;
}
 

/**
 * Uses the {@link ExtensionManager} to discover any {@link ControllerService} implementations that implement a specific
 * ControllerService API.
 *
 * @param parent the controller service API
 * @return a list of controller services that implement the controller service API
 */
private List<Class<? extends ControllerService>> lookupControllerServiceImpls(
        final Class<? extends ControllerService> parent) {

    final List<Class<? extends ControllerService>> implementations = new ArrayList<>();

    // first get all ControllerService implementations
    final Set<Class> controllerServices = ExtensionManager.getExtensions(ControllerService.class);

    // then iterate over all controller services looking for any that is a child of the parent
    // ControllerService API that was passed in as a parameter
    for (final Class<? extends ControllerService> controllerServiceClass : controllerServices) {
        if (parent.isAssignableFrom(controllerServiceClass)) {
            implementations.add(controllerServiceClass);
        }
    }

    return implementations;
}
 

/**
 * Gets the ControllerService types that this controller supports.
 *
 * @param serviceType type
 * @return the ControllerService types that this controller supports
 */
public Set<DocumentedTypeDTO> getControllerServiceTypes(final String serviceType) {
    final Set<Class> serviceImplementations = ExtensionManager.getExtensions(ControllerService.class);

    // identify the controller services that implement the specified serviceType if applicable
    final Set<Class> matchingServiceImplementions;
    if (serviceType != null) {
        matchingServiceImplementions = new HashSet<>();

        // check each type and remove those that aren't in the specified ancestry
        for (final Class type : serviceImplementations) {
            if (implementsServiceType(serviceType, type)) {
                matchingServiceImplementions.add(type);
            }
        }
    } else {
        matchingServiceImplementions = serviceImplementations;
    }

    return dtoFactory.fromDocumentedTypes(matchingServiceImplementions);
}
 

private void addControllerServiceFingerprint(final StringBuilder builder, final ControllerServiceDTO dto, final FlowController controller) {
    builder.append(dto.getId());
    builder.append(dto.getType());
    builder.append(dto.getName());
    builder.append(dto.getComments());
    builder.append(dto.getAnnotationData());
    builder.append(dto.getState());

    // create an instance of the ControllerService so that we know the default property values
    ControllerService controllerService = null;
    try {
        if (controller != null) {
            controllerService = controller.createControllerService(dto.getType(), UUID.randomUUID().toString(), false).getControllerServiceImplementation();
        }
    } catch (Exception e) {
        logger.warn("Unable to create ControllerService of type {} due to {}; its default properties will be fingerprinted instead of being ignored.", dto.getType(), e.toString());
        if (logger.isDebugEnabled()) {
            logger.warn("", e);
        }
    }

    addPropertiesFingerprint(builder, controllerService, dto.getProperties());
}
 

/**
 * Returns a ConfigurableComponentInitializer for the type of component.
 * Currently Processor, ControllerService and ReportingTask are supported.
 *
 * @param componentClass the class that requires a ConfigurableComponentInitializer
 * @return a ConfigurableComponentInitializer capable of initializing that specific type of class
 */
public static ConfigurableComponentInitializer createComponentInitializer(
        final ExtensionManager extensionManager, final Class<? extends ConfigurableComponent> componentClass) {
    if (Processor.class.isAssignableFrom(componentClass)) {
        return new ProcessorInitializer(extensionManager);
    } else if (ControllerService.class.isAssignableFrom(componentClass)) {
        return new ControllerServiceInitializer(extensionManager);
    } else if (ReportingTask.class.isAssignableFrom(componentClass)) {
        return new ReportingTaskingInitializer(extensionManager);
    }

    return null;
}
 

public MockConfigurationContext(final ControllerService service,
        final Map<PropertyDescriptor, String> properties,
        final ControllerServiceLookup serviceLookup,
        final VariableRegistry variableRegistry) {
    this.service = service;
    this.properties = properties;
    this.serviceLookup = serviceLookup;
    this.variableRegistry = variableRegistry;
}
 

public StandardControllerServiceNode(final LoggableComponent<ControllerService> implementation, final LoggableComponent<ControllerService> proxiedControllerService,
                                     final ControllerServiceInvocationHandler invocationHandler, final String id, final ValidationContextFactory validationContextFactory,
                                     final ControllerServiceProvider serviceProvider, final String componentType, final String componentCanonicalClass,
                                     final ComponentVariableRegistry variableRegistry, final ReloadComponent reloadComponent, final ExtensionManager extensionManager,
                                     final ValidationTrigger validationTrigger, final boolean isExtensionMissing) {

    super(id, validationContextFactory, serviceProvider, componentType, componentCanonicalClass, variableRegistry, reloadComponent, extensionManager, validationTrigger, isExtensionMissing);
    this.serviceProvider = serviceProvider;
    this.active = new AtomicBoolean();
    setControllerServiceAndProxy(implementation, proxiedControllerService, invocationHandler);
    stateTransition = new ServiceStateTransition();
}
 

/**
 * Generates documentation into the work/docs dir specified by
 * NiFiProperties.
 *
 * @param properties to lookup nifi properties
 * @param extensionMapping extension mapping
 */
public static void generate(final NiFiProperties properties, final ExtensionManager extensionManager, final ExtensionMapping extensionMapping) {
    final File explodedNiFiDocsDir = properties.getComponentDocumentationWorkingDirectory();

    logger.debug("Generating documentation for: " + extensionMapping.size() + " components in: " + explodedNiFiDocsDir);

    documentConfigurableComponent(extensionManager.getExtensions(Processor.class), explodedNiFiDocsDir, extensionManager);
    documentConfigurableComponent(extensionManager.getExtensions(ControllerService.class), explodedNiFiDocsDir, extensionManager);
    documentConfigurableComponent(extensionManager.getExtensions(ReportingTask.class), explodedNiFiDocsDir, extensionManager);
}
 

@Override
public void teardown(ConfigurableComponent component) {
    try (NarCloseable narCloseable = NarCloseable.withComponentNarLoader(component.getClass(), component.getIdentifier())) {
        ControllerService controllerService = (ControllerService) component;

        final ComponentLog logger = new MockComponentLogger();
        final MockConfigurationContext context = new MockConfigurationContext();
        ReflectionUtils.quietlyInvokeMethodsWithAnnotation(OnShutdown.class, controllerService, logger, context);
    } finally {
        ExtensionManager.removeInstanceClassLoaderIfExists(component.getIdentifier());
    }
}
 

@Test
public void testLookupShouldReturnQueriedService() {
    // GIVEN
    String mappedCreatedServiceID1 = "mappedCreatedServiceID1";
    String mappedCreatedServiceID2 = "mappedCreatedServiceID2";

    ControllerService mappedCreatedService1 = mock(SERVICE_TYPE);
    ControllerService mappedCreatedService2 = mock(SERVICE_TYPE);

    MockControllerServiceInitializationContext serviceLookup = new MockControllerServiceInitializationContext(mappedCreatedService1, mappedCreatedServiceID1);
    serviceLookup.addControllerService(mappedCreatedService2, mappedCreatedServiceID2);

    String dynamicProperty1 = "property1";
    String dynamicProperty2 = "property2";

    mapService(dynamicProperty1, mappedCreatedServiceID1);
    mapService(dynamicProperty2, mappedCreatedServiceID2);

    String lookupServiceKey = dynamicProperty2;
    ControllerService expected = mappedCreatedService2;

    // WHEN
    testSubject.onEnabled(new MockConfigurationContext(properties, serviceLookup));
    ControllerService actual = testSubject.lookupService(createAttributes(lookupServiceKey));

    // THEN
    assertEquals(expected, actual);
}
 

@Override
public ControllerService asControllerService() {
    if (rawValue == null || rawValue.equals("") || serviceLookup == null) {
        return null;
    }

    return serviceLookup.getControllerService(rawValue);
}
 

@Override
public void initialize(final ConfigurableComponent component) {
    try {
        if (component instanceof Processor) {
            initialize((Processor) component);
        } else if (component instanceof ControllerService) {
            initialize((ControllerService) component);
        } else if (component instanceof ReportingTask) {
            initialize((ReportingTask) component);
        }
    } catch (final InitializationException ie) {
        throw new RuntimeException("Failed to initialize " + component, ie);
    }
}
 

/**
 * Returns whether the specified type implements the specified serviceType.
 *
 * @param serviceType type
 * @param type type
 * @return whether the specified type implements the specified serviceType
 */
private boolean implementsServiceType(final String serviceType, final Class type) {
    final List<Class<?>> interfaces = ClassUtils.getAllInterfaces(type);
    for (final Class i : interfaces) {
        if (ControllerService.class.isAssignableFrom(i) && i.getName().equals(serviceType)) {
            return true;
        }
    }

    return false;
}
 

@Override
public void enableControllerService(final ControllerService service) {
    final ControllerServiceConfiguration configuration = context.getConfiguration(service.getIdentifier());
    if (configuration == null) {
        throw new IllegalArgumentException("Controller Service " + service + " is not known");
    }

    if (configuration.isEnabled()) {
        throw new IllegalStateException("Cannot enable Controller Service " + service + " because it is not disabled");
    }

    // ensure controller service is valid before enabling
    final ValidationContext validationContext = new MockValidationContext(context).getControllerServiceValidationContext(service);
    final Collection<ValidationResult> results = context.getControllerService(service.getIdentifier()).validate(validationContext);

    for (final ValidationResult result : results) {
        if (!result.isValid()) {
            throw new IllegalStateException("Cannot enable Controller Service " + service + " because it is in an invalid state: " + result.toString());
        }
    }

    try {
        final ConfigurationContext configContext = new MockConfigurationContext(service, configuration.getProperties(), context,variableRegistry);
        ReflectionUtils.invokeMethodsWithAnnotation(OnEnabled.class, service, configContext);
    } catch (final InvocationTargetException ite) {
        ite.getCause().printStackTrace();
        Assert.fail("Failed to enable Controller Service " + service + " due to " + ite.getCause());
    } catch (final Exception e) {
        e.printStackTrace();
        Assert.fail("Failed to enable Controller Service " + service + " due to " + e);
    }

    configuration.setEnabled(true);
}
 

public void removeControllerService(final ControllerService service) {
    final ControllerService canonical = getControllerService(service.getIdentifier());
    if (canonical == null || canonical != service) {
        throw new IllegalArgumentException("Controller Service " + service + " is not known");
    }

    controllerServiceMap.remove(service.getIdentifier());
}
 

@Test
public void testControllerServiceUpdateShouldCallOnSetProperty() {
    // Arrange
    final ControllerService testService = new SimpleTestService();
    final AddAttributeProcessor proc = new AddAttributeProcessor();
    final TestRunner runner = TestRunners.newTestRunner(proc);
    final String serviceIdentifier = "test";
    final String pdName = "name";
    final String pdValue = "exampleName";
    try {
        runner.addControllerService(serviceIdentifier, testService);
    } catch (InitializationException e) {
        fail(e.getMessage());
    }

    assertFalse("onPropertyModified has been called", ((SimpleTestService) testService).isOpmCalled());

    // Act
    ValidationResult vr = runner.setProperty(testService, pdName, pdValue);

    // Assert
    assertTrue(vr.isValid());

    ControllerServiceConfiguration csConf = ((MockProcessContext) runner.getProcessContext()).getConfiguration(serviceIdentifier);
    PropertyDescriptor propertyDescriptor = testService.getPropertyDescriptor(pdName);
    String retrievedPDValue = csConf.getProperties().get(propertyDescriptor);

    assertEquals(pdValue, retrievedPDValue);
    assertTrue("onPropertyModified has not been called", ((SimpleTestService) testService).isOpmCalled());
}
 

public MockConfigurationContext(final ControllerService service,
        final Map<PropertyDescriptor, String> properties,
        final ControllerServiceLookup serviceLookup,
        final VariableRegistry variableRegistry) {
    this.service = service;
    this.properties = properties;
    this.serviceLookup = serviceLookup;
    this.variableRegistry = variableRegistry;
}
 

@Test
public void validateServiceIsLocatableViaServiceLoader() {
    ServiceLoader<ControllerService> loader = ServiceLoader.<ControllerService>load(ControllerService.class);
    Iterator<ControllerService> iter = loader.iterator();
    boolean present = false;
    while (iter.hasNext()) {
        ControllerService cs = iter.next();
        if (cs instanceof JMSConnectionFactoryProviderDefinition) {
            present = true;
            break;
        }
    }
    assertTrue(present);
}
 

@Test
public void validateServiceIsLocatableViaServiceLoader() {
    ServiceLoader<ControllerService> loader = ServiceLoader.<ControllerService> load(ControllerService.class);
    Iterator<ControllerService> iter = loader.iterator();
    boolean present = false;
    while (iter.hasNext()) {
        ControllerService cs = iter.next();
        assertTrue(cs instanceof JMSConnectionFactoryProviderDefinition);
        present = true;
    }
    assertTrue(present);
}
 

@Test
public void testServiceIsNotWebSocketService() throws Exception {
    final TestRunner runner = TestRunners.newTestRunner(PutWebSocket.class);
    final ControllerService service = spy(ControllerService.class);

    final WebSocketSession webSocketSession = getWebSocketSession();

    final String serviceId = "ws-service";
    final String endpointId = "client-1";
    final String textMessageFromServer = "message from server.";
    when(service.getIdentifier()).thenReturn(serviceId);
    runner.addControllerService(serviceId, service);

    runner.enableControllerService(service);

    final Map<String, String> attributes = new HashMap<>();
    attributes.put(ATTR_WS_CS_ID, serviceId);
    attributes.put(ATTR_WS_ENDPOINT_ID, endpointId);
    attributes.put(ATTR_WS_SESSION_ID, webSocketSession.getSessionId());
    runner.enqueue(textMessageFromServer, attributes);

    runner.run();

    final List<MockFlowFile> succeededFlowFiles = runner.getFlowFilesForRelationship(PutWebSocket.REL_SUCCESS);
    assertEquals(0, succeededFlowFiles.size());

    final List<MockFlowFile> failedFlowFiles = runner.getFlowFilesForRelationship(PutWebSocket.REL_FAILURE);
    assertEquals(1, failedFlowFiles.size());
    final MockFlowFile failedFlowFile = failedFlowFiles.iterator().next();
    assertNotNull(failedFlowFile.getAttribute(ATTR_WS_FAILURE_DETAIL));

    final List<ProvenanceEventRecord> provenanceEvents = runner.getProvenanceEvents();
    assertEquals(0, provenanceEvents.size());

}
 

@Override
public boolean isValidationRequired(final ControllerService service) {
    // No need to validate services that are already enabled.
    final ControllerServiceState serviceState = controllerServiceProvider.getControllerServiceNode(service.getIdentifier()).getState();
    if (serviceState == ControllerServiceState.ENABLED || serviceState == ControllerServiceState.ENABLING) {
        return false;
    }

    return true;
}
 

@Override
public void removeControllerService(final ControllerService service) {
    disableControllerService(service);

    try {
        ReflectionUtils.invokeMethodsWithAnnotation(OnRemoved.class, service);
    } catch (final Exception e) {
        e.printStackTrace();
        Assert.fail("Failed to remove Controller Service " + service + " due to " + e);
    }

    context.removeControllerService(service);
}
 

@Test
public void testServiceIsNotFound() throws Exception {
    final TestRunner runner = TestRunners.newTestRunner(PutWebSocket.class);
    final ControllerService service = spy(ControllerService.class);

    final WebSocketSession webSocketSession = getWebSocketSession();

    final String serviceId = "ws-service";
    final String endpointId = "client-1";
    final String textMessageFromServer = "message from server.";
    when(service.getIdentifier()).thenReturn(serviceId);
    runner.addControllerService(serviceId, service);

    runner.enableControllerService(service);

    final Map<String, String> attributes = new HashMap<>();
    attributes.put(ATTR_WS_CS_ID, "different-service-id");
    attributes.put(ATTR_WS_ENDPOINT_ID, endpointId);
    attributes.put(ATTR_WS_SESSION_ID, webSocketSession.getSessionId());
    runner.enqueue(textMessageFromServer, attributes);

    runner.run();

    final List<MockFlowFile> succeededFlowFiles = runner.getFlowFilesForRelationship(PutWebSocket.REL_SUCCESS);
    assertEquals(0, succeededFlowFiles.size());

    final List<MockFlowFile> failedFlowFiles = runner.getFlowFilesForRelationship(PutWebSocket.REL_FAILURE);
    assertEquals(1, failedFlowFiles.size());
    final MockFlowFile failedFlowFile = failedFlowFiles.iterator().next();
    assertNotNull(failedFlowFile.getAttribute(ATTR_WS_FAILURE_DETAIL));

    final List<ProvenanceEventRecord> provenanceEvents = runner.getProvenanceEvents();
    assertEquals(0, provenanceEvents.size());

}
 

@Override
public ValidationResult setProperty(final ControllerService service, final PropertyDescriptor property, final AllowableValue value) {
    return setProperty(service, property, value.getValue());
}
 

@Override
public ControllerService getControllerService(final String identifier) {
    final StatelessControllerServiceConfiguration status = controllerServiceMap.get(identifier);
    return (status == null) ? null : status.getService();
}