import datagen from 'k6/x/frostfs/datagen'; import logging from 'k6/x/frostfs/logging'; import registry from 'k6/x/frostfs/registry'; import s3 from 'k6/x/frostfs/s3'; import {SharedArray} from 'k6/data'; import {sleep} from 'k6'; import {textSummary} from './libs/k6-summary-0.0.2.js'; import {parseEnv} from './libs/env-parser.js'; import {uuidv4} from './libs/k6-utils-1.4.0.js'; parseEnv(); const bucket_list = new SharedArray('bucket_list', function () { return JSON.parse(open(__ENV.PREGEN_JSON)).buckets; }); const summary_json = __ENV.SUMMARY_JSON || "/tmp/summary.json"; // Select random S3 endpoint for current VU const s3_endpoints = __ENV.S3_ENDPOINTS.split(','); const s3_endpoint = s3_endpoints[Math.floor(Math.random() * s3_endpoints.length)]; const s3_client = s3.connect(`http://${s3_endpoint}`); const log = logging.new().withField("endpoint", s3_endpoint); const registry_enabled = !!__ENV.REGISTRY_FILE; const obj_registry = registry_enabled ? registry.open(__ENV.REGISTRY_FILE) : undefined; const duration = __ENV.DURATION; const generator = datagen.generator(1024 * parseInt(__ENV.WRITE_OBJ_SIZE), __ENV.PAYLOAD_TYPE || ""); const scenarios = {}; const write_vu_count = parseInt(__ENV.WRITERS || '0'); if (write_vu_count < 1) { throw 'number of VUs (env WRITERS) performing write operations should be greater than 0'; } const write_multipart_vu_count = parseInt(__ENV.WRITERS_MULTIPART || '0'); if (write_multipart_vu_count < 1) { throw 'number of parts (env WRITERS_MULTIPART) to upload in parallel should be greater than 0'; } if (write_vu_count > 0) { scenarios.write_multipart = { executor: 'constant-vus', vus: write_vu_count, duration: `${duration}s`, exec: 'obj_write_multipart', gracefulStop: '5s', }; } export const options = { scenarios, setupTimeout: '5s', }; export function setup() { const total_vu_count = write_vu_count * write_multipart_vu_count; console.log(`Pregenerated buckets: ${bucket_list.length}`); console.log(`Writing VUs: ${write_vu_count}`); console.log(`Writing multipart VUs: ${write_multipart_vu_count}`); console.log(`Total VUs: ${total_vu_count}`); } export function teardown(data) { if (obj_registry) { obj_registry.close(); } } export function handleSummary(data) { return { 'stdout': textSummary(data, {indent: ' ', enableColors: false}), [summary_json]: JSON.stringify(data), }; } const write_multipart_part_size = 1024 * parseInt(__ENV.WRITE_OBJ_PART_SIZE || '0') if (write_multipart_part_size < 5 * 1024 * 1024) { throw 'part size (env WRITE_OBJ_PART_SIZE * 1024) must be greater than (5 MB)'; } export function obj_write_multipart() { if (__ENV.SLEEP_WRITE) { sleep(__ENV.SLEEP_WRITE); } const key = __ENV.OBJ_NAME || uuidv4(); const bucket = bucket_list[Math.floor(Math.random() * bucket_list.length)]; const {payload, hash} = generator.genPayload(registry_enabled); const resp = s3_client.multipart(bucket, key, write_multipart_part_size, write_multipart_vu_count, payload); if (!resp.success) { log.withFields({bucket: bucket, key: key}).error(resp.error); return; } if (obj_registry) { obj_registry.addObject("", "", bucket, key, hash); } }