Skip to content

Commit

Permalink
NIFI-14156 Improved Bootstrap Process on Windows with Attach API
Browse files Browse the repository at this point in the history
- Added Virtual Machine Attach API implementation of to support Bootstrap commands on Windows
- Improved nifi.cmd start handling to launch in minimized window

This closes #9683.

Signed-off-by: Mark Bathori <[email protected]>
  • Loading branch information
exceptionfactory authored and mark-bathori committed Feb 3, 2025
1 parent a7bc0e5 commit 2b1055e
Show file tree
Hide file tree
Showing 6 changed files with 299 additions and 6 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@
import org.apache.nifi.bootstrap.command.process.ManagementServerAddressProvider;
import org.apache.nifi.bootstrap.command.process.ProcessHandleManagementServerAddressProvider;
import org.apache.nifi.bootstrap.command.process.ProcessHandleProvider;
import org.apache.nifi.bootstrap.command.process.VirtualMachineManagementServerAddressProvider;
import org.apache.nifi.bootstrap.configuration.ApplicationClassName;
import org.apache.nifi.bootstrap.configuration.ManagementServerPath;
import org.slf4j.Logger;
Expand Down Expand Up @@ -107,7 +108,7 @@ public void run() {
}

protected void run(final ProcessHandle applicationProcessHandle) {
final ManagementServerAddressProvider managementServerAddressProvider = new ProcessHandleManagementServerAddressProvider(applicationProcessHandle);
final ManagementServerAddressProvider managementServerAddressProvider = getManagementServerAddressProvider(applicationProcessHandle);
final Optional<String> managementServerAddress = managementServerAddressProvider.getAddress();

final long pid = applicationProcessHandle.pid();
Expand Down Expand Up @@ -176,4 +177,19 @@ protected HttpClient getHttpClient() {
builder.connectTimeout(CONNECT_TIMEOUT);
return builder.build();
}

private ManagementServerAddressProvider getManagementServerAddressProvider(final ProcessHandle applicationProcessHandle) {
final ManagementServerAddressProvider managementServerAddressProvider;

final ProcessHandle.Info applicationProcessHandleInfo = applicationProcessHandle.info();
final Optional<String[]> arguments = applicationProcessHandleInfo.arguments();
if (arguments.isPresent()) {
managementServerAddressProvider = new ProcessHandleManagementServerAddressProvider(applicationProcessHandle);
} else {
// Use Virtual Machine Attach API when ProcessHandle does not support arguments as described in JDK-8176725
managementServerAddressProvider = new VirtualMachineManagementServerAddressProvider(applicationProcessHandle);
}

return managementServerAddressProvider;
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@
import org.apache.nifi.bootstrap.command.io.StandardBootstrapArgumentParser;
import org.apache.nifi.bootstrap.command.process.StandardProcessHandleProvider;
import org.apache.nifi.bootstrap.command.process.ProcessHandleProvider;
import org.apache.nifi.bootstrap.command.process.VirtualMachineProcessHandleProvider;
import org.apache.nifi.bootstrap.configuration.ApplicationClassName;
import org.apache.nifi.bootstrap.configuration.ConfigurationProvider;
import org.apache.nifi.bootstrap.configuration.StandardConfigurationProvider;
Expand Down Expand Up @@ -99,7 +100,7 @@ public BootstrapCommand getBootstrapCommand(final String[] arguments) {

private BootstrapCommand getBootstrapCommand(final BootstrapArgument bootstrapArgument, final String[] arguments) {
final ConfigurationProvider configurationProvider = new StandardConfigurationProvider(System.getenv(), System.getProperties());
final ProcessHandleProvider processHandleProvider = new StandardProcessHandleProvider(configurationProvider);
final ProcessHandleProvider processHandleProvider = getProcessHandleProvider(configurationProvider);
final ResponseStreamHandler commandLoggerStreamHandler = new LoggerResponseStreamHandler(commandLogger);
final BootstrapCommand stopBootstrapCommand = new StopBootstrapCommand(processHandleProvider, configurationProvider);

Expand Down Expand Up @@ -238,4 +239,21 @@ private ResponseStreamHandler getStatusHistoryResponseStreamHandler(final String

return responseStreamHandler;
}

private ProcessHandleProvider getProcessHandleProvider(final ConfigurationProvider configurationProvider) {
final ProcessHandleProvider processHandleProvider;

final ProcessHandle currentProcessHandle = ProcessHandle.current();
final ProcessHandle.Info currentProcessHandleInfo = currentProcessHandle.info();
final Optional<String[]> currentProcessArguments = currentProcessHandleInfo.arguments();

if (currentProcessArguments.isPresent()) {
processHandleProvider = new StandardProcessHandleProvider(configurationProvider);
} else {
// Use Virtual Machine Attach API when ProcessHandle does not support arguments as described in JDK-8176725
processHandleProvider = new VirtualMachineProcessHandleProvider(configurationProvider);
}

return processHandleProvider;
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,66 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one or more
* contributor license agreements. See the NOTICE file distributed with
* this work for additional information regarding copyright ownership.
* The ASF licenses this file to You under the Apache License, Version 2.0
* (the "License"); you may not use this file except in compliance with
* the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.apache.nifi.bootstrap.command.process;

import com.sun.tools.attach.VirtualMachine;
import org.apache.nifi.bootstrap.configuration.SystemProperty;

import java.io.IOException;
import java.util.Objects;
import java.util.Optional;
import java.util.Properties;

/**
* Provider implementation resolves the Management Server Address from System Properties of Application Virtual Machine using the Attach API
*/
public class VirtualMachineManagementServerAddressProvider implements ManagementServerAddressProvider {
private final ProcessHandle processHandle;

public VirtualMachineManagementServerAddressProvider(final ProcessHandle processHandle) {
this.processHandle = Objects.requireNonNull(processHandle);
}

/**
* Get Management Server Address from System Properties of Application Virtual Machine using Attach API
*
* @return Management Server Address or null when not found
*/
@Override
public Optional<String> getAddress() {
String managementServerAddress = null;

final String applicationProcessId = Long.toString(processHandle.pid());

try {
final VirtualMachine virtualMachine = VirtualMachine.attach(applicationProcessId);
try {
managementServerAddress = getAddress(virtualMachine);
} finally {
virtualMachine.detach();
}
} catch (final Exception ignored) {

}

return Optional.ofNullable(managementServerAddress);
}

private String getAddress(final VirtualMachine virtualMachine) throws IOException {
final Properties systemProperties = virtualMachine.getSystemProperties();
return systemProperties.getProperty(SystemProperty.MANAGEMENT_SERVER_ADDRESS.getProperty());
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,127 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one or more
* contributor license agreements. See the NOTICE file distributed with
* this work for additional information regarding copyright ownership.
* The ASF licenses this file to You under the Apache License, Version 2.0
* (the "License"); you may not use this file except in compliance with
* the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.apache.nifi.bootstrap.command.process;

import com.sun.tools.attach.VirtualMachine;
import com.sun.tools.attach.VirtualMachineDescriptor;
import com.sun.tools.attach.spi.AttachProvider;
import org.apache.nifi.bootstrap.configuration.ConfigurationProvider;
import org.apache.nifi.bootstrap.configuration.SystemProperty;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.nio.file.Path;
import java.util.List;
import java.util.Objects;
import java.util.Optional;
import java.util.Properties;

/**
* Virtual Machine implementation of ProcessHandle Provider using the Attach API with System Properties
*/
public class VirtualMachineProcessHandleProvider implements ProcessHandleProvider {
private static final Logger logger = LoggerFactory.getLogger(VirtualMachineProcessHandleProvider.class);

private final ConfigurationProvider configurationProvider;

public VirtualMachineProcessHandleProvider(final ConfigurationProvider configurationProvider) {
this.configurationProvider = Objects.requireNonNull(configurationProvider);
}

/**
* Find Process Handle for Application based on matching argument for path to application properties
*
* @return Application Process Handle or empty when not found
*/
@Override
public Optional<ProcessHandle> findApplicationProcessHandle() {
final Path applicationProperties = configurationProvider.getApplicationProperties();
return findProcessHandle(SystemProperty.APPLICATION_PROPERTIES, applicationProperties);
}

/**
* Find Process Handle for Bootstrap based on matching argument for path to bootstrap configuration
*
* @return Bootstrap Process Handle or empty when not found
*/
@Override
public Optional<ProcessHandle> findBootstrapProcessHandle() {
final Path bootstrapConfiguration = configurationProvider.getBootstrapConfiguration();
return findProcessHandle(SystemProperty.BOOTSTRAP_CONFIGURATION, bootstrapConfiguration);
}

private Optional<ProcessHandle> findProcessHandle(final SystemProperty systemProperty, final Path configuration) {
final ProcessHandle currentProcessHandle = ProcessHandle.current();
final String currentProcessId = Long.toString(currentProcessHandle.pid());

Optional<ProcessHandle> processHandleFound = Optional.empty();

final List<VirtualMachineDescriptor> virtualMachineDescriptors = VirtualMachine.list();
for (final VirtualMachineDescriptor virtualMachineDescriptor : virtualMachineDescriptors) {
final String virtualMachineId = virtualMachineDescriptor.id();
if (currentProcessId.equals(virtualMachineId)) {
continue;
}

processHandleFound = findProcessHandle(virtualMachineDescriptor, systemProperty, configuration);
if (processHandleFound.isPresent()) {
break;
}
}

return processHandleFound;
}

private Optional<ProcessHandle> findProcessHandle(final VirtualMachineDescriptor descriptor, final SystemProperty systemProperty, final Path configuration) {
final AttachProvider attachProvider = descriptor.provider();
final String virtualMachineId = descriptor.id();

Optional<ProcessHandle> processHandle = Optional.empty();
try {
final VirtualMachine virtualMachine = attachProvider.attachVirtualMachine(virtualMachineId);
logger.debug("Attached Virtual Machine [{}]", virtualMachine.id());
try {
processHandle = findProcessHandle(virtualMachine, systemProperty, configuration);
} finally {
virtualMachine.detach();
}
} catch (final Exception e) {
logger.debug("Attach Virtual Machine [{}] failed", virtualMachineId, e);
}

return processHandle;
}

private Optional<ProcessHandle> findProcessHandle(final VirtualMachine virtualMachine, final SystemProperty systemProperty, final Path configuration) throws IOException {
final Properties systemProperties = virtualMachine.getSystemProperties();
final String configurationProperty = systemProperties.getProperty(systemProperty.getProperty());
final String configurationPath = configuration.toString();

final Optional<ProcessHandle> processHandle;

if (configurationPath.equals(configurationProperty)) {
final String virtualMachineId = virtualMachine.id();
final long processId = Long.parseLong(virtualMachineId);
processHandle = ProcessHandle.of(processId);
} else {
processHandle = Optional.empty();
}

return processHandle;
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,63 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one or more
* contributor license agreements. See the NOTICE file distributed with
* this work for additional information regarding copyright ownership.
* The ASF licenses this file to You under the Apache License, Version 2.0
* (the "License"); you may not use this file except in compliance with
* the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.apache.nifi.bootstrap.command.process;

import org.apache.nifi.bootstrap.configuration.ConfigurationProvider;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;
import org.junit.jupiter.api.io.TempDir;
import org.mockito.Mock;
import org.mockito.junit.jupiter.MockitoExtension;

import java.nio.file.Path;
import java.util.Optional;

import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.mockito.Mockito.when;

@ExtendWith(MockitoExtension.class)
class VirtualMachineProcessHandleProviderTest {
@Mock
private ConfigurationProvider configurationProvider;

@TempDir
private Path tempDir;

private VirtualMachineProcessHandleProvider provider;

@BeforeEach
void setProvider() {
provider = new VirtualMachineProcessHandleProvider(configurationProvider);
}

@Test
void testFindApplicationProcessHandleEmpty() {
when(configurationProvider.getApplicationProperties()).thenReturn(tempDir);
final Optional<ProcessHandle> applicationProcessHandle = provider.findApplicationProcessHandle();

assertTrue(applicationProcessHandle.isEmpty());
}

@Test
void testFindBootstrapProcessHandleEmpty() {
when(configurationProvider.getBootstrapConfiguration()).thenReturn(tempDir);
final Optional<ProcessHandle> bootstrapProcessHandle = provider.findBootstrapProcessHandle();

assertTrue(bootstrapProcessHandle.isEmpty());
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,7 @@ set CONFIG_FILE_PROPERTY=-Dorg.apache.nifi.bootstrap.config.file=%CONF_DIR%\boot
set PROPERTIES_FILE_PROPERTY=-Dnifi.properties.file.path=%CONF_DIR%\nifi.properties
set BOOTSTRAP_HEAP_SIZE=48m

set JAVA_ARGS=%LOG_DIR_PROPERTY% %CONFIG_FILE_PROPERTY% %PROPERTIES_FILE_PROPERTY%
set JAVA_ARGS=%LOG_DIR_PROPERTY% %CONFIG_FILE_PROPERTY%
set JAVA_PARAMS=-cp %BOOTSTRAP_LIB_DIR%\*;%CONF_DIR% %JAVA_ARGS%
set JAVA_MEMORY=-Xms%BOOTSTRAP_HEAP_SIZE% -Xmx%BOOTSTRAP_HEAP_SIZE%

Expand All @@ -47,11 +47,14 @@ set RUN_COMMAND="%~1"
if %RUN_COMMAND% == "set-single-user-credentials" (
rem Set credentials with quoted arguments passed to Java command
set "CREDENTIALS=^"%~2^" ^"%~3^""
call "%JAVA_EXE%" %JAVA_PARAMS% org.apache.nifi.authentication.single.user.command.SetSingleUserCredentials %CREDENTIALS%
call "%JAVA_EXE%" %JAVA_PARAMS% %PROPERTIES_FILE_PROPERTY% org.apache.nifi.authentication.single.user.command.SetSingleUserCredentials %CREDENTIALS%
) else if %RUN_COMMAND% == "set-sensitive-properties-key" (
call "%JAVA_EXE%" %JAVA_PARAMS% org.apache.nifi.flow.encryptor.command.SetSensitivePropertiesKey %~2
call "%JAVA_EXE%" %JAVA_PARAMS% %PROPERTIES_FILE_PROPERTY% org.apache.nifi.flow.encryptor.command.SetSensitivePropertiesKey %~2
) else if %RUN_COMMAND% == "set-sensitive-properties-algorithm" (
call "%JAVA_EXE%" %JAVA_PARAMS% org.apache.nifi.flow.encryptor.command.SetSensitivePropertiesAlgorithm %~2
call "%JAVA_EXE%" %JAVA_PARAMS% %PROPERTIES_FILE_PROPERTY% org.apache.nifi.flow.encryptor.command.SetSensitivePropertiesAlgorithm %~2
) else if %RUN_COMMAND% == "start" (
rem Start bootstrap process in new minimized window
call start /MIN "Apache NiFi" "%JAVA_EXE%" %JAVA_MEMORY% %JAVA_PARAMS% org.apache.nifi.bootstrap.BootstrapProcess %RUN_COMMAND%
) else (
call "%JAVA_EXE%" %JAVA_MEMORY% %JAVA_PARAMS% org.apache.nifi.bootstrap.BootstrapProcess %RUN_COMMAND%
)
Expand Down

0 comments on commit 2b1055e

Please sign in to comment.