Support payload offsets

Support sending a file or stream payload from a non-zero offset.
This allows the clients to resume a transfer. Upper layers should use
this feature only when both sides support it.

PiperOrigin-RevId: 385137570

cpp/core/internal/internal_payload.h

@@ -79,6 +79,14 @@ class InternalPayload {
   // cleanup may be required by the concrete implementation.
   virtual Exception AttachNextChunk(const ByteArray& chunk) = 0;
 
+  // Skips current stream pointer to the offset.
+  //
+  // Used when this is a resume outgoing transfer, so we want to skip
+  // some data until the offset position.
+  //
+  // @return the offset really skipped
+  virtual ExceptionOr<size_t> SkipToOffset(size_t offset) = 0;
+
   // Cleans up any resources used by this Payload. Called when we're stopping
   // early, e.g. after being cancelled or having no more recipients left.
   virtual void Close() {}

cpp/core/internal/internal_payload_factory.cc

@@ -62,6 +62,11 @@ class BytesInternalPayload : public InternalPayload {
     return {Exception::kSuccess};
   }
 
+  ExceptionOr<size_t> SkipToOffset(size_t offset) override {
+    NEARBY_LOGS(WARNING) << "Bytes payload does not support offsets";
+    return {Exception::kIo};
+  }
+
  private:
   // We're caching the total size here because the backing payload will be
   // moved to another owner during the lifetime of an incoming
@@ -108,6 +113,25 @@ class OutgoingStreamInternalPayload : public InternalPayload {
     return {Exception::kIo};
   }
 
+  ExceptionOr<size_t> SkipToOffset(size_t offset) override {
+    InputStream* stream = payload_.AsStream();
+    if (stream == nullptr) return {Exception::kIo};
+
+    ExceptionOr<size_t> real_offset = stream->Skip(offset);
+    if (real_offset.ok() && real_offset.GetResult() == offset) {
+      return real_offset;
+    }
+    // Close the outgoing stream on any error
+    stream->Close();
+    if (!real_offset.ok()) {
+      return real_offset;
+    }
+    NEARBY_LOGS(WARNING) << "Skip offset: " << real_offset.GetResult()
+                         << ", expected offset: " << offset << " for payload "
+                         << this;
+    return {Exception::kIo};
+  }
+
   void Close() override {
     // Ignore the potential Exception returned by close(), as a counterpart
     // to Java's closeQuietly().
@@ -140,6 +164,12 @@ class IncomingStreamInternalPayload : public InternalPayload {
     return output_stream_->Write(chunk);
   }
 
+  ExceptionOr<size_t> SkipToOffset(size_t offset) override {
+    NEARBY_LOGS(WARNING) << "Cannot skip offset for an incoming Payload "
+                         << this;
+    return {Exception::kIo};
+  }
+
   void Close() override { output_stream_->Close(); }
 
  private:
@@ -183,6 +213,28 @@ class OutgoingFileInternalPayload : public InternalPayload {
     return {Exception::kIo};
   }
 
+  ExceptionOr<size_t> SkipToOffset(size_t offset) override {
+    NEARBY_LOGS(INFO) << "SkipToOffset " << offset;
+    InputFile* file = payload_.AsFile();
+    if (!file) {
+      return {Exception::kIo};
+    }
+
+    ExceptionOr<size_t> real_offset = file->Skip(offset);
+    if (real_offset.ok() && real_offset.GetResult() == offset) {
+      return real_offset;
+    }
+    // Close the outgoing file on any error
+    file->Close();
+    if (!real_offset.ok()) {
+      return real_offset;
+    }
+    NEARBY_LOGS(WARNING) << "Skip offset: " << real_offset.GetResult()
+                         << ", expected offset: " << offset
+                         << " for file payload " << this;
+    return {Exception::kIo};
+  }
+
   void Close() override {
     InputFile* file = payload_.AsFile();
     if (file) file->Close();
@@ -218,6 +270,12 @@ class IncomingFileInternalPayload : public InternalPayload {
     return output_file_.Write(chunk);
   }
 
+  ExceptionOr<size_t> SkipToOffset(size_t offset) override {
+    NEARBY_LOGS(WARNING) << "Cannot skip offset for an incoming file Payload "
+                         << this;
+    return {Exception::kIo};
+  }
+
   void Close() override { output_file_.Close(); }
 
  private:

cpp/core/internal/internal_payload_factory_test.cc

@@ -127,6 +127,55 @@ TEST(InternalPayloadFActoryTest, CanCreateIternalPayloadFromFileMessage) {
   EXPECT_EQ(payload.GetId(), payload.AsFile()->GetPayloadId());
 }
 
+void CreateFileWithContents(Payload::Id payload_id, const ByteArray& contents) {
+  OutputFile file(payload_id);
+  EXPECT_TRUE(file.Write(contents).Ok());
+  EXPECT_TRUE(file.Close().Ok());
+}
+
+TEST(InternalPayloadFActoryTest,
+     SkipToOffset_FilePayloadValidOffset_SkipsOffset) {
+  ByteArray contents("0123456789");
+  constexpr size_t kOffset = 4;
+  size_t size_after_skip = contents.size() - kOffset;
+  Payload::Id payload_id = Payload::GenerateId();
+  CreateFileWithContents(payload_id, contents);
+  std::unique_ptr<InternalPayload> internal_payload =
+      CreateOutgoingInternalPayload(
+          Payload{payload_id, InputFile(payload_id, contents.size())});
+  EXPECT_NE(internal_payload, nullptr);
+
+  ExceptionOr<size_t> result = internal_payload->SkipToOffset(kOffset);
+
+  EXPECT_TRUE(result.ok());
+  EXPECT_EQ(result.GetResult(), kOffset);
+  EXPECT_EQ(internal_payload->GetTotalSize(), contents.size());
+  ByteArray contents_after_skip =
+      internal_payload->DetachNextChunk(size_after_skip);
+  EXPECT_EQ(contents_after_skip, ByteArray("456789"));
+}
+
+TEST(InternalPayloadFActoryTest,
+     SkipToOffset_StreamPayloadValidOffset_SkipsOffset) {
+  ByteArray contents("0123456789");
+  constexpr size_t kOffset = 6;
+  auto pipe = std::make_shared<Pipe>();
+  std::unique_ptr<InternalPayload> internal_payload =
+      CreateOutgoingInternalPayload(Payload{[pipe]() -> InputStream& {
+        return pipe->GetInputStream();  // NOLINT
+      }});
+  EXPECT_NE(internal_payload, nullptr);
+  pipe->GetOutputStream().Write(contents);
+
+  ExceptionOr<size_t> result = internal_payload->SkipToOffset(kOffset);
+
+  EXPECT_TRUE(result.ok());
+  EXPECT_EQ(result.GetResult(), kOffset);
+  EXPECT_EQ(internal_payload->GetTotalSize(), -1);
+  ByteArray contents_after_skip = internal_payload->DetachNextChunk(512);
+  EXPECT_EQ(contents_after_skip, ByteArray("6789"));
+}
+
 }  // namespace
 }  // namespace connections
 }  // namespace nearby

