pw_transfer/java/main/dev/pigweed/pw_transfer/Transfer.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 static com.google.common.util.concurrent.MoreExecutors.directExecutor;
 import static dev.pigweed.pw_transfer.TransferProgress.UNKNOWN_TRANSFER_SIZE;

 import com.google.common.util.concurrent.SettableFuture;
 import dev.pigweed.pw_log.Logger;
 import dev.pigweed.pw_rpc.Status;
 import java.time.Duration;
 import java.time.Instant;
 import java.util.function.BooleanSupplier;
 import java.util.function.Consumer;

 /** Base class for tracking the state of a read or write transfer. */
 abstract class Transfer<T> {
   private static final Logger logger = Logger.forClass(Transfer.class);

   // Largest nanosecond instant. Used to block indefinitely when no transfers are pending.
   static final Instant NO_TIMEOUT = Instant.ofEpochSecond(0, Long.MAX_VALUE);

   private final int resourceId;
   private final int sessionId;
   private final TransferEventHandler.TransferInterface eventHandler;
   private final SettableFuture<T> future;
   private final int maxRetries;
   private final int timeoutMillis;
   private final int initialTimeoutMillis;
   private final Consumer<TransferProgress> progressCallback;
   private final BooleanSupplier shouldAbortCallback;
   private final Instant startTime;

   private State state = new Inactive();
   private Instant deadline = NO_TIMEOUT;
   private boolean isCleanedUp = false;

   // The number of times this transfer has retried due to an RPC disconnection. Limit this to
   // maxRetries to prevent repeated crashes if reading to / writing from a particular transfer is
   // causing crashes.
   private int disconnectionRetries = 0;

   /**
    * Creates a new read or write transfer.
    * @param resourceId The resource ID of the transfer
    * @param eventHandler Interface to use to send a chunk.
    * @param timeoutMillis Maximum time to wait for a chunk from the server.
    * @param initialTimeoutMillis Maximum time to wait for the first chunk from the server.
    * @param maxRetries Number of times to retry due to a timeout or RPC error.
    * @param progressCallback Called each time a packet is sent.
    * @param shouldAbortCallback BooleanSupplier that returns true if a transfer should be aborted.
    */
   Transfer(int resourceId,
       TransferEventHandler.TransferInterface eventHandler,
       int timeoutMillis,
       int initialTimeoutMillis,
       int maxRetries,
       Consumer<TransferProgress> progressCallback,
       BooleanSupplier shouldAbortCallback) {
     this.resourceId = resourceId;
     this.sessionId = resourceId; // TODO(b/227211777): Distinguish between resource & session ID
     this.eventHandler = eventHandler;

     this.future = SettableFuture.create();
     this.timeoutMillis = timeoutMillis;
     this.initialTimeoutMillis = initialTimeoutMillis;
     this.maxRetries = maxRetries;
     this.progressCallback = progressCallback;
     this.shouldAbortCallback = shouldAbortCallback;

     // If the future is cancelled, tell the TransferEventHandler to cancel the transfer.
     future.addListener(() -> {
       if (future.isCancelled()) {
         eventHandler.cancelTransfer(this);
       }
     }, directExecutor());

     startTime = Instant.now();
   }

   final int getResourceId() {
     return resourceId;
   }

   final int getSessionId() {
     return sessionId;
   }

   final Instant getDeadline() {
     return deadline;
   }

   final void setNextChunkTimeout() {
     deadline = Instant.now().plusMillis(timeoutMillis);
   }

   private void setInitialTimeout() {
     deadline = Instant.now().plusMillis(initialTimeoutMillis);
   }

   final void setTimeoutMicros(int timeoutMicros) {
     deadline = Instant.now().plusNanos((long) timeoutMicros * 1000);
   }

   final SettableFuture<T> getFuture() {
     return future;
   }

   final void start() {
     logger.atFine().log(
         "Transfer %d for resource %d starting with parameters: default timeout %d ms, initial timeout %d ms, %d max retires",
         getSessionId(),
         getResourceId(),
         timeoutMillis,
         initialTimeoutMillis,
         maxRetries);
     if (!sendChunk(getChunkForRetry())) {
       return; // Sending failed, transfer is cancelled
     }

     setState(getWaitingForDataState());
     setInitialTimeout();
   }

   /** Processes an incoming chunk from the server. */
   final void handleChunk(Chunk chunk) {
     // Since a packet has been received, don't allow retries on disconnection; abort instead.
     disconnectionRetries = Integer.MAX_VALUE;

     if (chunk.hasStatus()) {
       logger.atFinest().log("Event: handle final chunk");
       state.handleFinalChunk(chunk);
     } else {
       logger.atFinest().log("Event: handle data chunk");
       state.handleDataChunk(chunk);
     }
   }

   final void handleTimeoutIfDeadlineExceeded() {
     if (Instant.now().isAfter(deadline)) {
       logger.atFinest().log("Event: handleTimeout since %s is after %s", Instant.now(), deadline);
       state.handleTimeout();
     }
   }

   final void handleTermination() {
     state.handleTermination();
   }

   final void handleCancellation() {
     state.handleCancellation();
   }

   /** Restarts a transfer after an RPC disconnection. */
   final void handleDisconnection() {
     // disconnectionRetries is set to Int.MAX_VALUE when a packet is received to prevent retries
     // after the initial packet.
     if (disconnectionRetries++ < maxRetries) {
       logger.atFine().log("Restarting the pw_transfer RPC for transfer %d (attempt %d/%d)",
           sessionId,
           disconnectionRetries,
           maxRetries);
       if (sendChunk(getChunkForRetry())) {
         setInitialTimeout();
       }
     } else {
       cleanUp(new TransferError(
           "Transfer " + sessionId + " restarted " + maxRetries + " times, aborting",
           Status.INTERNAL));
     }
   }

   /** Returns the State to enter immediately after sending the first packet. */
   abstract State getWaitingForDataState();

   /**
    * Returns the chunk to send for a retry. Returns the initial chunk if no chunks have been sent.
    */
   abstract Chunk getChunkForRetry();

   /** Sets the result for the future after a successful transfer. */
   abstract void setFutureResult();

   final Chunk.Builder newChunk(Chunk.Type type) {
     // TODO(frolv): Properly set the session ID after it is configured by the server.
     return Chunk.newBuilder().setTransferId(getSessionId()).setType(type);
   }

   final void setState(State newState) {
     if (newState != state) {
       logger.atFinest().log(
           "Transfer %d state %s -> %s", getSessionId(), state.getName(), newState.getName());
     }
     state = newState;
   }

   /** Sends a chunk. Returns true if sent, false if sending failed and the transfer was aborted. */
   final boolean sendChunk(Chunk chunk) {
     if (shouldAbortCallback.getAsBoolean()) {
       logger.atWarning().log("Abort signal received.");
       cleanUp(new TransferError(sessionId, Status.ABORTED));
       return false;
     }

     try {
       eventHandler.sendChunk(chunk);
       return true;
     } catch (TransferError transferError) {
       cleanUp(transferError);
       return false;
     }
   }

   final boolean sendChunk(Chunk.Builder chunk) {
     return sendChunk(chunk.build());
   }

   /** Performs final cleanup of a completed transfer. No packets are sent to the server. */
   private void cleanUp(Status status) {
     eventHandler.unregisterTransfer(sessionId);

     logger.atInfo().log("Transfer %d completed with status %s", sessionId, status);
     if (status.ok()) {
       setFutureResult();
     } else {
       future.setException(new TransferError(sessionId, status));
     }
     isCleanedUp = true;
   }

   /** Finishes the transfer due to an exception. No packets are sent to the server. */
   final void cleanUp(TransferError exception) {
     eventHandler.unregisterTransfer(sessionId);

     logger.atWarning().withCause(exception).log("Transfer %d terminated with exception", sessionId);
     future.setException(exception);
     isCleanedUp = true;
   }

   /** Sends a status chunk to the server and finishes the transfer. */
   final void sendFinalChunk(Status status) {
     logger.atFine().log("Sending final chunk for transfer %d with status %s", sessionId, status);
     // If the transfer was completed due to concurrently handling chunks, don't send.
     if (isCleanedUp) {
       logger.atFine().log("Skipping sending final chunk on already-completed transfer");
       return;
     }

     // Only call finish() if the sendChunk was successful. If it wasn't, the exception would have
     // already terminated the transfer.
     if (sendChunk(newChunk(Chunk.Type.COMPLETION).setStatus(status.code()))) {
       cleanUp(status);
     }
   }

   /** Invokes the transfer progress callback and logs the progress. */
   final void updateProgress(long bytesSent, long bytesConfirmedReceived, long totalSizeBytes) {
     TransferProgress progress =
         TransferProgress.create(bytesSent, bytesConfirmedReceived, totalSizeBytes);
     progressCallback.accept(progress);

     long durationNanos = Duration.between(startTime, Instant.now()).toNanos();
     long totalRate = durationNanos == 0 ? 0 : (bytesSent * 1_000_000_000 / durationNanos);

     logger.atFine().log("Transfer %d progress: "
             + "%5.1f%% (%d B sent, %d B confirmed received of %s B total) at %d B/s",
         sessionId,
         progress.percentReceived(),
         bytesSent,
         bytesConfirmedReceived,
         totalSizeBytes == UNKNOWN_TRANSFER_SIZE ? "unknown" : totalSizeBytes,
         totalRate);
   }

   /** Represents a state in the transfer state machine. */
   abstract class State {
     private String getName() {
       return getClass().getSimpleName();
     }

     /** Called to handle a non-final chunk for this transfer. */
     void handleDataChunk(Chunk chunk) {
       logger.atFine().log(
           "Transfer %d [%s state]: Received unexpected data chunk", getSessionId(), getName());
     }

     /** Called to handle the final chunk for this transfer. */
     void handleFinalChunk(Chunk chunk) {
       Status status = Status.fromCode(chunk.getStatus());
       if (status != null) {
         cleanUp(status);
       } else {
         logger.atWarning().log("Received invalid status value %d", chunk.getStatus());
         sendFinalChunk(Status.INVALID_ARGUMENT);
       }
       setState(new Completed());
     }

     /** Called when this transfer's deadline expires. */
     void handleTimeout() {
       logger.atFine().log("Transfer %d [%s state]: Ignoring timeout", getSessionId(), getName());
     }

     /** Called if the transfer is cancelled by the user. */
     void handleCancellation() {
       sendFinalChunk(Status.CANCELLED);
       setState(new Completed());
     }

     /** Called when the transfer thread is shutting down. */
     void handleTermination() {
       sendFinalChunk(Status.ABORTED);
       setState(new Completed());
     }
   }

   private class Inactive extends State {}

   /** Recovering from an expired timeout. */
   class Recovery extends State {
     private int retries;

     @Override
     void handleDataChunk(Chunk chunk) {
       setState(getWaitingForDataState());
       state.handleDataChunk(chunk);
     }

     @Override
     void handleTimeout() {
       if (retries < maxRetries) {
         logger.atFiner().log("Transfer %d received no chunks for %d ms; retrying %d/%d",
             getSessionId(),
             timeoutMillis,
             retries,
             maxRetries);
         sendChunk(getChunkForRetry());
         retries += 1;
         setNextChunkTimeout();
         return;
       }

       sendFinalChunk(Status.DEADLINE_EXCEEDED);
       setState(new Completed());
     }
   }

   /** Transfer completed. Do nothing if the transfer is terminated or cancelled. */
   class Completed extends State {
     Completed() {
       deadline = NO_TIMEOUT;
     }

     @Override
     void handleTermination() {}

     @Override
     void handleCancellation() {}
   }
 }
	// 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 static com.google.common.util.concurrent.MoreExecutors.directExecutor;
	import static dev.pigweed.pw_transfer.TransferProgress.UNKNOWN_TRANSFER_SIZE;

	import com.google.common.util.concurrent.SettableFuture;
	import dev.pigweed.pw_log.Logger;
	import dev.pigweed.pw_rpc.Status;
	import java.time.Duration;
	import java.time.Instant;
	import java.util.function.BooleanSupplier;
	import java.util.function.Consumer;

	/** Base class for tracking the state of a read or write transfer. */
	abstract class Transfer<T> {
	private static final Logger logger = Logger.forClass(Transfer.class);

	// Largest nanosecond instant. Used to block indefinitely when no transfers are pending.
	static final Instant NO_TIMEOUT = Instant.ofEpochSecond(0, Long.MAX_VALUE);

	private final int resourceId;
	private final int sessionId;
	private final TransferEventHandler.TransferInterface eventHandler;
	private final SettableFuture<T> future;
	private final int maxRetries;
	private final int timeoutMillis;
	private final int initialTimeoutMillis;
	private final Consumer<TransferProgress> progressCallback;
	private final BooleanSupplier shouldAbortCallback;
	private final Instant startTime;

	private State state = new Inactive();
	private Instant deadline = NO_TIMEOUT;
	private boolean isCleanedUp = false;

	// The number of times this transfer has retried due to an RPC disconnection. Limit this to
	// maxRetries to prevent repeated crashes if reading to / writing from a particular transfer is
	// causing crashes.
	private int disconnectionRetries = 0;

	/**
	* Creates a new read or write transfer.
	* @param resourceId The resource ID of the transfer
	* @param eventHandler Interface to use to send a chunk.
	* @param timeoutMillis Maximum time to wait for a chunk from the server.
	* @param initialTimeoutMillis Maximum time to wait for the first chunk from the server.
	* @param maxRetries Number of times to retry due to a timeout or RPC error.
	* @param progressCallback Called each time a packet is sent.
	* @param shouldAbortCallback BooleanSupplier that returns true if a transfer should be aborted.
	*/
	Transfer(int resourceId,
	TransferEventHandler.TransferInterface eventHandler,
	int timeoutMillis,
	int initialTimeoutMillis,
	int maxRetries,
	Consumer<TransferProgress> progressCallback,
	BooleanSupplier shouldAbortCallback) {
	this.resourceId = resourceId;
	this.sessionId = resourceId; // TODO(b/227211777): Distinguish between resource & session ID
	this.eventHandler = eventHandler;

	this.future = SettableFuture.create();
	this.timeoutMillis = timeoutMillis;
	this.initialTimeoutMillis = initialTimeoutMillis;
	this.maxRetries = maxRetries;
	this.progressCallback = progressCallback;
	this.shouldAbortCallback = shouldAbortCallback;

	// If the future is cancelled, tell the TransferEventHandler to cancel the transfer.
	future.addListener(() -> {
	if (future.isCancelled()) {
	eventHandler.cancelTransfer(this);
	}
	}, directExecutor());

	startTime = Instant.now();
	}

	final int getResourceId() {
	return resourceId;
	}

	final int getSessionId() {
	return sessionId;
	}

	final Instant getDeadline() {
	return deadline;
	}

	final void setNextChunkTimeout() {
	deadline = Instant.now().plusMillis(timeoutMillis);
	}

	private void setInitialTimeout() {
	deadline = Instant.now().plusMillis(initialTimeoutMillis);
	}

	final void setTimeoutMicros(int timeoutMicros) {
	deadline = Instant.now().plusNanos((long) timeoutMicros * 1000);
	}

	final SettableFuture<T> getFuture() {
	return future;
	}

	final void start() {
	logger.atFine().log(
	"Transfer %d for resource %d starting with parameters: default timeout %d ms, initial timeout %d ms, %d max retires",
	getSessionId(),
	getResourceId(),
	timeoutMillis,
	initialTimeoutMillis,
	maxRetries);
	if (!sendChunk(getChunkForRetry())) {
	return; // Sending failed, transfer is cancelled
	}

	setState(getWaitingForDataState());
	setInitialTimeout();
	}

	/** Processes an incoming chunk from the server. */
	final void handleChunk(Chunk chunk) {
	// Since a packet has been received, don't allow retries on disconnection; abort instead.
	disconnectionRetries = Integer.MAX_VALUE;

	if (chunk.hasStatus()) {
	logger.atFinest().log("Event: handle final chunk");
	state.handleFinalChunk(chunk);
	} else {
	logger.atFinest().log("Event: handle data chunk");
	state.handleDataChunk(chunk);
	}
	}

	final void handleTimeoutIfDeadlineExceeded() {
	if (Instant.now().isAfter(deadline)) {
	logger.atFinest().log("Event: handleTimeout since %s is after %s", Instant.now(), deadline);
	state.handleTimeout();
	}
	}

	final void handleTermination() {
	state.handleTermination();
	}

	final void handleCancellation() {
	state.handleCancellation();
	}

	/** Restarts a transfer after an RPC disconnection. */
	final void handleDisconnection() {
	// disconnectionRetries is set to Int.MAX_VALUE when a packet is received to prevent retries
	// after the initial packet.
	if (disconnectionRetries++ < maxRetries) {
	logger.atFine().log("Restarting the pw_transfer RPC for transfer %d (attempt %d/%d)",
	sessionId,
	disconnectionRetries,
	maxRetries);
	if (sendChunk(getChunkForRetry())) {
	setInitialTimeout();
	}
	} else {
	cleanUp(new TransferError(
	"Transfer " + sessionId + " restarted " + maxRetries + " times, aborting",
	Status.INTERNAL));
	}
	}

	/** Returns the State to enter immediately after sending the first packet. */
	abstract State getWaitingForDataState();

	/**
	* Returns the chunk to send for a retry. Returns the initial chunk if no chunks have been sent.
	*/
	abstract Chunk getChunkForRetry();

	/** Sets the result for the future after a successful transfer. */
	abstract void setFutureResult();

	final Chunk.Builder newChunk(Chunk.Type type) {
	// TODO(frolv): Properly set the session ID after it is configured by the server.
	return Chunk.newBuilder().setTransferId(getSessionId()).setType(type);
	}

	final void setState(State newState) {
	if (newState != state) {
	logger.atFinest().log(
	"Transfer %d state %s -> %s", getSessionId(), state.getName(), newState.getName());
	}
	state = newState;
	}

	/** Sends a chunk. Returns true if sent, false if sending failed and the transfer was aborted. */
	final boolean sendChunk(Chunk chunk) {
	if (shouldAbortCallback.getAsBoolean()) {
	logger.atWarning().log("Abort signal received.");
	cleanUp(new TransferError(sessionId, Status.ABORTED));
	return false;
	}

	try {
	eventHandler.sendChunk(chunk);
	return true;
	} catch (TransferError transferError) {
	cleanUp(transferError);
	return false;
	}
	}

	final boolean sendChunk(Chunk.Builder chunk) {
	return sendChunk(chunk.build());
	}

	/** Performs final cleanup of a completed transfer. No packets are sent to the server. */
	private void cleanUp(Status status) {
	eventHandler.unregisterTransfer(sessionId);

	logger.atInfo().log("Transfer %d completed with status %s", sessionId, status);
	if (status.ok()) {
	setFutureResult();
	} else {
	future.setException(new TransferError(sessionId, status));
	}
	isCleanedUp = true;
	}

	/** Finishes the transfer due to an exception. No packets are sent to the server. */
	final void cleanUp(TransferError exception) {
	eventHandler.unregisterTransfer(sessionId);

	logger.atWarning().withCause(exception).log("Transfer %d terminated with exception", sessionId);
	future.setException(exception);
	isCleanedUp = true;
	}

	/** Sends a status chunk to the server and finishes the transfer. */
	final void sendFinalChunk(Status status) {
	logger.atFine().log("Sending final chunk for transfer %d with status %s", sessionId, status);
	// If the transfer was completed due to concurrently handling chunks, don't send.
	if (isCleanedUp) {
	logger.atFine().log("Skipping sending final chunk on already-completed transfer");
	return;
	}

	// Only call finish() if the sendChunk was successful. If it wasn't, the exception would have
	// already terminated the transfer.
	if (sendChunk(newChunk(Chunk.Type.COMPLETION).setStatus(status.code()))) {
	cleanUp(status);
	}
	}

	/** Invokes the transfer progress callback and logs the progress. */
	final void updateProgress(long bytesSent, long bytesConfirmedReceived, long totalSizeBytes) {
	TransferProgress progress =
	TransferProgress.create(bytesSent, bytesConfirmedReceived, totalSizeBytes);
	progressCallback.accept(progress);

	long durationNanos = Duration.between(startTime, Instant.now()).toNanos();
	long totalRate = durationNanos == 0 ? 0 : (bytesSent * 1_000_000_000 / durationNanos);

	logger.atFine().log("Transfer %d progress: "
	+ "%5.1f%% (%d B sent, %d B confirmed received of %s B total) at %d B/s",
	sessionId,
	progress.percentReceived(),
	bytesSent,
	bytesConfirmedReceived,
	totalSizeBytes == UNKNOWN_TRANSFER_SIZE ? "unknown" : totalSizeBytes,
	totalRate);
	}

	/** Represents a state in the transfer state machine. */
	abstract class State {
	private String getName() {
	return getClass().getSimpleName();
	}

	/** Called to handle a non-final chunk for this transfer. */
	void handleDataChunk(Chunk chunk) {
	logger.atFine().log(
	"Transfer %d [%s state]: Received unexpected data chunk", getSessionId(), getName());
	}

	/** Called to handle the final chunk for this transfer. */
	void handleFinalChunk(Chunk chunk) {
	Status status = Status.fromCode(chunk.getStatus());
	if (status != null) {
	cleanUp(status);
	} else {
	logger.atWarning().log("Received invalid status value %d", chunk.getStatus());
	sendFinalChunk(Status.INVALID_ARGUMENT);
	}
	setState(new Completed());
	}

	/** Called when this transfer's deadline expires. */
	void handleTimeout() {
	logger.atFine().log("Transfer %d [%s state]: Ignoring timeout", getSessionId(), getName());
	}

	/** Called if the transfer is cancelled by the user. */
	void handleCancellation() {
	sendFinalChunk(Status.CANCELLED);
	setState(new Completed());
	}

	/** Called when the transfer thread is shutting down. */
	void handleTermination() {
	sendFinalChunk(Status.ABORTED);
	setState(new Completed());
	}
	}

	private class Inactive extends State {}

	/** Recovering from an expired timeout. */
	class Recovery extends State {
	private int retries;

	@Override
	void handleDataChunk(Chunk chunk) {
	setState(getWaitingForDataState());
	state.handleDataChunk(chunk);
	}

	@Override
	void handleTimeout() {
	if (retries < maxRetries) {
	logger.atFiner().log("Transfer %d received no chunks for %d ms; retrying %d/%d",
	getSessionId(),
	timeoutMillis,
	retries,
	maxRetries);
	sendChunk(getChunkForRetry());
	retries += 1;
	setNextChunkTimeout();
	return;
	}

	sendFinalChunk(Status.DEADLINE_EXCEEDED);
	setState(new Completed());
	}
	}

	/** Transfer completed. Do nothing if the transfer is terminated or cancelled. */
	class Completed extends State {
	Completed() {
	deadline = NO_TIMEOUT;
	}

	@Override
	void handleTermination() {}

	@Override
	void handleCancellation() {}
	}
	}