| /* |
| * 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. |
| */ |
| |
| import CommonTestProperties.Runner |
| import CommonTestProperties.SDK |
| import LoadTestConfig.SerializableOption |
| import groovy.json.JsonBuilder |
| import org.codehaus.groovy.runtime.InvokerHelper |
| |
| import java.util.function.Predicate |
| |
| import static java.util.Objects.nonNull |
| import static java.util.Objects.requireNonNull |
| |
| /** |
| * This class contains simple DSL for load tests configuration. Configuration as Map<String, Serializable> |
| * [{@link LoadTestConfig#config config} -- returns configuration map] |
| * [{@link LoadTestConfig#templateConfig templateConfig} -- return LoadTestConfig reusable object] |
| * [{@link LoadTestConfig#fromTemplate fromTemplate} -- returns configuration from given template].<br><br> |
| * |
| * Example: |
| * <blockquote><pre> |
| * LoadTestConfig template = templateConfig { |
| * title 'Load test' |
| * test 'org.apache.beam.sdk.loadtests.SomeLoadTests' |
| * dataflow() |
| * pipelineOptions { |
| * python() |
| * jobName 'Any job name' |
| * publishToBigQuery true |
| * //other fields |
| * } |
| * specificParameters([ |
| * fanout: 4 |
| * ]) |
| * } |
| * Map<String, Serializable> configMap = fromTemplate(template) { |
| * //fields can be changed or/and added |
| * portable() |
| * pipelineOptions { |
| * parallelism 5 |
| * inputOptions { |
| * numRecords 20000 |
| * keySize 1000 |
| * valueSize 10 |
| * } |
| * } |
| * } |
| * </pre></blockquote> |
| */ |
| class LoadTestConfig implements SerializableOption<Map<String, Serializable>> { |
| |
| private String _title |
| private String _test |
| private Runner _runner |
| private PipelineOptions _pipelineOptions |
| |
| private LoadTestConfig() {} |
| |
| void title(final String title) { |
| _title = title |
| } |
| void test(final String test) { |
| _test = test |
| } |
| |
| //runners |
| void dataflow() { setRunnerAndUpdatePipelineOptions(Runner.DATAFLOW)} |
| void portable() { setRunnerAndUpdatePipelineOptions(Runner.PORTABLE) } |
| |
| private void setRunnerAndUpdatePipelineOptions(final Runner runner) { |
| _runner = runner |
| final def pipeline = _pipelineOptions ?: new PipelineOptions() |
| pipeline.i_runner = runner |
| _pipelineOptions = pipeline |
| } |
| |
| void pipelineOptions(final Closure cl = {}) { |
| final def options = _pipelineOptions ?: new PipelineOptions() |
| delegateAndInvoke(options, cl) |
| _pipelineOptions = options |
| } |
| |
| /** |
| * Returns load test config object which can be reusable.</br> |
| * All possible fields that can be set: |
| * <blockquote><pre> |
| * templateConfig { |
| * title [String] |
| * test [String] |
| * [dataflow(), portable()] -- runner |
| * pipelineOptions { |
| * [python(), java()] -- sdk |
| * jobName [String] |
| * appName [String] |
| * project [String] |
| * publishToBigQuery [boolean] |
| * metricsDataset (python) [String] |
| * metricsTable (python) [String] |
| * bigQueryDataset (java) [String] |
| * bigQueryTable (java) [String] |
| * numWorkers [int] |
| * parallelism [int] |
| * tempLocation [String] |
| * autoscalingAlgorithm [String] |
| * jobEndpoint [String] |
| * environmentType [String] |
| * environmentConfig [String] |
| * inputOptions/coInputOptions (for python) { |
| * numRecords [int] |
| * keySize [int] |
| * valueSize [int] |
| * numHotKeys [int] |
| * hotKeyFraction [int] |
| * } |
| * sourceOptions/coSourceOptions (for java) { |
| * numRecords [int] |
| * keySizeBytes [int] |
| * valueSizeBytes [int] |
| * numHotKeys [int] |
| * hotKeyFraction [int] |
| * splitPointFrequencyRecords [int] |
| * } |
| * stepOptions { |
| * outputRecordsPerInputRecord [int] |
| * preservesInputKeyDistribution [boolean] |
| * } |
| * specificParameters [Map<String, Object>] |
| * } |
| * } |
| * </pre></blockquote> |
| * @param cl Closure with fields setting |
| * @return LoadTestConfig object |
| */ |
| static LoadTestConfig templateConfig(final Closure cl = {}) { |
| final def config = new LoadTestConfig() |
| delegateAndInvoke(config, cl) |
| return config |
| } |
| |
| /** |
| * Returns configuration map from given template. Any field can be changed or/and added. Validation is performed |
| * before final map is returned (ex. Flink runner requires <b>environmentConfig</b> to be set). In case of |
| * validation failure exception is thrown.<br> |
| * Example result: |
| *<blockquote><pre> |
| * [ |
| * title : 'any given title', |
| * test : 'org.apache.beam.sdk.loadtests.SomeLoadTests', |
| * runner : CommonTestProperties.Runner.DATAFLOW, |
| * pipelineOptions: [ |
| * job_name : 'any given job name', |
| * publish_to_big_query: true, |
| * project : 'apache-beam-testing', |
| * metrics_dataset : 'given_dataset_name', |
| * metrics_table : 'given_table_name', |
| * input_options : '\'{"num_records": 200000000,"key_size": 1,"value_size":9}\'', |
| * iterations : 1, |
| * fanout : 1, |
| * parallelism : 5, |
| * job_endpoint : 'localhost:1234', |
| * environment_config : 'given_environment_config', |
| * environment_type : 'given_environment_type' |
| * ] |
| * ] |
| * </blockquote></pre> |
| * @param templateConfig LoadTestConfig instance |
| * @param cl Closure with fields setting |
| * @return configuration map |
| * @see LoadTestConfig |
| * @see LoadTestConfig#templateConfig |
| */ |
| static Map<String, Serializable> fromTemplate(final LoadTestConfig templateConfig, final Closure cl = {}) { |
| final def newConfig = of(templateConfig) |
| delegateAndInvoke(newConfig, cl) |
| final def properties = newConfig.propertiesMap |
| verifyProperties(properties) |
| return ConfigHelper.convertProperties(properties) |
| } |
| |
| /** |
| * Returns configuration map (see {@link LoadTestConfig#fromTemplate}) directly from given settings |
| * @param cl Closure with settings |
| * @return configuration map |
| */ |
| static Map<String, Serializable> config(final Closure cl = {}) { |
| final def config = new LoadTestConfig() |
| delegateAndInvoke(config, cl) |
| final def properties = config.propertiesMap |
| verifyProperties(properties) |
| return ConfigHelper.convertProperties(config.propertiesMap) |
| } |
| |
| private static void delegateAndInvoke(final delegate, final Closure cl = {}) { |
| final def code = cl.rehydrate(delegate, this, this) |
| code.resolveStrategy = Closure.DELEGATE_ONLY |
| code() |
| } |
| |
| private static LoadTestConfig of(final LoadTestConfig oldConfig) { |
| final def newConfig = new LoadTestConfig() |
| |
| //primitive values |
| InvokerHelper.setProperties(newConfig, oldConfig.propertiesMap) |
| |
| //non-primitive values |
| newConfig._pipelineOptions = oldConfig._pipelineOptions ? PipelineOptions.of(oldConfig._pipelineOptions) : null |
| |
| return newConfig |
| } |
| |
| @Override |
| Map<String, Serializable> toPrimitiveValues() { |
| final def map = propertiesMap |
| verifyProperties(map) |
| return ConfigHelper.convertProperties(map) |
| } |
| |
| LinkedHashMap<String, Object> getPropertiesMap() { |
| return [ |
| _title: _title, |
| _test: _test, |
| _runner: _runner, |
| _pipelineOptions: _pipelineOptions |
| ] |
| } |
| |
| private static void verifyProperties(final LinkedHashMap<String, Object> map) { |
| for (entry in map.entrySet()) { |
| requireNonNull(entry.value, "Missing ${entry.key.substring(1)} in configuration") |
| } |
| } |
| |
| private static class PipelineOptions implements SerializableOption<Map<String, Serializable>> { |
| private Map<String, Object> _specificParameters = new HashMap<>() |
| private boolean _streaming = false |
| private SourceOptions _coSourceOptions |
| private InputOptions _coInputOptions |
| private StepOptions _stepOptions |
| |
| //required |
| private String _project |
| private String _publishToBigQuery |
| |
| //java required |
| private String _bigQueryDataset |
| private String _bigQueryTable |
| private String _appName |
| private SourceOptions _sourceOptions |
| |
| //python required |
| private String _metricsDataset |
| private String _metricsTable |
| private String _jobName |
| private InputOptions _inputOptions |
| |
| //internal usage |
| private SDK i_sdk |
| private Runner i_runner |
| private static final i_required = [ |
| "_project", |
| "_publishToBigQuery" |
| ] |
| private static final i_dataflowRequired = [ |
| "_numWorkers", |
| "_tempLocation", |
| "_autoscalingAlgorithm", |
| "_region" |
| ] |
| private static final i_portableRequired = [ |
| "_jobEndpoint", |
| "_environmentType", |
| "_environmentConfig", |
| "_parallelism" |
| ] |
| private static final i_javaRequired = [ |
| "_bigQueryDataset", |
| "_bigQueryTable", |
| "_sourceOptions", |
| "_appName" |
| ] |
| private static final i_pythonRequired = [ |
| "_metricsDataset", |
| "_metricsTable", |
| "_inputOptions", |
| "_jobName" |
| ] |
| |
| //dataflow required |
| private def _numWorkers |
| private String _tempLocation |
| private String _autoscalingAlgorithm |
| private String _region = 'us-central1' |
| |
| //flink required |
| private String _jobEndpoint |
| private String _environmentType |
| private String _environmentConfig |
| private def _parallelism |
| |
| void jobName(final String name) { _jobName = name } |
| void appName(final String name) { _appName = name } |
| void project(final String project) { _project = project } |
| void tempLocation(final String location) { _tempLocation = location } |
| void publishToBigQuery(final boolean publish) { _publishToBigQuery = publish } |
| void metricsDataset(final String dataset) { _metricsDataset = dataset } |
| void metricsTable(final String table) { _metricsTable = table } |
| void inputOptions(final InputOptions options) { _inputOptions = options } |
| void numWorkers(final int workers) { _numWorkers = workers } |
| void autoscalingAlgorithm(final String algorithm) { _autoscalingAlgorithm = algorithm } |
| void region(final String region) { _region = region } |
| void jobEndpoint(final String endpoint) { _jobEndpoint = endpoint } |
| void environmentType(final String type) { _environmentType = type } |
| void environmentConfig(final String config) { _environmentConfig = config } |
| void parallelism(final int parallelism) { _parallelism = parallelism } |
| void bigQueryDataset(final String dataset) { _bigQueryDataset = dataset } |
| void bigQueryTable(final String table) { _bigQueryTable = table } |
| void streaming(final boolean isStreaming) { _streaming = isStreaming } |
| void sourceOptions(final Closure cl = {}) { _sourceOptions = makeSourceOptions(cl) } |
| void coSourceOptions(final Closure cl = {}) { _coSourceOptions = makeSourceOptions(cl) } |
| void inputOptions(final Closure cl = {}) { _inputOptions = makeInputOptions(cl) } |
| void coInputOptions(final Closure cl = {}) { _coInputOptions = makeInputOptions(cl) } |
| void stepOptions(final Closure cl = {}) { _stepOptions = makeStepOptions(cl) } |
| void specificParameters(final Map<String, Object> map) { _specificParameters.putAll(map) } |
| |
| //sdk -- snake_case vs camelCase |
| void python() { i_sdk = SDK.PYTHON } |
| void java() { i_sdk = SDK.JAVA } |
| |
| |
| private InputOptions makeInputOptions(final Closure cl = {}) { |
| return makeOptions(cl, _inputOptions ?: InputOptions.withSDK(i_sdk)) |
| } |
| |
| private SourceOptions makeSourceOptions(final Closure cl = {}) { |
| return makeOptions(cl, _sourceOptions ?: SourceOptions.withSDK(i_sdk)) |
| } |
| |
| private StepOptions makeStepOptions(final Closure cl = {}) { |
| return makeOptions(cl, _stepOptions ?: StepOptions.withSDK(i_sdk)) |
| } |
| |
| private <T> T makeOptions(final Closure cl = {}, final T options) { |
| final def code = cl.rehydrate(options, this, this) |
| code.resolveStrategy = Closure.DELEGATE_ONLY |
| code() |
| return options |
| } |
| |
| @Override |
| Map<String, Serializable> toPrimitiveValues() { |
| final def map = propertiesMap |
| verifyPipelineProperties(map) |
| return ConfigHelper.convertProperties(map, i_sdk) |
| } |
| |
| private void verifyPipelineProperties(final Map<String, Object> map) { |
| verifyRequired(map) |
| switch (i_runner) { |
| case Runner.DATAFLOW: |
| verifyDataflowProperties(map) |
| break |
| case Runner.PORTABLE: |
| verifyPortableProperties(map) |
| break |
| default: |
| break |
| } |
| } |
| |
| private void verifyRequired(final Map<String, Object> map) { |
| verifyCommonRequired(map) |
| switch (i_sdk) { |
| case SDK.PYTHON: |
| verifyPythonRequired(map) |
| break |
| case SDK.JAVA: |
| verifyJavaRequired(map) |
| break |
| default: |
| break |
| } |
| } |
| |
| private static void verifyCommonRequired(final Map<String, Object> map) { |
| verify(map, "") { i_required.contains(it.key) } |
| } |
| |
| private static void verifyPythonRequired(final Map<String, Object> map) { |
| verify(map, "for Python SDK") { i_pythonRequired.contains(it.key) } |
| } |
| |
| private static void verifyJavaRequired(final Map<String, Object> map) { |
| verify(map, "for Java SDK") { i_javaRequired.contains(it.key) } |
| } |
| |
| private static void verifyDataflowProperties(final Map<String, Object> map) { |
| verify(map, "for Dataflow runner") { i_dataflowRequired.contains(it.key) } |
| } |
| |
| private static void verifyPortableProperties(final Map<String, Object> map) { |
| verify(map, "for Portable runner") { i_portableRequired.contains(it.key) } |
| } |
| |
| private static void verify(final Map<String, Object> map, final String message, final Predicate<Map.Entry<String, Object>> predicate) { |
| map.entrySet() |
| .stream() |
| .filter(predicate) |
| .forEach{ requireNonNull(it.value, "${it.key.substring(1)} is required " + message) } |
| } |
| |
| static PipelineOptions of(final PipelineOptions options) { |
| final def newOptions = new PipelineOptions() |
| |
| //primitive values |
| InvokerHelper.setProperties(newOptions, options.propertiesMap) |
| |
| //non-primitive |
| newOptions._inputOptions = options._inputOptions ? InputOptions.of(options._inputOptions) : null |
| newOptions._coInputOptions = options._coInputOptions ? InputOptions.of(options._coInputOptions) : null |
| newOptions._sourceOptions = options._sourceOptions ? SourceOptions.of(options._sourceOptions) : null |
| newOptions._coSourceOptions = options._coSourceOptions ? SourceOptions.of(options._coSourceOptions) : null |
| newOptions._stepOptions = options._stepOptions ? StepOptions.of(options._stepOptions) : null |
| newOptions._specificParameters = new HashMap<>(options._specificParameters) |
| |
| return newOptions |
| } |
| |
| Map<String, Object> getPropertiesMap() { |
| return [ |
| i_sdk: i_sdk, |
| i_runner: i_runner, |
| _jobName: _jobName, |
| _appName: _appName, |
| _project: _project, |
| _tempLocation: _tempLocation, |
| _publishToBigQuery: _publishToBigQuery, |
| _metricsDataset: _metricsDataset, |
| _metricsTable: _metricsTable, |
| _numWorkers: _numWorkers, |
| _autoscalingAlgorithm: _autoscalingAlgorithm, |
| _region: _region, |
| _inputOptions: _inputOptions, |
| _coInputOptions: _coInputOptions, |
| _jobEndpoint: _jobEndpoint, |
| _environmentType: _environmentType, |
| _environmentConfig: _environmentConfig, |
| _parallelism: _parallelism, |
| _bigQueryDataset: _bigQueryDataset, |
| _bigQueryTable: _bigQueryTable, |
| _streaming: _streaming, |
| _sourceOptions: _sourceOptions, |
| _coSourceOptions: _coSourceOptions, |
| _stepOptions: _stepOptions |
| ].putAll(_specificParameters.entrySet()) |
| } |
| |
| private static class InputOptions implements SerializableOption<String> { |
| private def _numRecords |
| private def _keySize |
| private def _valueSize |
| private def _numHotKeys |
| private def _hotKeyFraction |
| |
| //internal usage |
| private SDK i_sdk |
| |
| private InputOptions() {} |
| |
| static withSDK(final SDK sdk) { |
| final def input = new InputOptions() |
| input.i_sdk = sdk |
| return input |
| } |
| |
| void numRecords(final int num) { _numRecords = num } |
| void keySize(final int size) { _keySize = size } |
| void valueSize(final int size) { _valueSize = size } |
| void numHotsKeys(final int num) { _numHotKeys = num } |
| void hotKeyFraction(final int fraction) { _hotKeyFraction = fraction } |
| |
| @Override |
| String toPrimitiveValues() { |
| return "'${new JsonBuilder(ConfigHelper.convertProperties(propertiesMap, i_sdk)).toString()}'" |
| } |
| |
| static InputOptions of(final InputOptions oldOptions) { |
| final def newOptions = new InputOptions() |
| InvokerHelper.setProperties(newOptions, oldOptions.propertiesMap) |
| return newOptions |
| } |
| |
| LinkedHashMap<String, Object> getPropertiesMap() { |
| return [ |
| i_sdk: i_sdk, |
| _numRecords: _numRecords, |
| _keySize: _keySize, |
| _valueSize: _valueSize, |
| _numHotKeys: _numHotKeys, |
| _hotKeyFraction: _hotKeyFraction |
| ] as LinkedHashMap<String, Object> |
| } |
| } |
| |
| private static class SourceOptions implements SerializableOption<String> { |
| private def _numRecords |
| private def _keySizeBytes |
| private def _valueSizeBytes |
| private def _numHotKeys |
| private def _hotKeyFraction |
| private def _splitPointFrequencyRecords |
| |
| //internal usage |
| private SDK i_sdk |
| |
| private SourceOptions() {} |
| |
| static withSDK(final SDK sdk) { |
| final def input = new SourceOptions() |
| input.i_sdk = sdk |
| return input |
| } |
| |
| void numRecords(final int num) { _numRecords = num } |
| void keySizeBytes(final int size) { _keySizeBytes = size } |
| void valueSizeBytes(final int size) { _valueSizeBytes = size } |
| void numHotsKeys(final int num) { _numHotKeys = num } |
| void hotKeyFraction(final int fraction) { _hotKeyFraction = fraction } |
| void splitPointFrequencyRecords(final int splitPoint) { _splitPointFrequencyRecords = splitPoint } |
| |
| @Override |
| String toPrimitiveValues() { |
| return new JsonBuilder(ConfigHelper.convertProperties(propertiesMap, i_sdk)).toString() |
| } |
| |
| static SourceOptions of(final SourceOptions oldOptions) { |
| final def newOptions = new SourceOptions() |
| InvokerHelper.setProperties(newOptions, oldOptions.propertiesMap) |
| return newOptions |
| } |
| |
| Map<String, Object> getPropertiesMap() { |
| return [ |
| i_sdk: i_sdk, |
| _numRecords: _numRecords, |
| _keySizeBytes: _keySizeBytes, |
| _valueSizeBytes: _valueSizeBytes, |
| _numHotKeys: _numHotKeys, |
| _hotKeyFraction: _hotKeyFraction, |
| _splitPointFrequencyRecords: _splitPointFrequencyRecords |
| ] |
| } |
| } |
| |
| private static class StepOptions implements SerializableOption<String> { |
| private def _outputRecordsPerInputRecord |
| private boolean _preservesInputKeyDistribution |
| |
| //internal usage |
| private SDK i_sdk |
| |
| private StepOptions() {} |
| |
| static withSDK(final SDK sdk) { |
| final def option = new StepOptions() |
| option.i_sdk = sdk |
| return option |
| } |
| |
| void outputRecordsPerInputRecord(final int records) { _outputRecordsPerInputRecord = records } |
| void preservesInputKeyDistribution(final boolean shouldPreserve) { _preservesInputKeyDistribution = shouldPreserve } |
| |
| @Override |
| String toPrimitiveValues() { |
| return new JsonBuilder(ConfigHelper.convertProperties(propertiesMap, i_sdk)).toString() |
| } |
| |
| Map<String, Object> getPropertiesMap() { |
| return [ |
| i_sdk: i_sdk, |
| _outputRecordsPerInputRecord: _outputRecordsPerInputRecord, |
| _preservesInputKeyDistribution: _preservesInputKeyDistribution |
| ] as Map<String, Object> |
| } |
| |
| static StepOptions of(final StepOptions oldOption) { |
| final def newOption = new StepOptions() |
| InvokerHelper.setProperties(newOption, oldOption.propertiesMap) |
| return newOption |
| } |
| } |
| } |
| |
| private interface SerializableOption<T> { |
| T toPrimitiveValues() |
| } |
| |
| private static class ConfigHelper { |
| private static final List<String> FIELDS_TO_REMOVE = ["class", "i_sdk", "i_runner"] |
| |
| static Map<String, Serializable> convertProperties(final Map<String, Object> propertyMap, final SDK sdk = SDK.JAVA) { |
| return propertyMap |
| .findAll { nonNull(it.value) } |
| .findAll { !FIELDS_TO_REMOVE.contains(it.key) } |
| .collectEntries { key, value -> |
| [ |
| modifyKey(key, sdk), |
| toPrimitive(value) |
| ] |
| } as Map<String, Serializable> |
| } |
| |
| private static String modifyKey(final String key, final SDK sdk) { |
| final def result = key.startsWith('_') ? key.substring(1) : key |
| switch (sdk) { |
| case SDK.PYTHON: |
| return toSnakeCase(result) |
| case SDK.JAVA: |
| return toCamelCase(result) |
| default: |
| throw new IllegalArgumentException("SDK not specified") |
| } |
| } |
| |
| private static String toSnakeCase(final String text) { |
| return text.replaceAll(/([A-Z])/, /_$1/).toLowerCase().replaceAll(/^_/, '') |
| } |
| |
| private static String toCamelCase(final String text) { |
| return text.replaceAll( "(_)([A-Za-z0-9])", { Object[] it -> ((String) it[2]).toUpperCase() }) |
| } |
| |
| private static def toPrimitive(value) { |
| return value instanceof SerializableOption |
| ? value.toPrimitiveValues() |
| : value |
| } |
| } |
| } |