cpp/core/internal/payload_manager.cc

@@ -40,7 +40,7 @@ constexpr const absl::Duration PayloadManager::kWaitCloseTimeout;
 bool PayloadManager::SendPayloadLoop(
     ClientProxy* client, PendingPayload& pending_payload,
     PayloadTransferFrame::PayloadHeader& payload_header,
-    std::int64_t& next_chunk_offset) {
+    std::int64_t& next_chunk_offset, size_t resume_offset) {
   // in lieu of structured binding:
   auto pair = GetAvailableAndUnavailableEndpoints(pending_payload);
   const EndpointIds& available_endpoint_ids =
@@ -81,6 +81,24 @@ bool PayloadManager::SendPayloadLoop(
   // payload. For the sake of accuracy, we update the pending payload here
   // because it's after all payload terminating events are handled, but
   // right before we actually start detaching the next chunk.
+  if (next_chunk_offset == 0 && resume_offset > 0) {
+    ExceptionOr<size_t> real_offset =
+        pending_payload.GetInternalPayload()->SkipToOffset(resume_offset);
+    if (!real_offset.ok()) {
+      // Stop sending since it may cause remote file merging failed.
+      NEARBY_LOGS(WARNING) << "PayloadManager failed to skip offset "
+                           << resume_offset << " on payload_id "
+                           << pending_payload.GetInternalPayload()->GetId();
+      HandleFinishedOutgoingPayload(
+          client, available_endpoint_ids, payload_header, next_chunk_offset,
+          proto::connections::PayloadStatus::LOCAL_ERROR);
+      return false;
+    }
+    NEARBY_LOGS(VERBOSE) << "PayloadManager successfully skipped "
+                         << real_offset.GetResult() << " bytes on payload_id "
+                         << pending_payload.GetInternalPayload()->GetId();
+    next_chunk_offset = real_offset.GetResult();
+  }
   for (const auto& endpoint_id : available_endpoint_ids) {
     pending_payload.SetOffsetForEndpoint(endpoint_id, next_chunk_offset);
   }
@@ -105,8 +123,12 @@ bool PayloadManager::SendPayloadLoop(
     return false;
   }
 
-  PayloadTransferFrame::PayloadChunk payload_chunk(
-      CreatePayloadChunk(next_chunk_offset, std::move(next_chunk)));
+  // Only need to handle outgoing data chunk offset, because the offset will be
+  // used to decide if the received chunk is the initial payload chunk.
+  // In other cases, the offset should only be used in both side logs when error
+  // happened.
+  PayloadTransferFrame::PayloadChunk payload_chunk(CreatePayloadChunk(
+      next_chunk_offset - resume_offset, std::move(next_chunk)));
   const EndpointIds& failed_endpoint_ids = endpoint_manager_->SendPayloadChunk(
       payload_header, payload_chunk, available_endpoint_ids);
   // Check whether at least one endpoint failed.
@@ -346,10 +368,15 @@ void PayloadManager::SendPayload(ClientProxy* client,
   // completely done with. If we ever want to provide isolation across
   // ClientProxy objects this will need to be significantly re-architected.
   Payload::Type payload_type = payload.GetType();
+  size_t resume_offset =
+      FeatureFlags::GetInstance().GetFlags().enable_send_payload_offset
+          ? payload.GetOffset()
+          : 0;
+
   Payload::Id payload_id =
       CreateOutgoingPayload(std::move(payload), endpoint_ids);
   executor->Execute("send-payload", [this, client, endpoint_ids, payload_id,
-                                     payload_type]() {
+                                     payload_type, resume_offset]() {
     if (shutdown_.Get()) return;
     PendingPayload* pending_payload = GetPayload(payload_id);
     if (!pending_payload) {
@@ -363,12 +390,13 @@ void PayloadManager::SendPayload(ClientProxy* client,
     auto* internal_payload = pending_payload->GetInternalPayload();
     if (!internal_payload) return;
     PayloadTransferFrame::PayloadHeader payload_header{
-        CreatePayloadHeader(*internal_payload)};
+        CreatePayloadHeader(*internal_payload, resume_offset)};
     bool should_continue = true;
     std::int64_t next_chunk_offset = 0;
     while (should_continue && !shutdown_.Get()) {
-      should_continue = SendPayloadLoop(client, *pending_payload,
-                                        payload_header, next_chunk_offset);
+      should_continue =
+          SendPayloadLoop(client, *pending_payload, payload_header,
+                          next_chunk_offset, resume_offset);
     }
     RunOnStatusUpdateThread("destroy-payload",
                             [this, payload_id]()
@@ -544,12 +572,16 @@ int PayloadManager::GetOptimalChunkSize(EndpointIds endpoint_ids) {
 }
 
 PayloadTransferFrame::PayloadHeader PayloadManager::CreatePayloadHeader(
-    const InternalPayload& internal_payload) {
+    const InternalPayload& internal_payload, size_t offset) {
   PayloadTransferFrame::PayloadHeader payload_header;
+  size_t payload_size = internal_payload.GetTotalSize();
 
   payload_header.set_id(internal_payload.GetId());
   payload_header.set_type(internal_payload.GetType());
-  payload_header.set_total_size(internal_payload.GetTotalSize());
+  payload_header.set_total_size(payload_size ==
+                                        InternalPayload::kIndeterminateSize
+                                    ? InternalPayload::kIndeterminateSize
+                                    : payload_size - offset);
 
   return payload_header;
 }

cpp/core/internal/payload_manager.h

@@ -189,7 +189,7 @@ class PayloadManager : public EndpointManager::FrameProcessor {
 
   bool SendPayloadLoop(ClientProxy* client, PendingPayload& pending_payload,
                        PayloadTransferFrame::PayloadHeader& payload_header,
-                       std::int64_t& next_chunk_offset);
+                       std::int64_t& next_chunk_offset, size_t resume_offset);
   void SendClientCallbacksForFinishedIncomingPayloadRunnable(
       ClientProxy* client, const std::string& endpoint_id,
       const PayloadTransferFrame::PayloadHeader& payload_header,
@@ -211,7 +211,7 @@ class PayloadManager : public EndpointManager::FrameProcessor {
   int GetOptimalChunkSize(EndpointIds endpoint_ids);
 
   PayloadTransferFrame::PayloadHeader CreatePayloadHeader(
-      const InternalPayload& payload);
+      const InternalPayload& payload, size_t offset);
   PayloadTransferFrame::PayloadChunk CreatePayloadChunk(std::int64_t offset,
                                                         ByteArray body);
 

cpp/core/internal/payload_manager_test.cc

@@ -311,6 +311,59 @@ TEST_P(PayloadManagerTest, CanCancelPayloadOnSenderSide) {
   env_.Stop();
 }
 
+TEST_P(PayloadManagerTest, SendPayloadWithSkip_StreamPayload) {
+  constexpr size_t kOffset = 3;
+  env_.Start();
+  PayloadSimulationUser user_a(kDeviceA, GetParam());
+  PayloadSimulationUser user_b(kDeviceB, GetParam());
+  ASSERT_TRUE(SetupConnection(user_a, user_b));
+
+  auto pipe = std::make_shared<Pipe>();
+  OutputStream& tx = pipe->GetOutputStream();
+
+  user_a.ExpectPayload(payload_latch_);
+  const ByteArray message{std::string(kMessage)};
+  // The first write to the output stream will send the first PAYLOAD_TRANSFER
+  // packet with payload info and message data.
+  tx.Write(message);
+
+  Payload payload([pipe]() -> InputStream& {
+    return pipe->GetInputStream();  // NOLINT
+  });
+  payload.SetOffset(kOffset);
+  user_b.SendPayload(std::move(payload));
+  ASSERT_TRUE(payload_latch_.Await(kDefaultTimeout).result());
+  ASSERT_NE(user_a.GetPayload().AsStream(), nullptr);
+  InputStream& rx = *user_a.GetPayload().AsStream();
+  NEARBY_LOG(INFO, "Stream extracted.");
+
+  EXPECT_TRUE(user_a.WaitForProgress(
+      [&message](const PayloadProgressInfo& info) {
+        return info.bytes_transferred >= message.size() - kOffset;
+      },
+      kProgressTimeout));
+  ByteArray result = rx.Read(Pipe::kChunkSize).result();
+  EXPECT_EQ(result, ByteArray("sage"));
+  NEARBY_LOG(INFO, "Packet 1 handled.");
+
+  tx.Write(message);
+  EXPECT_TRUE(user_a.WaitForProgress(
+      [&message](const PayloadProgressInfo& info) {
+        return info.bytes_transferred >= 2 * message.size() - kOffset;
+      },
+      kProgressTimeout));
+  ByteArray result2 = rx.Read(Pipe::kChunkSize).result();
+  EXPECT_EQ(result2, message);
+  NEARBY_LOG(INFO, "Packet 2 handled.");
+
+  rx.Close();
+  tx.Close();
+  NEARBY_LOG(INFO, "Test completed.");
+  user_a.Stop();
+  user_b.Stop();
+  env_.Stop();
+}
+
 INSTANTIATE_TEST_SUITE_P(ParametrisedPayloadManagerTest, PayloadManagerTest,
                          ::testing::ValuesIn(kTestCases));
 

cpp/core/payload.h

@@ -89,17 +89,28 @@ class Payload {
   // Returns Payload type.
   Type GetType() const { return type_; }
 
+  // Sets the payload offset in bytes
+  void SetOffset(size_t offset) {
+    CHECK(type_ == Type::kFile || type_ == Type::kStream);
+    InputFile* file = AsFile();
+    if (file != nullptr) {
+      CHECK(offset < file->GetTotalSize());
+    }
+    offset_ = offset;
+  }
+
+  size_t GetOffset() { return offset_; }
+
   // Generate Payload Id; to be passed to outgoing file constructor.
   static Id GenerateId() { return Prng().NextInt64(); }
 
  private:
-  Type FindType(const Content& content) const {
-    return static_cast<Type>(content_.index());
-  }
+  Type FindType() const { return static_cast<Type>(content_.index()); }
 
   Content content_;
   Id id_{GenerateId()};
-  Type type_{FindType(content_)};
+  Type type_{FindType()};
+  size_t offset_{0};
 };
 
 }  // namespace connections

cpp/core/payload_test.cc

@@ -43,18 +43,25 @@ TEST(PayloadTest, SupportsByteArrayType) {
 }
 
 TEST(PayloadTest, SupportsFileType) {
+  constexpr size_t kOffset = 99;
   const auto payload_id = Payload::GenerateId();
   InputFile file(payload_id, 100);
   InputStream& stream = file.GetInputStream();
+
   Payload payload(payload_id, std::move(file));
+  payload.SetOffset(kOffset);
+
   EXPECT_EQ(payload.GetType(), Payload::Type::kFile);
   EXPECT_EQ(payload.AsStream(), nullptr);
   EXPECT_EQ(&payload.AsFile()->GetInputStream(), &stream);
   EXPECT_EQ(payload.AsBytes(), ByteArray{});
+  EXPECT_EQ(payload.GetOffset(), kOffset);
 }
 
 TEST(PayloadTest, SupportsStreamType) {
+  constexpr size_t kOffset = 1234456;
   auto pipe = std::make_shared<Pipe>();
+
   Payload payload(
       [streamable = pipe]() -> InputStream& {
         // For some reason, linter warns us that we return a dangling reference.
@@ -63,10 +70,13 @@ TEST(PayloadTest, SupportsStreamType) {
         // shared_ptr<Pipe> is captured by value.
         return streamable->GetInputStream();  // NOLINT
       });
+  payload.SetOffset(kOffset);
+
   EXPECT_EQ(payload.GetType(), Payload::Type::kStream);
   EXPECT_EQ(payload.AsStream(), &pipe->GetInputStream());
   EXPECT_EQ(payload.AsFile(), nullptr);
   EXPECT_EQ(payload.AsBytes(), ByteArray{});
+  EXPECT_EQ(payload.GetOffset(), kOffset);
 }
 
 TEST(PayloadTest, PayloadIsMoveable) {