hadoop-ozone/integration-test/src/test/java/org/apache/hadoop/ozone/MiniOzoneLoadGenerator.java - hadoop - Git at Google

 /**
  * 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.hadoop.ozone;

 import org.apache.commons.lang3.RandomUtils;
 import org.apache.hadoop.conf.StorageUnit;
 import org.apache.hadoop.hdds.client.ReplicationFactor;
 import org.apache.hadoop.hdds.client.ReplicationType;
 import org.apache.hadoop.ozone.chaos.TestProbability;
 import org.apache.hadoop.ozone.client.OzoneBucket;
 import org.apache.hadoop.ozone.client.io.OzoneInputStream;
 import org.apache.hadoop.ozone.client.io.OzoneOutputStream;
 import org.apache.hadoop.util.Time;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;

 import java.io.IOException;
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Optional;
 import java.util.concurrent.ThreadPoolExecutor;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.ArrayBlockingQueue;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;

 /**
  * A Simple Load generator for testing.
  */
 public class MiniOzoneLoadGenerator {

   private static final Logger LOG =
       LoggerFactory.getLogger(MiniOzoneLoadGenerator.class);

   private static String keyNameDelimiter = "_";

   private ThreadPoolExecutor writeExecutor;
   private int numWriteThreads;
   // number of buffer to be allocated, each is allocated with length which
   // is multiple of 2, each buffer is populated with random data.
   private int numBuffers;
   private List<ByteBuffer> buffers;

   private AtomicBoolean isWriteThreadRunning;

   private final List<OzoneBucket> ozoneBuckets;

   private final AtomicInteger agedFileWrittenIndex;
   private final ExecutorService agedFileExecutor;
   private final OzoneBucket agedLoadBucket;
   private final TestProbability agedWriteProbability;

   MiniOzoneLoadGenerator(List<OzoneBucket> bucket,
                          OzoneBucket agedLoadBucket, int numThreads,
       int numBuffers) {
     this.ozoneBuckets = bucket;
     this.numWriteThreads = numThreads;
     this.numBuffers = numBuffers;
     this.writeExecutor = new ThreadPoolExecutor(numThreads, numThreads, 100,
         TimeUnit.SECONDS, new ArrayBlockingQueue<>(1024),
         new ThreadPoolExecutor.CallerRunsPolicy());
     this.writeExecutor.prestartAllCoreThreads();

     this.agedFileWrittenIndex = new AtomicInteger(0);
     this.agedFileExecutor = Executors.newSingleThreadExecutor();
     this.agedLoadBucket = agedLoadBucket;
     this.agedWriteProbability = TestProbability.valueOf(10);

     this.isWriteThreadRunning = new AtomicBoolean(false);

     // allocate buffers and populate random data.
     buffers = new ArrayList<>();
     for (int i = 0; i < numBuffers; i++) {
       int size = (int) StorageUnit.KB.toBytes(1 << i);
       ByteBuffer buffer = ByteBuffer.allocate(size);
       buffer.put(RandomUtils.nextBytes(size));
       buffers.add(buffer);
     }
   }

   // Start IO load on an Ozone bucket.
   private void load(long runTimeMillis) {
     long threadID = Thread.currentThread().getId();
     LOG.info("Started IO Thread:{}.", threadID);
     String threadName = Thread.currentThread().getName();
     long startTime = Time.monotonicNow();

     while (isWriteThreadRunning.get() &&
         (Time.monotonicNow() < startTime + runTimeMillis)) {
       OzoneBucket bucket =
           ozoneBuckets.get((int) (Math.random() * ozoneBuckets.size()));
       try {
         int index = RandomUtils.nextInt();
         String keyName = writeData(index, bucket, threadName);

         readData(bucket, keyName, index);

         deleteKey(bucket, keyName);
       } catch (Exception e) {
         LOG.error("LOADGEN: Exiting due to exception", e);
         break;
       }
     }
     // This will terminate other threads too.
     isWriteThreadRunning.set(false);
     LOG.info("Terminating IO thread:{}.", threadID);
   }


   private String writeData(int keyIndex, OzoneBucket bucket, String threadName)
       throws Exception {
     // choose a random buffer.
     ByteBuffer buffer = buffers.get(keyIndex % numBuffers);
     int bufferCapacity = buffer.capacity();

     String keyName = getKeyName(keyIndex, threadName);
     LOG.trace("LOADGEN: Writing key {}", keyName);
     try (OzoneOutputStream stream = bucket.createKey(keyName,
         bufferCapacity, ReplicationType.RATIS, ReplicationFactor.THREE,
         new HashMap<>())) {
       stream.write(buffer.array());
       LOG.trace("LOADGEN: Written key {}", keyName);
     } catch (Throwable t) {
       LOG.error("LOADGEN: Create key:{} failed with exception, skipping",
           keyName, t);
       throw t;
     }

     return keyName;
   }

   private void readData(OzoneBucket bucket, String keyName, int index)
       throws Exception {
     LOG.trace("LOADGEN: Reading key {}", keyName);

     ByteBuffer buffer = buffers.get(index % numBuffers);
     int bufferCapacity = buffer.capacity();

     try (OzoneInputStream stream = bucket.readKey(keyName)) {
       byte[] readBuffer = new byte[bufferCapacity];
       int readLen = stream.read(readBuffer);

       if (readLen < bufferCapacity) {
         throw new IOException("Read mismatch, key:" + keyName +
             " read data length:" + readLen +
             " is smaller than excepted:" + bufferCapacity);
       }

       if (!Arrays.equals(readBuffer, buffer.array())) {
         throw new IOException("Read mismatch, key:" + keyName +
             " read data does not match the written data");
       }
       LOG.trace("LOADGEN: Read key {}", keyName);
     } catch (Throwable t) {
       LOG.error("LOADGEN: Read key:{} failed with exception", keyName, t);
       throw t;
     }
   }

   private void deleteKey(OzoneBucket bucket, String keyName) throws Exception {
     LOG.trace("LOADGEN: Deleting key {}", keyName);
     try {
       bucket.deleteKey(keyName);
       LOG.trace("LOADGEN: Deleted key {}", keyName);
     } catch (Throwable t) {
       LOG.error("LOADGEN: Unable to delete key:{}", keyName, t);
       throw t;
     }
   }

   private Optional<Integer> randomKeyToRead() {
     int currentIndex = agedFileWrittenIndex.get();
     return currentIndex != 0
       ? Optional.of(RandomUtils.nextInt(0, currentIndex))
       : Optional.empty();
   }

   private void startAgedFilesLoad(long runTimeMillis) {
     long threadID = Thread.currentThread().getId();
     LOG.info("AGED LOADGEN: Started Aged IO Thread:{}.", threadID);
     String threadName = Thread.currentThread().getName();
     long startTime = Time.monotonicNow();

     while (isWriteThreadRunning.get() &&
         (Time.monotonicNow() < startTime + runTimeMillis)) {

       String keyName = null;
       try {
         if (agedWriteProbability.isTrue()) {
           keyName = writeData(agedFileWrittenIndex.getAndIncrement(),
               agedLoadBucket, threadName);
         } else {
           Optional<Integer> index = randomKeyToRead();
           if (index.isPresent()) {
             keyName = getKeyName(index.get(), threadName);
             readData(agedLoadBucket, keyName, index.get());
           }
         }
       } catch (Throwable t) {
         LOG.error("AGED LOADGEN: {} Exiting due to exception", keyName, t);
         break;
       }
     }
     // This will terminate other threads too.
     isWriteThreadRunning.set(false);
     LOG.info("Terminating IO thread:{}.", threadID);
   }

   void startIO(long time, TimeUnit timeUnit) {
     List<CompletableFuture<Void>> writeFutures = new ArrayList<>();
     LOG.info("Starting MiniOzoneLoadGenerator for time {}:{} with {} buffers " +
             "and {} threads", time, timeUnit, numBuffers, numWriteThreads);
     if (isWriteThreadRunning.compareAndSet(false, true)) {
       // Start the IO thread
       for (int i = 0; i < numWriteThreads; i++) {
         writeFutures.add(
             CompletableFuture.runAsync(() -> load(timeUnit.toMillis(time)),
                 writeExecutor));
       }

       writeFutures.add(CompletableFuture.runAsync(() ->
               startAgedFilesLoad(timeUnit.toMillis(time)), agedFileExecutor));

       // Wait for IO to complete
       for (CompletableFuture<Void> f : writeFutures) {
         try {
           f.get();
         } catch (Throwable t) {
           LOG.error("startIO failed with exception", t);
         }
       }
     }
   }

   public void shutdownLoadGenerator() {
     try {
       writeExecutor.shutdown();
       writeExecutor.awaitTermination(1, TimeUnit.DAYS);
     } catch (Exception e) {
       LOG.error("error while closing ", e);
     }
   }

   private static String getKeyName(int keyIndex, String threadName) {
     return threadName + keyNameDelimiter + keyIndex;
   }
 }
	/**
	* 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.hadoop.ozone;

	import org.apache.commons.lang3.RandomUtils;
	import org.apache.hadoop.conf.StorageUnit;
	import org.apache.hadoop.hdds.client.ReplicationFactor;
	import org.apache.hadoop.hdds.client.ReplicationType;
	import org.apache.hadoop.ozone.chaos.TestProbability;
	import org.apache.hadoop.ozone.client.OzoneBucket;
	import org.apache.hadoop.ozone.client.io.OzoneInputStream;
	import org.apache.hadoop.ozone.client.io.OzoneOutputStream;
	import org.apache.hadoop.util.Time;
	import org.slf4j.Logger;
	import org.slf4j.LoggerFactory;

	import java.io.IOException;
	import java.nio.ByteBuffer;
	import java.util.ArrayList;
	import java.util.Arrays;
	import java.util.HashMap;
	import java.util.List;
	import java.util.Optional;
	import java.util.concurrent.ThreadPoolExecutor;
	import java.util.concurrent.CompletableFuture;
	import java.util.concurrent.TimeUnit;
	import java.util.concurrent.ArrayBlockingQueue;
	import java.util.concurrent.ExecutorService;
	import java.util.concurrent.Executors;
	import java.util.concurrent.atomic.AtomicBoolean;
	import java.util.concurrent.atomic.AtomicInteger;

	/**
	* A Simple Load generator for testing.
	*/
	public class MiniOzoneLoadGenerator {

	private static final Logger LOG =
	LoggerFactory.getLogger(MiniOzoneLoadGenerator.class);

	private static String keyNameDelimiter = "_";

	private ThreadPoolExecutor writeExecutor;
	private int numWriteThreads;
	// number of buffer to be allocated, each is allocated with length which
	// is multiple of 2, each buffer is populated with random data.
	private int numBuffers;
	private List<ByteBuffer> buffers;

	private AtomicBoolean isWriteThreadRunning;

	private final List<OzoneBucket> ozoneBuckets;

	private final AtomicInteger agedFileWrittenIndex;
	private final ExecutorService agedFileExecutor;
	private final OzoneBucket agedLoadBucket;
	private final TestProbability agedWriteProbability;

	MiniOzoneLoadGenerator(List<OzoneBucket> bucket,
	OzoneBucket agedLoadBucket, int numThreads,
	int numBuffers) {
	this.ozoneBuckets = bucket;
	this.numWriteThreads = numThreads;
	this.numBuffers = numBuffers;
	this.writeExecutor = new ThreadPoolExecutor(numThreads, numThreads, 100,
	TimeUnit.SECONDS, new ArrayBlockingQueue<>(1024),
	new ThreadPoolExecutor.CallerRunsPolicy());
	this.writeExecutor.prestartAllCoreThreads();

	this.agedFileWrittenIndex = new AtomicInteger(0);
	this.agedFileExecutor = Executors.newSingleThreadExecutor();
	this.agedLoadBucket = agedLoadBucket;
	this.agedWriteProbability = TestProbability.valueOf(10);

	this.isWriteThreadRunning = new AtomicBoolean(false);

	// allocate buffers and populate random data.
	buffers = new ArrayList<>();
	for (int i = 0; i < numBuffers; i++) {
	int size = (int) StorageUnit.KB.toBytes(1 << i);
	ByteBuffer buffer = ByteBuffer.allocate(size);
	buffer.put(RandomUtils.nextBytes(size));
	buffers.add(buffer);
	}
	}

	// Start IO load on an Ozone bucket.
	private void load(long runTimeMillis) {
	long threadID = Thread.currentThread().getId();
	LOG.info("Started IO Thread:{}.", threadID);
	String threadName = Thread.currentThread().getName();
	long startTime = Time.monotonicNow();

	while (isWriteThreadRunning.get() &&
	(Time.monotonicNow() < startTime + runTimeMillis)) {
	OzoneBucket bucket =
	ozoneBuckets.get((int) (Math.random() * ozoneBuckets.size()));
	try {
	int index = RandomUtils.nextInt();
	String keyName = writeData(index, bucket, threadName);

	readData(bucket, keyName, index);

	deleteKey(bucket, keyName);
	} catch (Exception e) {
	LOG.error("LOADGEN: Exiting due to exception", e);
	break;
	}
	}
	// This will terminate other threads too.
	isWriteThreadRunning.set(false);
	LOG.info("Terminating IO thread:{}.", threadID);
	}


	private String writeData(int keyIndex, OzoneBucket bucket, String threadName)
	throws Exception {
	// choose a random buffer.
	ByteBuffer buffer = buffers.get(keyIndex % numBuffers);
	int bufferCapacity = buffer.capacity();

	String keyName = getKeyName(keyIndex, threadName);
	LOG.trace("LOADGEN: Writing key {}", keyName);
	try (OzoneOutputStream stream = bucket.createKey(keyName,
	bufferCapacity, ReplicationType.RATIS, ReplicationFactor.THREE,
	new HashMap<>())) {
	stream.write(buffer.array());
	LOG.trace("LOADGEN: Written key {}", keyName);
	} catch (Throwable t) {
	LOG.error("LOADGEN: Create key:{} failed with exception, skipping",
	keyName, t);
	throw t;
	}

	return keyName;
	}

	private void readData(OzoneBucket bucket, String keyName, int index)
	throws Exception {
	LOG.trace("LOADGEN: Reading key {}", keyName);

	ByteBuffer buffer = buffers.get(index % numBuffers);
	int bufferCapacity = buffer.capacity();

	try (OzoneInputStream stream = bucket.readKey(keyName)) {
	byte[] readBuffer = new byte[bufferCapacity];
	int readLen = stream.read(readBuffer);

	if (readLen < bufferCapacity) {
	throw new IOException("Read mismatch, key:" + keyName +
	" read data length:" + readLen +
	" is smaller than excepted:" + bufferCapacity);
	}

	if (!Arrays.equals(readBuffer, buffer.array())) {
	throw new IOException("Read mismatch, key:" + keyName +
	" read data does not match the written data");
	}
	LOG.trace("LOADGEN: Read key {}", keyName);
	} catch (Throwable t) {
	LOG.error("LOADGEN: Read key:{} failed with exception", keyName, t);
	throw t;
	}
	}

	private void deleteKey(OzoneBucket bucket, String keyName) throws Exception {
	LOG.trace("LOADGEN: Deleting key {}", keyName);
	try {
	bucket.deleteKey(keyName);
	LOG.trace("LOADGEN: Deleted key {}", keyName);
	} catch (Throwable t) {
	LOG.error("LOADGEN: Unable to delete key:{}", keyName, t);
	throw t;
	}
	}

	private Optional<Integer> randomKeyToRead() {
	int currentIndex = agedFileWrittenIndex.get();
	return currentIndex != 0
	? Optional.of(RandomUtils.nextInt(0, currentIndex))
	: Optional.empty();
	}

	private void startAgedFilesLoad(long runTimeMillis) {
	long threadID = Thread.currentThread().getId();
	LOG.info("AGED LOADGEN: Started Aged IO Thread:{}.", threadID);
	String threadName = Thread.currentThread().getName();
	long startTime = Time.monotonicNow();

	while (isWriteThreadRunning.get() &&
	(Time.monotonicNow() < startTime + runTimeMillis)) {

	String keyName = null;
	try {
	if (agedWriteProbability.isTrue()) {
	keyName = writeData(agedFileWrittenIndex.getAndIncrement(),
	agedLoadBucket, threadName);
	} else {
	Optional<Integer> index = randomKeyToRead();
	if (index.isPresent()) {
	keyName = getKeyName(index.get(), threadName);
	readData(agedLoadBucket, keyName, index.get());
	}
	}
	} catch (Throwable t) {
	LOG.error("AGED LOADGEN: {} Exiting due to exception", keyName, t);
	break;
	}
	}
	// This will terminate other threads too.
	isWriteThreadRunning.set(false);
	LOG.info("Terminating IO thread:{}.", threadID);
	}

	void startIO(long time, TimeUnit timeUnit) {
	List<CompletableFuture<Void>> writeFutures = new ArrayList<>();
	LOG.info("Starting MiniOzoneLoadGenerator for time {}:{} with {} buffers " +
	"and {} threads", time, timeUnit, numBuffers, numWriteThreads);
	if (isWriteThreadRunning.compareAndSet(false, true)) {
	// Start the IO thread
	for (int i = 0; i < numWriteThreads; i++) {
	writeFutures.add(
	CompletableFuture.runAsync(() -> load(timeUnit.toMillis(time)),
	writeExecutor));
	}

	writeFutures.add(CompletableFuture.runAsync(() ->
	startAgedFilesLoad(timeUnit.toMillis(time)), agedFileExecutor));

	// Wait for IO to complete
	for (CompletableFuture<Void> f : writeFutures) {
	try {
	f.get();
	} catch (Throwable t) {
	LOG.error("startIO failed with exception", t);
	}
	}
	}
	}

	public void shutdownLoadGenerator() {
	try {
	writeExecutor.shutdown();
	writeExecutor.awaitTermination(1, TimeUnit.DAYS);
	} catch (Exception e) {
	LOG.error("error while closing ", e);
	}
	}

	private static String getKeyName(int keyIndex, String threadName) {
	return threadName + keyNameDelimiter + keyIndex;
	}
	}