pw_transfer/java/main/dev/pigweed/pw_transfer/TransferEventHandler.java - pigweed/pigweed - Git at Google

 // Copyright 2022 The Pigweed Authors
 //
 // Licensed 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
 //
 //     https://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 dev.pigweed.pw_transfer;

 import com.google.common.util.concurrent.ListenableFuture;
 import dev.pigweed.pw_log.Logger;
 import dev.pigweed.pw_rpc.Call;
 import dev.pigweed.pw_rpc.ChannelOutputException;
 import dev.pigweed.pw_rpc.MethodClient;
 import dev.pigweed.pw_rpc.RpcError;
 import dev.pigweed.pw_rpc.Status;
 import dev.pigweed.pw_rpc.StreamObserver;
 import java.time.Instant;
 import java.time.temporal.ChronoUnit;
 import java.time.temporal.TemporalUnit;
 import java.util.ArrayList;
 import java.util.Comparator;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.Semaphore;
 import java.util.concurrent.TimeUnit;
 import java.util.function.BooleanSupplier;
 import java.util.function.Consumer;
 import javax.annotation.Nullable;

 /** Manages the active transfers and dispatches events to them. */
 class TransferEventHandler {
   private static final Logger logger = Logger.forClass(TransferEventHandler.class);

   // Instant and BlockingQueue use different time unit types.
   private static final TemporalUnit TIME_UNIT = ChronoUnit.MICROS;
   private static final TimeUnit POLL_TIME_UNIT = TimeUnit.MICROSECONDS;

   private final MethodClient readMethod;
   private final MethodClient writeMethod;

   private final BlockingQueue<Event> events = new LinkedBlockingQueue<>();

   // Maps resource ID to transfer
   private final Map<Integer, Transfer<?>> transfers = new HashMap<>();

   @Nullable private Call.ClientStreaming<Chunk> readStream = null;
   @Nullable private Call.ClientStreaming<Chunk> writeStream = null;
   private boolean processEvents = true;

   TransferEventHandler(MethodClient readMethod, MethodClient writeMethod) {
     this.readMethod = readMethod;
     this.writeMethod = writeMethod;
   }

   ListenableFuture<Void> startWriteTransferAsClient(int resourceId,
       int transferTimeoutMillis,
       int initialTransferTimeoutMillis,
       int maxRetries,
       byte[] data,
       Consumer<TransferProgress> progressCallback,
       BooleanSupplier shouldAbortCallback) {
     WriteTransfer transfer = new WriteTransfer(resourceId,
         new TransferInterface() {
           @Override
           Call.ClientStreaming<Chunk> getStream() throws ChannelOutputException {
             if (writeStream == null) {
               writeStream = writeMethod.invokeBidirectionalStreaming(new ChunkHandler() {
                 @Override
                 void resetStream() {
                   writeStream = null;
                 }
               });
             }
             return writeStream;
           }
         },
         transferTimeoutMillis,
         initialTransferTimeoutMillis,
         maxRetries,
         data,
         progressCallback,
         shouldAbortCallback);
     startTransferAsClient(transfer);
     return transfer.getFuture();
   }

   ListenableFuture<byte[]> startReadTransferAsClient(int resourceId,
       int transferTimeoutMillis,
       int initialTransferTimeoutMillis,
       int maxRetries,
       TransferParameters parameters,
       Consumer<TransferProgress> progressCallback,
       BooleanSupplier shouldAbortCallback) {
     ReadTransfer transfer = new ReadTransfer(resourceId,
         new TransferInterface() {
           @Override
           Call.ClientStreaming<Chunk> getStream() throws ChannelOutputException {
             if (readStream == null) {
               readStream = readMethod.invokeBidirectionalStreaming(new ChunkHandler() {
                 @Override
                 void resetStream() {
                   readStream = null;
                 }
               });
             }
             return readStream;
           }
         },
         transferTimeoutMillis,
         initialTransferTimeoutMillis,
         maxRetries,
         parameters,
         progressCallback,
         shouldAbortCallback);
     startTransferAsClient(transfer);
     return transfer.getFuture();
   }

   private void startTransferAsClient(Transfer<?> transfer) {
     enqueueEvent(() -> {
       if (transfers.put(transfer.getResourceId(), transfer) != null) {
         transfer.cleanUp(new TransferError("A transfer for resource ID " + transfer.getResourceId()
                 + " is already in progress! Only one read/write transfer per resource is supported at a time",
             Status.ALREADY_EXISTS));
         return;
       }
       transfer.start();
     });
   }

   /** Handles events until stop() is called. */
   void run() {
     while (processEvents) {
       handleNextEvent();
     }
   }

   /** Stops the transfer event handler from processing events. */
   void stop() {
     enqueueEvent(() -> {
       logger.atFine().log("Terminating TransferEventHandler");
       transfers.values().forEach(Transfer::handleTermination);
       processEvents = false;
     });
   }

   /** Blocks until all events currently in the queue are processed; for test use only. */
   void waitUntilEventsAreProcessedForTest() {
     Semaphore semaphore = new Semaphore(0);
     enqueueEvent(semaphore::release);
     try {
       semaphore.acquire();
     } catch (InterruptedException e) {
       throw new AssertionError("Unexpectedly interrupted", e);
     }
   }

   private void enqueueEvent(Event event) {
     while (true) {
       try {
         events.put(event);
         return;
       } catch (InterruptedException e) {
         // Ignore and keep trying.
       }
     }
   }

   private void handleNextEvent() {
     final long sleepFor = Math.min(0, Instant.now().until(getNextTimeout(), TIME_UNIT));
     try {
       Event event = events.poll(sleepFor, POLL_TIME_UNIT);
       if (event != null) {
         event.handle();
       }
     } catch (InterruptedException e) {
       // If interrupted, check for timeouts anyway.
     }

     for (Transfer<?> transfer : transfers.values()) {
       transfer.handleTimeoutIfDeadlineExceeded();
     }
   }

   private Instant getNextTimeout() {
     Optional<Transfer<?>> transfer =
         transfers.values().stream().min(Comparator.comparing(Transfer::getDeadline));
     return transfer.isPresent() ? transfer.get().getDeadline() : Transfer.NO_TIMEOUT;
   }

   /** This interface gives a Transfer access to the TransferEventHandler. */
   abstract class TransferInterface {
     private TransferInterface() {}

     /**
      *  Sends the provided transfer chunk.
      *
      *  Must be called on the transfer therad.
      */
     void sendChunk(Chunk chunk) throws TransferError {
       try {
         getStream().send(chunk);
       } catch (ChannelOutputException | RpcError e) {
         throw new TransferError("Failed to send chunk for write transfer", e);
       }
     }

     /**
      *  Removes this transfer from the list of active transfers.
      *
      *  Must be called on the transfer therad.
      */
     void unregisterTransfer(int sessionId) {
       transfers.remove(sessionId);
     }

     /**
      * Initiates the cancellation process for the provided transfer.
      *
      * May be called from any thread.
      */
     void cancelTransfer(Transfer<?> transfer) {
       enqueueEvent(transfer::handleCancellation);
     }

     /** Gets either the read or write stream. */
     abstract Call.ClientStreaming<Chunk> getStream() throws ChannelOutputException;
   }

   /** Handles responses on the pw_transfer RPCs. */
   private abstract class ChunkHandler implements StreamObserver<Chunk> {
     @Override
     public final void onNext(Chunk chunk) {
       enqueueEvent(() -> {
         Transfer<?> transfer = transfers.get(chunk.getTransferId());
         if (transfer != null) {
           logger.atFinest().log(
               "Transfer %d received chunk: %s", transfer.getSessionId(), chunkToString(chunk));
           transfer.handleChunk(chunk);
         } else {
           logger.atWarning().log(
               "Ignoring unrecognized transfer session ID %d", chunk.getTransferId());
         }
       });
     }

     @Override
     public final void onCompleted(Status status) {
       onError(Status.INTERNAL); // This RPC should never complete: treat as an internal error.
     }

     @Override
     public final void onError(Status status) {
       enqueueEvent(() -> {
         resetStream();

         // The transfers remove themselves from the Map during cleanup, iterate over a copied list.
         List<Transfer<?>> activeTransfers = new ArrayList<>(transfers.values());

         // FAILED_PRECONDITION indicates that the stream packet was not recognized as the stream is
         // not open. This could occur if the server resets. Notify pending transfers that this has
         // occurred so they can restart.
         if (status.equals(Status.FAILED_PRECONDITION)) {
           activeTransfers.forEach(Transfer::handleDisconnection);
         } else {
           TransferError error = new TransferError(
               "Transfer stream RPC closed unexpectedly with status " + status, Status.INTERNAL);
           activeTransfers.forEach(t -> t.cleanUp(error));
         }
       });
     }

     abstract void resetStream();
   }

   private static String chunkToString(Chunk chunk) {
     StringBuilder str = new StringBuilder();
     str.append("transferId:").append(chunk.getTransferId()).append(" ");
     str.append("windowEndOffset:").append(chunk.getWindowEndOffset()).append(" ");
     str.append("offset:").append(chunk.getOffset()).append(" ");
     // Don't include the actual data; it's too much.
     str.append("len(data):").append(chunk.getData().size()).append(" ");
     if (chunk.hasPendingBytes()) {
       str.append("pendingBytes:").append(chunk.getPendingBytes()).append(" ");
     }
     if (chunk.hasMaxChunkSizeBytes()) {
       str.append("maxChunkSizeBytes:").append(chunk.getMaxChunkSizeBytes()).append(" ");
     }
     if (chunk.hasMinDelayMicroseconds()) {
       str.append("minDelayMicroseconds:").append(chunk.getMinDelayMicroseconds()).append(" ");
     }
     if (chunk.hasRemainingBytes()) {
       str.append("remainingBytes:").append(chunk.getRemainingBytes()).append(" ");
     }
     if (chunk.hasStatus()) {
       str.append("status:").append(chunk.getStatus()).append(" ");
     }
     if (chunk.hasType()) {
       str.append("type:").append(chunk.getTypeValue()).append(" ");
     }
     if (chunk.hasResourceId()) {
       str.append("resourceId:").append(chunk.getSessionId()).append(" ");
     }
     if (chunk.hasSessionId()) {
       str.append("sessionId:").append(chunk.getSessionId()).append(" ");
     }
     return str.toString();
   }

   // Represents an event that occurs during a transfer
   private interface Event {
     void handle();
   }
 }
	// Copyright 2022 The Pigweed Authors
	//
	// Licensed 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
	//
	// https://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 dev.pigweed.pw_transfer;

	import com.google.common.util.concurrent.ListenableFuture;
	import dev.pigweed.pw_log.Logger;
	import dev.pigweed.pw_rpc.Call;
	import dev.pigweed.pw_rpc.ChannelOutputException;
	import dev.pigweed.pw_rpc.MethodClient;
	import dev.pigweed.pw_rpc.RpcError;
	import dev.pigweed.pw_rpc.Status;
	import dev.pigweed.pw_rpc.StreamObserver;
	import java.time.Instant;
	import java.time.temporal.ChronoUnit;
	import java.time.temporal.TemporalUnit;
	import java.util.ArrayList;
	import java.util.Comparator;
	import java.util.HashMap;
	import java.util.List;
	import java.util.Map;
	import java.util.Optional;
	import java.util.concurrent.BlockingQueue;
	import java.util.concurrent.LinkedBlockingQueue;
	import java.util.concurrent.Semaphore;
	import java.util.concurrent.TimeUnit;
	import java.util.function.BooleanSupplier;
	import java.util.function.Consumer;
	import javax.annotation.Nullable;

	/** Manages the active transfers and dispatches events to them. */
	class TransferEventHandler {
	private static final Logger logger = Logger.forClass(TransferEventHandler.class);

	// Instant and BlockingQueue use different time unit types.
	private static final TemporalUnit TIME_UNIT = ChronoUnit.MICROS;
	private static final TimeUnit POLL_TIME_UNIT = TimeUnit.MICROSECONDS;

	private final MethodClient readMethod;
	private final MethodClient writeMethod;

	private final BlockingQueue<Event> events = new LinkedBlockingQueue<>();

	// Maps resource ID to transfer
	private final Map<Integer, Transfer<?>> transfers = new HashMap<>();

	@Nullable private Call.ClientStreaming<Chunk> readStream = null;
	@Nullable private Call.ClientStreaming<Chunk> writeStream = null;
	private boolean processEvents = true;

	TransferEventHandler(MethodClient readMethod, MethodClient writeMethod) {
	this.readMethod = readMethod;
	this.writeMethod = writeMethod;
	}

	ListenableFuture<Void> startWriteTransferAsClient(int resourceId,
	int transferTimeoutMillis,
	int initialTransferTimeoutMillis,
	int maxRetries,
	byte[] data,
	Consumer<TransferProgress> progressCallback,
	BooleanSupplier shouldAbortCallback) {
	WriteTransfer transfer = new WriteTransfer(resourceId,
	new TransferInterface() {
	@Override
	Call.ClientStreaming<Chunk> getStream() throws ChannelOutputException {
	if (writeStream == null) {
	writeStream = writeMethod.invokeBidirectionalStreaming(new ChunkHandler() {
	@Override
	void resetStream() {
	writeStream = null;
	}
	});
	}
	return writeStream;
	}
	},
	transferTimeoutMillis,
	initialTransferTimeoutMillis,
	maxRetries,
	data,
	progressCallback,
	shouldAbortCallback);
	startTransferAsClient(transfer);
	return transfer.getFuture();
	}

	ListenableFuture<byte[]> startReadTransferAsClient(int resourceId,
	int transferTimeoutMillis,
	int initialTransferTimeoutMillis,
	int maxRetries,
	TransferParameters parameters,
	Consumer<TransferProgress> progressCallback,
	BooleanSupplier shouldAbortCallback) {
	ReadTransfer transfer = new ReadTransfer(resourceId,
	new TransferInterface() {
	@Override
	Call.ClientStreaming<Chunk> getStream() throws ChannelOutputException {
	if (readStream == null) {
	readStream = readMethod.invokeBidirectionalStreaming(new ChunkHandler() {
	@Override
	void resetStream() {
	readStream = null;
	}
	});
	}
	return readStream;
	}
	},
	transferTimeoutMillis,
	initialTransferTimeoutMillis,
	maxRetries,
	parameters,
	progressCallback,
	shouldAbortCallback);
	startTransferAsClient(transfer);
	return transfer.getFuture();
	}

	private void startTransferAsClient(Transfer<?> transfer) {
	enqueueEvent(() -> {
	if (transfers.put(transfer.getResourceId(), transfer) != null) {
	transfer.cleanUp(new TransferError("A transfer for resource ID " + transfer.getResourceId()
	+ " is already in progress! Only one read/write transfer per resource is supported at a time",
	Status.ALREADY_EXISTS));
	return;
	}
	transfer.start();
	});
	}

	/** Handles events until stop() is called. */
	void run() {
	while (processEvents) {
	handleNextEvent();
	}
	}

	/** Stops the transfer event handler from processing events. */
	void stop() {
	enqueueEvent(() -> {
	logger.atFine().log("Terminating TransferEventHandler");
	transfers.values().forEach(Transfer::handleTermination);
	processEvents = false;
	});
	}

	/** Blocks until all events currently in the queue are processed; for test use only. */
	void waitUntilEventsAreProcessedForTest() {
	Semaphore semaphore = new Semaphore(0);
	enqueueEvent(semaphore::release);
	try {
	semaphore.acquire();
	} catch (InterruptedException e) {
	throw new AssertionError("Unexpectedly interrupted", e);
	}
	}

	private void enqueueEvent(Event event) {
	while (true) {
	try {
	events.put(event);
	return;
	} catch (InterruptedException e) {
	// Ignore and keep trying.
	}
	}
	}

	private void handleNextEvent() {
	final long sleepFor = Math.min(0, Instant.now().until(getNextTimeout(), TIME_UNIT));
	try {
	Event event = events.poll(sleepFor, POLL_TIME_UNIT);
	if (event != null) {
	event.handle();
	}
	} catch (InterruptedException e) {
	// If interrupted, check for timeouts anyway.
	}

	for (Transfer<?> transfer : transfers.values()) {
	transfer.handleTimeoutIfDeadlineExceeded();
	}
	}

	private Instant getNextTimeout() {
	Optional<Transfer<?>> transfer =
	transfers.values().stream().min(Comparator.comparing(Transfer::getDeadline));
	return transfer.isPresent() ? transfer.get().getDeadline() : Transfer.NO_TIMEOUT;
	}

	/** This interface gives a Transfer access to the TransferEventHandler. */
	abstract class TransferInterface {
	private TransferInterface() {}

	/**
	* Sends the provided transfer chunk.
	*
	* Must be called on the transfer therad.
	*/
	void sendChunk(Chunk chunk) throws TransferError {
	try {
	getStream().send(chunk);
	} catch (ChannelOutputException \| RpcError e) {
	throw new TransferError("Failed to send chunk for write transfer", e);
	}
	}

	/**
	* Removes this transfer from the list of active transfers.
	*
	* Must be called on the transfer therad.
	*/
	void unregisterTransfer(int sessionId) {
	transfers.remove(sessionId);
	}

	/**
	* Initiates the cancellation process for the provided transfer.
	*
	* May be called from any thread.
	*/
	void cancelTransfer(Transfer<?> transfer) {
	enqueueEvent(transfer::handleCancellation);
	}

	/** Gets either the read or write stream. */
	abstract Call.ClientStreaming<Chunk> getStream() throws ChannelOutputException;
	}

	/** Handles responses on the pw_transfer RPCs. */
	private abstract class ChunkHandler implements StreamObserver<Chunk> {
	@Override
	public final void onNext(Chunk chunk) {
	enqueueEvent(() -> {
	Transfer<?> transfer = transfers.get(chunk.getTransferId());
	if (transfer != null) {
	logger.atFinest().log(
	"Transfer %d received chunk: %s", transfer.getSessionId(), chunkToString(chunk));
	transfer.handleChunk(chunk);
	} else {
	logger.atWarning().log(
	"Ignoring unrecognized transfer session ID %d", chunk.getTransferId());
	}
	});
	}

	@Override
	public final void onCompleted(Status status) {
	onError(Status.INTERNAL); // This RPC should never complete: treat as an internal error.
	}

	@Override
	public final void onError(Status status) {
	enqueueEvent(() -> {
	resetStream();

	// The transfers remove themselves from the Map during cleanup, iterate over a copied list.
	List<Transfer<?>> activeTransfers = new ArrayList<>(transfers.values());

	// FAILED_PRECONDITION indicates that the stream packet was not recognized as the stream is
	// not open. This could occur if the server resets. Notify pending transfers that this has
	// occurred so they can restart.
	if (status.equals(Status.FAILED_PRECONDITION)) {
	activeTransfers.forEach(Transfer::handleDisconnection);
	} else {
	TransferError error = new TransferError(
	"Transfer stream RPC closed unexpectedly with status " + status, Status.INTERNAL);
	activeTransfers.forEach(t -> t.cleanUp(error));
	}
	});
	}

	abstract void resetStream();
	}

	private static String chunkToString(Chunk chunk) {
	StringBuilder str = new StringBuilder();
	str.append("transferId:").append(chunk.getTransferId()).append(" ");
	str.append("windowEndOffset:").append(chunk.getWindowEndOffset()).append(" ");
	str.append("offset:").append(chunk.getOffset()).append(" ");
	// Don't include the actual data; it's too much.
	str.append("len(data):").append(chunk.getData().size()).append(" ");
	if (chunk.hasPendingBytes()) {
	str.append("pendingBytes:").append(chunk.getPendingBytes()).append(" ");
	}
	if (chunk.hasMaxChunkSizeBytes()) {
	str.append("maxChunkSizeBytes:").append(chunk.getMaxChunkSizeBytes()).append(" ");
	}
	if (chunk.hasMinDelayMicroseconds()) {
	str.append("minDelayMicroseconds:").append(chunk.getMinDelayMicroseconds()).append(" ");
	}
	if (chunk.hasRemainingBytes()) {
	str.append("remainingBytes:").append(chunk.getRemainingBytes()).append(" ");
	}
	if (chunk.hasStatus()) {
	str.append("status:").append(chunk.getStatus()).append(" ");
	}
	if (chunk.hasType()) {
	str.append("type:").append(chunk.getTypeValue()).append(" ");
	}
	if (chunk.hasResourceId()) {
	str.append("resourceId:").append(chunk.getSessionId()).append(" ");
	}
	if (chunk.hasSessionId()) {
	str.append("sessionId:").append(chunk.getSessionId()).append(" ");
	}
	return str.toString();
	}

	// Represents an event that occurs during a transfer
	private interface Event {
	void handle();
	}
	}