package com.bazaarvoice.emodb.web.scanner.writer; import com.bazaarvoice.emodb.common.dropwizard.time.ClockTicker; import com.codahale.metrics.Counter; import com.codahale.metrics.MetricRegistry; import com.fasterxml.jackson.databind.ObjectMapper; import com.google.common.base.Function; import com.google.common.base.Optional; import com.google.common.base.Predicate; import com.google.common.base.Stopwatch; import com.google.common.collect.FluentIterable; import com.google.common.collect.ImmutableMap; import com.google.common.collect.Lists; import com.google.common.collect.Maps; import com.google.common.util.concurrent.FutureCallback; import com.google.common.util.concurrent.Futures; import com.google.common.util.concurrent.ListenableFuture; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import javax.annotation.Nullable; import java.io.File; import java.io.IOException; import java.io.OutputStream; import java.net.URI; import java.time.Duration; import java.time.Instant; import java.util.List; import java.util.Map; import java.util.concurrent.TimeUnit; import java.util.concurrent.locks.Condition; import java.util.concurrent.locks.ReentrantLock; import static com.google.common.base.Preconditions.checkArgument; import static com.google.common.base.Preconditions.checkNotNull; import static java.lang.String.format; /** * ScanWriter that writes to a temporary file first, then transfers it. */ abstract public class TemporaryFileScanWriter extends AbstractScanWriter { private final Logger _log = LoggerFactory.getLogger(TemporaryFileScanWriter.class); private final Counter _openTransfers; private final Counter _blockedNewShards; // Default number of shards that can be written and/or transferred concurrently private final static int DEFAULT_MAX_OPEN_SHARDS = 20; private final Map<TransferKey, ShardFiles> _openShardFiles = Maps.newHashMap(); private final int _maxOpenShards; private final ReentrantLock _lock = new ReentrantLock(); private final Condition _shardFilesClosedOrExceptionCaught = _lock.newCondition(); private volatile IOException _uploadException = null; private final ObjectMapper _mapper; protected TemporaryFileScanWriter(String type, int taskId, URI baseUri, Compression compression, MetricRegistry metricRegistry, Optional<Integer> maxOpenShards, ObjectMapper objectMapper) { super(type, taskId, baseUri, compression, metricRegistry); checkNotNull(maxOpenShards, "maxOpenShards"); _maxOpenShards = maxOpenShards.or(DEFAULT_MAX_OPEN_SHARDS); checkArgument(_maxOpenShards > 0, "maxOpenShards <= 0"); _openTransfers = metricRegistry.counter(MetricRegistry.name("bv.emodb.scan", "ScanUploader", "open-transfers")); _blockedNewShards = metricRegistry.counter(MetricRegistry.name("bv.emodb.scan", "ScanUploader", "blocked-new-shards")); _mapper = checkNotNull(objectMapper, "objectMapper"); } /** * Implementation-specific call to asynchronously transfer the temporary file to the given URI. */ abstract protected ListenableFuture<?> transfer(TransferKey transferKey, URI uri, File file); /** * Implementation-specific call to get TransferStatus objects for all active transfers. */ abstract protected Map<TransferKey, TransferStatus> getStatusForActiveTransfers(); @Override public ScanDestinationWriter writeShardRows(String tableName, final String placement, int shardId, long tableUuid) throws IOException, InterruptedException { final ShardFiles shardFiles = getShardFiles(shardId, tableUuid); // Fail immediately if there has already been a failure propagateExceptionIfPresent(); // Don't allow another shard to start if there are too many outstanding transfers blockNewShardIfNecessary(shardFiles); if (shardFiles.isCanceled()) { throw new IOException("Shard file closed: " + shardFiles.getKey()); } if (_closed) { throw new IOException("Writer closed"); } final File shardFile = shardFiles.addShardFile(); try { final URI uri = getUriForShard(tableName, shardId, tableUuid); OutputStream out = open(shardFile, getCounterForPlacement(placement)); return new ShardWriter(out, _mapper) { @Override synchronized protected void ready(boolean isEmpty, Optional<Integer> finalPartCount) throws IOException { if (shardFiles.isCanceled()) { throw new IOException("Shard file closed: " + shardFiles.getKey()); } if (finalPartCount.isPresent()) { shardFiles.setFinalPartCount(finalPartCount); } Optional<File> completeFile = shardFiles.shardFileComplete(isEmpty, shardFile); if (completeFile == null) { // File is not complete return; } if (completeFile.isPresent()) { // File is complete and has content asyncTransfer(completeFile.get()); } else { // File is complete but was empty so there is nothing to transfer. closeShardFiles(shardFiles); } } public void asyncTransfer(final File completeFile) { _log.debug("Initiating async transfer: id={}, file={}, uri={}", _taskId, completeFile, uri); _openTransfers.inc(); ListenableFuture<?> future = transfer(shardFiles.getKey(), uri, completeFile); Futures.addCallback(future, new FutureCallback<Object>() { @Override public void onSuccess(Object result) { transferComplete(shardFiles); } @Override public void onFailure(Throwable t) { if (!_closed) { // Frequently a spurious exception related to cancellation is thrown if this instance // has been closed. Only register the exception if the instance has not been closed. registerException(placement, t); } transferComplete(shardFiles); } }); } @Override synchronized protected void cancel() { shardFiles.setCanceled(true); closeShardFiles(shardFiles); } public String toString() { return uri.toString(); } }; } catch (IOException e) { shardFiles.deleteShardFile(shardFile); throw e; } } private ShardFiles getShardFiles(int shardId, long tableUuid) { TransferKey key = new TransferKey(tableUuid, shardId); ShardFiles shardFiles; _lock.lock(); try { shardFiles = _openShardFiles.get(key); if (shardFiles == null) { shardFiles = new ShardFiles(key); _openShardFiles.put(key, shardFiles); } } finally { _lock.unlock(); } return shardFiles; } /** Cleans up any files associated with the ShardFiles instance and removes it from then open set of ShardFiles. */ private void closeShardFiles(ShardFiles shardFiles) { shardFiles.deleteAllShardFiles(); _lock.lock(); try { _openShardFiles.remove(shardFiles.getKey()); _shardFilesClosedOrExceptionCaught.signalAll(); } finally { _lock.unlock(); } } /** Called when a shard file has been transferred, successfully or otherwise. This also closes the ShardFiles. */ private void transferComplete(ShardFiles shardFiles) { _log.debug("Transfer complete: id={}, file={}", _taskId, shardFiles.getFirstFile()); closeShardFiles(shardFiles); _openTransfers.dec(); } /** * If a previous upload has already failed then don't allow any more shards to be written. */ private void propagateExceptionIfPresent() throws IOException { if (_uploadException != null) { throw _uploadException; } } private void blockNewShardIfNecessary(ShardFiles shardFiles) throws IOException, InterruptedException { if (maxOpenShardsGreaterThan(_maxOpenShards - 1, shardFiles.getKey())) { _blockedNewShards.inc(); try { blockUntilOpenShardsAtMost(_maxOpenShards - 1, shardFiles.getKey()); } finally { _blockedNewShards.dec(); } } } /** Blocks until the number of open shards is equal to or less than the provided threshold. */ private void blockUntilOpenShardsAtMost(int maxOpenShards, @Nullable TransferKey permittedKey) throws IOException, InterruptedException { blockUntilOpenShardsAtMost(maxOpenShards, permittedKey, Instant.MAX); } /** * Blocks until the number of open shards is equal to or less than the provided threshold or the current time * is after the timeout timestamp. */ private boolean blockUntilOpenShardsAtMost(int maxOpenShards, @Nullable TransferKey permittedKey, Instant timeout) throws IOException, InterruptedException { Stopwatch stopWatch = Stopwatch.createStarted(); Instant now; _lock.lock(); try { while (!_closed && maxOpenShardsGreaterThan(maxOpenShards, permittedKey) && (now = Instant.now()).isBefore(timeout)) { // Stop blocking if there is an exception propagateExceptionIfPresent(); // Wait no longer than 30 seconds; we want to log at least every 30 seconds we've been waiting. long waitTime = Math.min(Duration.ofSeconds(30).toMillis(), Duration.between(now, timeout).toMillis()); _shardFilesClosedOrExceptionCaught.await(waitTime, TimeUnit.MILLISECONDS); if (!maxOpenShardsGreaterThan(maxOpenShards, permittedKey)) { propagateExceptionIfPresent(); return true; } _log.debug("After {} seconds task {} still has {} open shards", stopWatch.elapsed(TimeUnit.SECONDS), _taskId, _openShardFiles.size()); } propagateExceptionIfPresent(); return !maxOpenShardsGreaterThan(maxOpenShards, permittedKey); } finally { _lock.unlock(); } } private boolean maxOpenShardsGreaterThan(int maxOpenShards, @Nullable TransferKey permittedKey) { return _openShardFiles.size() > maxOpenShards && // Too many open shards (permittedKey == null || !_openShardFiles.containsKey(permittedKey)); // Never block if the permitted key's shard is already open } private void registerException(String placement, Throwable e) { _log.error("Async transfer failed for task id={}, placement {}", _taskId, placement, e); IOException ioException; if (e instanceof IOException) { ioException = (IOException) e; } else { //noinspection ThrowableInstanceNeverThrown ioException = new IOException(e); } _lock.lock(); try { if (_uploadException == null) { _uploadException = ioException; _shardFilesClosedOrExceptionCaught.signalAll(); } } finally { _lock.unlock(); } } @Override public WaitForAllTransfersCompleteResult waitForAllTransfersComplete(Duration duration) throws IOException, InterruptedException { boolean complete = blockUntilOpenShardsAtMost(0, null, Instant.now().plus(duration)); if (complete) { return new WaitForAllTransfersCompleteResult(ImmutableMap.<TransferKey, TransferStatus>of()); } return new WaitForAllTransfersCompleteResult(getStatusForActiveTransfers()); } @Override public void close() { super.close(); _lock.lock(); try { for (ShardFiles shardFiles : _openShardFiles.values()) { shardFiles.setCanceled(true); shardFiles.deleteAllShardFiles(); } _openShardFiles.clear(); _shardFilesClosedOrExceptionCaught.signalAll(); } finally { _lock.unlock(); } } private class ShardFiles { private final TransferKey _key; private final List<ShardFile> _parts = Lists.newArrayList(); private Optional<Integer> _finalPartCount = Optional.absent(); private volatile boolean _canceled; private ShardFiles(TransferKey key) { _key = key; } private TransferKey getKey() { return _key; } synchronized public void setFinalPartCount(Optional<Integer> finalPartCount) { if (finalPartCount.isPresent()) { if (!_finalPartCount.isPresent()) { _finalPartCount = finalPartCount; } else if (!_finalPartCount.get().equals(finalPartCount.get())) { throw new IllegalStateException("Shard set with inconsistent final part counts: " + _key); } } } public File addShardFile() throws IOException { File file = createTemporaryShardFile(); ShardFile shardFile = new ShardFile(file); synchronized (this) { _parts.add(shardFile); } return file; } private File createTemporaryShardFile() throws IOException { return File.createTempFile( format("emoshard_%02x_%016x", _key.getShardId(), _key.getTableUuid()), _compression.getExtension()); } /** * Called when an individual shard file has been fully written. * @param isEmpty True if no data was written to the file. This can happen if a shard contained only deleted content * @param file The file to transfer, assuming it is not empty * @return null if the shard contains at least one incomplete file, absent if the all files are complete but * were all empty, or the File to transfer if the entire shard file is complete and not empty. */ @Nullable synchronized public Optional<File> shardFileComplete(boolean isEmpty, File file) throws IOException { boolean allPartsAvailable = _finalPartCount.or(-1) == _parts.size(); boolean allComplete = true; for (ShardFile shardFile : _parts) { if (file.equals(shardFile.getFile())) { shardFile.complete(isEmpty); } allComplete = allComplete && shardFile.isComplete(); } if (allPartsAvailable && allComplete) { // All parts have been written. Reduce to the list of files which actually have content. List<File> files = FluentIterable.from(_parts) .filter(_shardFileNotEmpty) .transform(_getShardFileFile) .toList(); switch (files.size()) { case 0: // All files were empty return Optional.absent(); case 1: // There is only a single file; return it return Optional.of(files.get(0)); default: // Combine the files into a single file File combinedFile = concatenateFiles(files, createTemporaryShardFile()); // Delete the uncombined files and set the combined file as the only file deleteAllShardFiles(); _parts.clear(); _parts.add(new ShardFile(combinedFile, true, false)); return Optional.of(combinedFile); } } else { // File not ready to transfer return null; } } private void deleteAllShardFiles() { for (ShardFile shardFile : _parts) { File file = shardFile.getFile(); deleteShardFile(file); } } public void deleteShardFile(File file) { if (!file.delete()) { _log.warn("Failed to delete file: {}", file); } else { _log.debug("Deleted temporary file : {}", file); } } public File getFirstFile() { return _parts.get(0).getFile(); } public boolean isCanceled() { return _canceled; } public void setCanceled(boolean canceled) { _canceled = canceled; } } private static class ShardFile { private final File _file; private boolean _complete; private boolean _isEmpty; private ShardFile(File file) { this(file, false, false); } private ShardFile(File file, boolean complete, boolean isEmpty) { _file = file; _complete = complete; _isEmpty = isEmpty; } private File getFile() { return _file; } private boolean isComplete() { return _complete; } private void complete(boolean isEmpty) { _complete = true; _isEmpty = isEmpty; } public boolean isEmpty() { return _isEmpty; } } /** Simple predicate to filter out ShardFiles that are empty. */ private static Predicate<ShardFile> _shardFileNotEmpty = new Predicate<ShardFile>() { @Override public boolean apply(ShardFile shardFile) { return !shardFile.isEmpty(); } }; /** Simple function to return the local temporary File associated with a ShardFile. */ private static Function<ShardFile, File> _getShardFileFile = new Function<ShardFile, File>() { @Override public File apply(ShardFile shardFile) { return shardFile.getFile(); } }; }