diff --git a/tpu_raiden/transport/peregrine/src/internal/channel/BUILD b/tpu_raiden/transport/peregrine/src/internal/channel/BUILD new file mode 100644 index 00000000..bf27dcb5 --- /dev/null +++ b/tpu_raiden/transport/peregrine/src/internal/channel/BUILD @@ -0,0 +1,187 @@ +# Copyright 2026 Google LLC. +# +# 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 +# +# 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. + +load("@rules_cc//cc:defs.bzl", "cc_library", "cc_test") + +package(default_visibility = ["//visibility:public"]) + +cc_library( + name = "channel_types", + hdrs = ["channel_types.h"], +) + +cc_library( + name = "channel", + hdrs = ["channel.h"], + deps = [ + ":channel_types", + "//tpu_raiden/transport/peregrine/src/internal/base:types", + "@com_google_absl//absl/types:span", + ], +) + +cc_library( + name = "channel_test_util", + testonly = True, + srcs = ["channel_test_util.cc"], + hdrs = ["channel_test_util.h"], + deps = [ + ":channel", + ":channel_msg", + ":channel_stream", + ":channel_util", + ":pipe", + "//tpu_raiden/transport/peregrine/src/internal/socket:socket_test_util", + "@com_google_absl//absl/log:check", + ], +) + +cc_library( + name = "channel_tcp", + srcs = ["channel_tcp.cc"], + hdrs = ["channel_tcp.h"], + deps = [ + ":channel", + ":channel_types", + "//tpu_raiden/transport/peregrine/src/internal/base:types", + "//tpu_raiden/transport/peregrine/src/internal/socket:socket_tcp", + "//tpu_raiden/transport/peregrine/src/internal/util", + "@com_google_absl//absl/base:core_headers", + "@com_google_absl//absl/log:check", + "@com_google_absl//absl/strings", + "@com_google_absl//absl/types:span", + ], +) + +cc_library( + name = "channel_udp", + srcs = ["channel_udp.cc"], + hdrs = ["channel_udp.h"], + deps = [ + ":channel", + ":channel_types", + "//tpu_raiden/transport/peregrine/src/internal/base:types", + "//tpu_raiden/transport/peregrine/src/internal/socket:socket_udp", + "//tpu_raiden/transport/peregrine/src/internal/util", + "@com_google_absl//absl/base:core_headers", + "@com_google_absl//absl/log:check", + "@com_google_absl//absl/strings", + "@com_google_absl//absl/types:span", + ], +) + +cc_library( + name = "pipe", + testonly = True, + srcs = ["pipe.cc"], + hdrs = ["pipe.h"], + deps = [ + "//tpu_raiden/transport/peregrine/src/internal/util:test_iov", + "@com_google_absl//absl/base:core_headers", + "@com_google_absl//absl/synchronization", + ], +) + +cc_library( + name = "channel_msg", + testonly = True, + srcs = ["channel_msg.cc"], + hdrs = ["channel_msg.h"], + deps = [ + ":channel", + ":channel_types", + ":pipe", + "//tpu_raiden/transport/peregrine/src/internal/base:types", + "//tpu_raiden/transport/peregrine/src/internal/util", + "//tpu_raiden/transport/peregrine/src/internal/util:test_iov", + "//tpu_raiden/transport/peregrine/src/util", + "//util/random:shared_bit_gen", + "@com_google_absl//absl/base:core_headers", + "@com_google_absl//absl/log:check", + "@com_google_absl//absl/strings:str_format", + "@com_google_absl//absl/synchronization", + "@com_google_absl//absl/types:span", + ], +) + +cc_library( + name = "channel_stream", + testonly = True, + srcs = ["channel_stream.cc"], + hdrs = ["channel_stream.h"], + deps = [ + ":channel", + ":channel_types", + ":pipe", + "//tpu_raiden/transport/peregrine/src/internal/base:types", + "//tpu_raiden/transport/peregrine/src/internal/util", + "//tpu_raiden/transport/peregrine/src/internal/util:test_iov", + "//tpu_raiden/transport/peregrine/src/util", + "//util/random:shared_bit_gen", + "@com_google_absl//absl/base:core_headers", + "@com_google_absl//absl/log:check", + "@com_google_absl//absl/strings:str_format", + "@com_google_absl//absl/synchronization", + "@com_google_absl//absl/types:span", + ], +) + +cc_test( + name = "channel_test", + srcs = ["channel_test.cc"], + shard_count = 6, + deps = [ + ":channel", + ":channel_test_util", + "//tpu_raiden/transport/peregrine/src/internal/base:types", + "//tpu_raiden/transport/peregrine/src/util", + "@com_google_absl//absl/log", + "@com_google_absl//absl/log:check", + "@com_google_absl//absl/strings:str_format", + "@com_google_absl//absl/types:span", + "@com_google_googletest//:gtest_main", + ], +) + +cc_library( + name = "channel_util", + srcs = ["channel_util.cc"], + hdrs = ["channel_util.h"], + deps = [ + ":channel", + ":channel_tcp", + ":channel_udp", + "//tpu_raiden/transport/peregrine/src/internal/base:endpoint", + "//tpu_raiden/transport/peregrine/src/internal/socket:connector", + "//tpu_raiden/transport/peregrine/src/internal/socket:socket_tcp", + "//tpu_raiden/transport/peregrine/src/internal/socket:socket_udp", + "@com_google_absl//absl/log:check", + ], +) + +cc_test( + name = "channel_util_test", + srcs = ["channel_util_test.cc"], + shard_count = 2, + deps = [ + ":channel_util", + "//tpu_raiden/transport/peregrine/src/internal/base:endpoint", + "//tpu_raiden/transport/peregrine/src/internal/socket:acceptor", + "//tpu_raiden/transport/peregrine/src/internal/socket:socket_tcp", + "//tpu_raiden/transport/peregrine/src/internal/util:test_util", + "@com_google_absl//absl/log:check", + "@com_google_absl//absl/strings:str_format", + "@com_google_googletest//:gtest_main", + ], +) diff --git a/tpu_raiden/transport/peregrine/src/internal/channel/channel.h b/tpu_raiden/transport/peregrine/src/internal/channel/channel.h new file mode 100644 index 00000000..b83bd406 --- /dev/null +++ b/tpu_raiden/transport/peregrine/src/internal/channel/channel.h @@ -0,0 +1,81 @@ +// Copyright 2026 Google LLC. +// +// 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 +// +// 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. + +#ifndef THIRD_PARTY_PEREGRINE_SRC_INTERNAL_CHANNEL_CHANNEL_H_ +#define THIRD_PARTY_PEREGRINE_SRC_INTERNAL_CHANNEL_CHANNEL_H_ + +#include +#include +#include + +#include "absl/types/span.h" +#include "tpu_raiden/transport/peregrine/src/internal/base/types.h" +#include "tpu_raiden/transport/peregrine/src/internal/channel/channel_types.h" + +namespace peregrine::internal { + +// An interface that specifies a channel abstraction. It is used for +// communications between two endpoints. +class Channel { + public: + // Destructor. + virtual ~Channel() = default; + + // Returns the channel type. + virtual ChannelType Type() const = 0; + + // Writes `len` bytes of data from `buf` to the channel. + // Returns the number of bytes actually written if successful. Zero byte means + // no data has been written due to non-error reasons. Returns -1 on error. + virtual ssize_t Write(const Byte* buf, size_t len) = 0; + + // Writes a number of buffers described by the `iovecs` to the channel. + // Returns the number of bytes actually written if successful. Zero byte means + // no data has been written due to non-error reasons. Returns -1 on error. + virtual ssize_t WriteV(absl::Span iovecs) = 0; + + // Reads data from the channel into the `buf`. + // For stream channel, it reads exactly `len` bytes of data. + // For message channel, it reads one message of up to `len` bytes. + // Returns the number of bytes actually read if successful. + // For stream channel, returns 0 if the peer side has closed the connection. + // For message channel, returns 0 if the received packet has no payload. + // Returns -1 on error. + virtual ssize_t Read(Byte* buf, size_t len) = 0; + + // Reads data from the channel into the `iovecs` buffers. + // For stream channel, it reads exactly `length(iovecs)` bytes of data. + // For message channel, it reads one message of up to `length(iovecs)` bytes. + // Returns the number of bytes actually read if successful. + // For stream channel, returns 0 if the peer side has closed the connection. + // For message channel, returns 0 if the received packet has no payload. + // Returns -1 on error. + virtual ssize_t ReadV(absl::Span iovecs) = 0; + + // Shuts down the channel. After the channel is shutdown, write calls will + // return -1, so no more data can be injected into the channel. Read calls + // will continue to read the remaining data in the channel, if any. + virtual void Shutdown() = 0; + + // Returns a string representation for the channel. + virtual std::string ToString() const = 0; +}; + +inline std::ostream& operator<<(std::ostream& os, const Channel& c) { + return os << c.ToString(); +} + +} // namespace peregrine::internal + +#endif // THIRD_PARTY_PEREGRINE_SRC_INTERNAL_CHANNEL_CHANNEL_H_ diff --git a/tpu_raiden/transport/peregrine/src/internal/channel/channel_msg.cc b/tpu_raiden/transport/peregrine/src/internal/channel/channel_msg.cc new file mode 100644 index 00000000..fe833aa5 --- /dev/null +++ b/tpu_raiden/transport/peregrine/src/internal/channel/channel_msg.cc @@ -0,0 +1,135 @@ +// Copyright 2026 Google LLC. +// +// 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 +// +// 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. + +#include "tpu_raiden/transport/peregrine/src/internal/channel/channel_msg.h" + +#include +#include +#include +#include + +#include "absl/base/optimization.h" +#include "absl/log/check.h" +#include "absl/synchronization/mutex.h" +#include "absl/types/span.h" +#include "tpu_raiden/transport/peregrine/src/internal/base/types.h" +#include "tpu_raiden/transport/peregrine/src/internal/channel/pipe.h" +#include "tpu_raiden/transport/peregrine/src/internal/util/test_iov.h" +#include "tpu_raiden/transport/peregrine/src/internal/util/util.h" +#include "tpu_raiden/transport/peregrine/src/util/util.h" +#include "util/random/shared_bit_gen.h" + +namespace peregrine::internal::testing { + +MemMsgChannel::MemMsgChannel(const BidiPipe& bidi, const int error_rate) + : error_rate_(std::clamp(error_rate, 0, 100)), + in_pipe_(bidi.InputPipe()), + out_pipe_(bidi.OutputPipe()) { + DCHECK_NE(in_pipe_, nullptr); + DCHECK_NE(out_pipe_, nullptr); +} + +ssize_t MemMsgChannel::WriteV(const absl::Span iovecs) { + DCHECK(IsValid(iovecs)); + + const size_t len = TotalLength(iovecs); + DCHECK_GE(len, 1); + + // To keep the message boundary, merge multiple iovecs into a single one. + OwnedIoVec owned_iov = TestOnly_Linearize(iovecs); + + absl::MutexLock lock(out_pipe_->mu); + if (out_pipe_->shutdown) return -1; + out_pipe_->queue.push_back(std::move(owned_iov)); + return len; +} + +bool MemMsgChannel::hasIncomingData() const { return in_pipe_->HasData(); } + +ssize_t MemMsgChannel::Read(Byte* const buf, const size_t len) { + DCHECK_NE(buf, nullptr); + DCHECK_GE(len, 1); + + absl::MutexLock lock(in_pipe_->mu); + in_pipe_->mu.Await(absl::Condition(this, &MemMsgChannel::hasIncomingData)); + if (in_pipe_->queue.empty()) { + DCHECK(in_pipe_->shutdown); + return 0; + } + + OwnedIoVec owned_iov = std::move(in_pipe_->queue.front()); + in_pipe_->queue.pop_front(); + if (error()) return -1; + + const size_t size = owned_iov.size; + if (size == 0) { + return 0; + } else if (size > len) { + return -1; + } else { + DCHECK(1 <= size && size <= len); + std::memcpy(buf, owned_iov.data.get(), size); + return size; + } +} + +ssize_t MemMsgChannel::ReadV(absl::Span iovecs) { + DCHECK(IsValid(iovecs)); + DCHECK_GE(TotalLength(iovecs), 1); + + absl::MutexLock lock(in_pipe_->mu); + in_pipe_->mu.Await(absl::Condition(this, &MemMsgChannel::hasIncomingData)); + if (in_pipe_->queue.empty()) { + DCHECK(in_pipe_->shutdown); + return 0; + } + + OwnedIoVec owned_iov = std::move(in_pipe_->queue.front()); + in_pipe_->queue.pop_front(); + if (error()) return -1; + + size_t size = owned_iov.size; + if (size == 0) return 0; + if (size > TotalLength(iovecs)) return -1; + + size_t offset = 0; + Byte* data = owned_iov.data.get(); + for (auto& iov : iovecs) { + if (const size_t len = iov.iov_len; size <= len) { + std::memcpy(iov.iov_base, data + offset, size); + offset += size; + size = 0; + break; + } else { + std::memcpy(iov.iov_base, data + offset, len); + offset += len; + size -= len; + } + } + DCHECK_EQ(size, 0); + return offset; +} + +void MemMsgChannel::Shutdown() { + in_pipe_->Shutdown(); + out_pipe_->Shutdown(); +} + +bool MemMsgChannel::error() const { + if ABSL_PREDICT_TRUE (error_rate_ <= 0) return false; + util_random::SharedBitGen bitgen; + return util::Random(bitgen, 1, 100) <= error_rate_; +} + +} // namespace peregrine::internal::testing diff --git a/tpu_raiden/transport/peregrine/src/internal/channel/channel_msg.h b/tpu_raiden/transport/peregrine/src/internal/channel/channel_msg.h new file mode 100644 index 00000000..6dcf41fd --- /dev/null +++ b/tpu_raiden/transport/peregrine/src/internal/channel/channel_msg.h @@ -0,0 +1,91 @@ +// Copyright 2026 Google LLC. +// +// 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 +// +// 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. + +#ifndef THIRD_PARTY_PEREGRINE_SRC_INTERNAL_CHANNEL_CHANNEL_MSG_H_ +#define THIRD_PARTY_PEREGRINE_SRC_INTERNAL_CHANNEL_CHANNEL_MSG_H_ + +#include +#include +#include + +#include "absl/base/thread_annotations.h" +#include "absl/strings/str_format.h" +#include "absl/types/span.h" +#include "tpu_raiden/transport/peregrine/src/internal/base/types.h" +#include "tpu_raiden/transport/peregrine/src/internal/channel/channel.h" +#include "tpu_raiden/transport/peregrine/src/internal/channel/channel_types.h" +#include "tpu_raiden/transport/peregrine/src/internal/channel/pipe.h" + +namespace peregrine::internal::testing { + +// An unreliable memory channel to help testing: lossless/lossy, message. +// It is thread-safe. +class MemMsgChannel final : public Channel { + public: + // Constructor. + explicit MemMsgChannel(const BidiPipe& bidi, int error_rate); + + // Returns the channel type. + constexpr ChannelType Type() const override { + return ChannelType::kUnreliableMessage; + } + + // Writes a buffer of `len` bytes to the channel. + // Returns the number of bytes actually written if successful. Zero byte means + // no data has been written due to non-error reasons. Returns -1 on error. + ssize_t Write(const Byte* buf, size_t len) override { + return WriteV({{(void*)buf, len}}); + } + + // Writes a number of buffers described by the `iovecs` to the channel. + // Returns the number of bytes actually written if successful. Zero byte means + // no data has been written due to non-error reasons. Returns -1 on error. + ssize_t WriteV(absl::Span iovecs) override; + + // Reads one message of up to `len` bytes from the channel into the `buf`. + // Returns the number of bytes actually read if successful. Returns 0 if the + // received packet has no payload. Returns -1 on error. + ssize_t Read(Byte* buf, size_t len) override; + + // Reads one message of up to `length(iovecs)` bytes from the channel into + // the buffers. Returns the number of bytes actually read if successful. + // Returns 0 if the received packet has no payload. Returns -1 on error. + ssize_t ReadV(absl::Span iovecs) override; + + // Shuts down the channel. After the channel is shutdown, write calls will + // return -1, so no more data can be injected into the channel. Read calls + // will continue to read the remaining data in the channel, if any. + void Shutdown() override; + + // Returns a string representation for the channel. + std::string ToString() const override { + return absl::StrFormat("MemMsgChannel: error_rate=%d%%", error_rate_); + } + + private: + // Returns true iff the `in_pipe_` has data. + bool hasIncomingData() const ABSL_EXCLUSIVE_LOCKS_REQUIRED(in_pipe_->mu); + + // Returns true iff the channel read/write should emulate an error. + bool error() const; + + private: + const int error_rate_; + std::shared_ptr in_pipe_; + std::shared_ptr out_pipe_; +}; + +} // namespace peregrine::internal::testing + +#endif // THIRD_PARTY_PEREGRINE_SRC_INTERNAL_CHANNEL_CHANNEL_MSG_H_ diff --git a/tpu_raiden/transport/peregrine/src/internal/channel/channel_stream.cc b/tpu_raiden/transport/peregrine/src/internal/channel/channel_stream.cc new file mode 100644 index 00000000..c69c4bac --- /dev/null +++ b/tpu_raiden/transport/peregrine/src/internal/channel/channel_stream.cc @@ -0,0 +1,139 @@ +// Copyright 2026 Google LLC. +// +// 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 +// +// 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. + +#include "tpu_raiden/transport/peregrine/src/internal/channel/channel_stream.h" + +#include +#include +#include +#include +#include + +#include "absl/base/optimization.h" +#include "absl/log/check.h" +#include "absl/strings/str_format.h" +#include "absl/synchronization/mutex.h" +#include "absl/types/span.h" +#include "tpu_raiden/transport/peregrine/src/internal/base/types.h" +#include "tpu_raiden/transport/peregrine/src/internal/channel/pipe.h" +#include "tpu_raiden/transport/peregrine/src/internal/util/test_iov.h" +#include "tpu_raiden/transport/peregrine/src/internal/util/util.h" +#include "tpu_raiden/transport/peregrine/src/util/util.h" +#include "util/random/shared_bit_gen.h" + +namespace peregrine::internal::testing { + +MemStreamChannel::MemStreamChannel(const BidiPipe& bidi, const int error_rate) + : error_rate_(std::clamp(error_rate, 0, 100)), + in_pipe_(bidi.InputPipe()), + out_pipe_(bidi.OutputPipe()) { + DCHECK_NE(in_pipe_, nullptr); + DCHECK_NE(out_pipe_, nullptr); +} + +ssize_t MemStreamChannel::WriteV(const absl::Span iovecs) { + DCHECK(IsValid(iovecs)); + + const size_t len = TotalLength(iovecs); + DCHECK_GE(len, 1); + + absl::MutexLock lock(out_pipe_->mu); + if (out_pipe_->shutdown) return -1; + for (const auto& v : iovecs) { + // Do not merge multiple iovecs into a single one. + DCHECK(IsValid(v)); + OwnedIoVec owned_iov = TestOnly_Linearize({v}); + out_pipe_->queue.push_back(std::move(owned_iov)); + } + return len; +} + +bool MemStreamChannel::hasIncomingData() const { return in_pipe_->HasData(); } + +ssize_t MemStreamChannel::Read(Byte* const buf, const size_t len) { + DCHECK_NE(buf, nullptr); + DCHECK_GE(len, 1); + + absl::MutexLock lock(in_pipe_->mu); + in_pipe_->mu.Await(absl::Condition(this, &MemStreamChannel::hasIncomingData)); + if (in_pipe_->queue.empty()) { + DCHECK(in_pipe_->shutdown); + return 0; + } + + size_t rcvd = 0; + size_t left = len; + Byte* ptr = buf; + while (left > 0) { + if (in_pipe_->queue.empty()) return rcvd ?: -1; + + OwnedIoVec owned_iov = std::move(in_pipe_->queue.front()); + in_pipe_->queue.pop_front(); + if (error()) return rcvd ?: -1; + + const Byte* const iov_ptr = owned_iov.data.get(); + const size_t size = owned_iov.size; + if (size == 0) { + continue; + } else if (size <= left) { + std::memcpy(ptr, iov_ptr, size); + ptr += size; + left -= size; + rcvd += size; + } else { + std::memcpy(ptr, iov_ptr, left); + const void* const p = iov_ptr + left; + const IoVec v(const_cast(p), size - left); + OwnedIoVec iov = TestOnly_Linearize({v}); + in_pipe_->queue.push_front(std::move(iov)); + rcvd += left; + left = 0; + break; + } + } + return rcvd; +} + +ssize_t MemStreamChannel::ReadV(absl::Span iovecs) { + DCHECK(IsValid(iovecs)); + DCHECK_GE(TotalLength(iovecs), 1); + + ssize_t total = 0; + for (const auto& iov : iovecs) { + Byte* buf = reinterpret_cast(iov.iov_base); + const size_t len = iov.iov_len; + const ssize_t n = Read(buf, len); + if (n <= 0) return total ?: n; + total += n; + if (static_cast(n) < len) break; + } + return total; +} + +void MemStreamChannel::Shutdown() { + in_pipe_->Shutdown(); + out_pipe_->Shutdown(); +} + +std::string MemStreamChannel::ToString() const { + return absl::StrFormat("MemStreamChannel: error_rate=%d%%", error_rate_); +} + +bool MemStreamChannel::error() const { + if ABSL_PREDICT_TRUE (error_rate_ <= 0) return false; + util_random::SharedBitGen bitgen; + return util::Random(bitgen, 1, 100) <= error_rate_; +} + +} // namespace peregrine::internal::testing diff --git a/tpu_raiden/transport/peregrine/src/internal/channel/channel_stream.h b/tpu_raiden/transport/peregrine/src/internal/channel/channel_stream.h new file mode 100644 index 00000000..114b07c5 --- /dev/null +++ b/tpu_raiden/transport/peregrine/src/internal/channel/channel_stream.h @@ -0,0 +1,88 @@ +// Copyright 2026 Google LLC. +// +// 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 +// +// 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. + +#ifndef THIRD_PARTY_PEREGRINE_SRC_INTERNAL_CHANNEL_CHANNEL_STREAM_H_ +#define THIRD_PARTY_PEREGRINE_SRC_INTERNAL_CHANNEL_CHANNEL_STREAM_H_ + +#include +#include +#include + +#include "absl/base/thread_annotations.h" +#include "absl/types/span.h" +#include "tpu_raiden/transport/peregrine/src/internal/base/types.h" +#include "tpu_raiden/transport/peregrine/src/internal/channel/channel.h" +#include "tpu_raiden/transport/peregrine/src/internal/channel/channel_types.h" +#include "tpu_raiden/transport/peregrine/src/internal/channel/pipe.h" + +namespace peregrine::internal::testing { + +// A reliable memory channel to help testing: lossless, stream. +// It is thread-safe. +class MemStreamChannel final : public Channel { + public: + // Constructor for paired bidirectional pipes. + explicit MemStreamChannel(const BidiPipe& bidi, int error_rate); + + // Returns the channel type. + constexpr ChannelType Type() const override { + return ChannelType::kReliableStream; + } + + // Writes a buffer of `len` bytes to the channel. + // Returns the number of bytes actually written if successful. Zero byte means + // no data has been written due to non-error reasons. Returns -1 on error. + ssize_t Write(const Byte* buf, size_t len) override { + return WriteV({{(void*)buf, len}}); + } + + // Writes a number of buffers described by the `iovecs` to the channel. + // Returns the number of bytes actually written if successful. Zero byte means + // no data has been written due to non-error reasons. Returns -1 on error. + ssize_t WriteV(absl::Span iovecs) override; + + // Reads exactly `len` bytes of data from the channel into the `buf`. + // Returns the number of bytes actually read if successful. Returns 0 if + // the peer side has closed the connection. Returns -1 on error. + ssize_t Read(Byte* buf, size_t len) override; + + // Reads exactly `length(iovecs)` bytes of data from the channel into the + // buffers. Returns the number of bytes actually read if successful. + // Returns 0 if the peer side has closed the connection. Returns -1 on error. + ssize_t ReadV(absl::Span iovecs) override; + + // Shuts down the channel. After the channel is shutdown, write calls will + // return -1, so no more data can be injected into the channel. Read calls + // will continue to read the remaining data in the channel, if any. + void Shutdown() override; + + // Returns a string representation for the channel. + std::string ToString() const override; + + private: + // Returns true iff the `in_pipe_` has data. + bool hasIncomingData() const ABSL_EXCLUSIVE_LOCKS_REQUIRED(in_pipe_->mu); + + // Returns true iff the channel read/write should emulate an error. + bool error() const; + + private: + const int error_rate_; + std::shared_ptr in_pipe_; + std::shared_ptr out_pipe_; +}; + +} // namespace peregrine::internal::testing + +#endif // THIRD_PARTY_PEREGRINE_SRC_INTERNAL_CHANNEL_CHANNEL_STREAM_H_ diff --git a/tpu_raiden/transport/peregrine/src/internal/channel/channel_tcp.cc b/tpu_raiden/transport/peregrine/src/internal/channel/channel_tcp.cc new file mode 100644 index 00000000..15fd054f --- /dev/null +++ b/tpu_raiden/transport/peregrine/src/internal/channel/channel_tcp.cc @@ -0,0 +1,62 @@ +// Copyright 2026 Google LLC. +// +// 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 +// +// 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. + +#include "tpu_raiden/transport/peregrine/src/internal/channel/channel_tcp.h" + +#include + +#include "absl/base/optimization.h" +#include "absl/log/check.h" +#include "absl/strings/str_cat.h" +#include "absl/types/span.h" +#include "tpu_raiden/transport/peregrine/src/internal/base/types.h" +#include "tpu_raiden/transport/peregrine/src/internal/util/util.h" + +namespace peregrine::internal { + +ssize_t TcpChannel::WriteV(const absl::Span iovecs) { + DCHECK(IsValid(iovecs)); + DCHECK_GE(iovecs.size(), 1); + DCHECK_LE(iovecs.size(), IOV_MAX); + DCHECK_GE(TotalLength(iovecs), 1); + + if ABSL_PREDICT_FALSE (iovecs.size() == 1) { + const auto [buf, len] = BufLen(iovecs[0]); + return socket_->Send(buf, len); + } else { + DCHECK_GE(iovecs.size(), 2); + return socket_->SendV(iovecs); + } +} + +ssize_t TcpChannel::ReadV(absl::Span iovecs) { + DCHECK(IsValid(iovecs)); + DCHECK_GE(iovecs.size(), 1); + DCHECK_LE(iovecs.size(), IOV_MAX); + DCHECK_GE(TotalLength(iovecs), 1); + + if ABSL_PREDICT_FALSE (iovecs.size() == 1) { + const auto [buf, len] = BufLen(iovecs[0]); + return socket_->Recv(buf, len); + } else { + DCHECK_GE(iovecs.size(), 2); + return socket_->RecvV(iovecs); + } +} + +std::string TcpChannel::ToString() const { + return absl::StrCat("TcpChannel: ", socket_->ToString()); +} + +} // namespace peregrine::internal diff --git a/tpu_raiden/transport/peregrine/src/internal/channel/channel_tcp.h b/tpu_raiden/transport/peregrine/src/internal/channel/channel_tcp.h new file mode 100644 index 00000000..6231ca7d --- /dev/null +++ b/tpu_raiden/transport/peregrine/src/internal/channel/channel_tcp.h @@ -0,0 +1,91 @@ +// Copyright 2026 Google LLC. +// +// 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 +// +// 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. + +#ifndef THIRD_PARTY_PEREGRINE_SRC_INTERNAL_CHANNEL_CHANNEL_TCP_H_ +#define THIRD_PARTY_PEREGRINE_SRC_INTERNAL_CHANNEL_CHANNEL_TCP_H_ + +#include + +#include +#include +#include +#include + +#include "absl/log/check.h" +#include "absl/types/span.h" +#include "tpu_raiden/transport/peregrine/src/internal/base/types.h" +#include "tpu_raiden/transport/peregrine/src/internal/channel/channel.h" +#include "tpu_raiden/transport/peregrine/src/internal/channel/channel_types.h" +#include "tpu_raiden/transport/peregrine/src/internal/socket/socket_tcp.h" + +namespace peregrine::internal { + +// A tcp socket based channel: reliable, stream. +// It is thread-compatible but not thread-safe. +class TcpChannel final : public Channel { + public: + // Constructor. + explicit TcpChannel(std::unique_ptr socket) + : socket_(std::move(socket)) { + DCHECK_NE(socket_, nullptr); + } + + // Returns the channel type. + constexpr ChannelType Type() const override { + return ChannelType::kReliableStream; + } + + // Writes a buffer of `len` bytes to the channel. + // Returns the number of bytes actually written if successful. Zero byte means + // no data has been written due to non-error reasons. Returns -1 on error. + ssize_t Write(const Byte* buf, size_t len) override { + DCHECK_NE(buf, nullptr); + DCHECK_GE(len, 1); + return socket_->Send(buf, len); + } + + // Writes a number of buffers described by the `iovecs` to the channel. + // Returns the number of bytes actually written if successful. Zero byte means + // no data has been written due to non-error reasons. Returns -1 on error. + ssize_t WriteV(absl::Span iovecs) override; + + // Reads exactly `len` bytes of data from the channel into the `buf`. + // Returns the number of bytes actually read if successful. Returns 0 if + // the peer side has closed the connection. Returns -1 on error. + ssize_t Read(Byte* buf, size_t len) override { + DCHECK_NE(buf, nullptr); + DCHECK_GE(len, 1); + return socket_->Recv(buf, len); + } + + // Reads exactly `length(iovecs)` bytes of data from the channel into the + // buffers. Returns the number of bytes actually read if successful. + // Returns 0 if the peer side has closed the connection. Returns -1 on error. + ssize_t ReadV(absl::Span iovecs) override; + + // Shuts down the channel. After the channel is shutdown, write calls will + // return -1, so no more data can be injected into the channel. Read calls + // will continue to read the remaining data in the channel, if any. + void Shutdown() override { socket_->Shutdown(); } + + // Returns a string representation for the channel. + std::string ToString() const override; + + private: + std::unique_ptr socket_; +}; + +} // namespace peregrine::internal + +#endif // THIRD_PARTY_PEREGRINE_SRC_INTERNAL_CHANNEL_CHANNEL_TCP_H_ diff --git a/tpu_raiden/transport/peregrine/src/internal/channel/channel_test.cc b/tpu_raiden/transport/peregrine/src/internal/channel/channel_test.cc new file mode 100644 index 00000000..06f8ce45 --- /dev/null +++ b/tpu_raiden/transport/peregrine/src/internal/channel/channel_test.cc @@ -0,0 +1,161 @@ +// Copyright 2026 Google LLC. +// +// 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 +// +// 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. + +#include "tpu_raiden/transport/peregrine/src/internal/channel/channel.h" + +#include + +#include +#include +#include +#include +#include + +#include +#include +#include "absl/log/check.h" +#include "absl/log/log.h" +#include "absl/strings/str_format.h" +#include "absl/types/span.h" +#include "tpu_raiden/transport/peregrine/src/internal/base/types.h" +#include "tpu_raiden/transport/peregrine/src/internal/channel/channel_test_util.h" +#include "tpu_raiden/transport/peregrine/src/util/util.h" + +namespace peregrine::internal::testing { +namespace { + +using ::testing::Eq; +using ::testing::Ne; +using ::testing::Pointwise; +using ::testing::Values; + +using Param = std::tuple; + +std::string FamilyToString(const int family) { + if (family == AF_INET) return "IPv4"; + if (family == AF_INET6) return "IPv6"; + return ""; +} + +std::string ToString(const ::testing::TestParamInfo& info) { + const TestChannelType type = std::get<0>(info.param); + const int family = std::get<1>(info.param); + const size_t size = std::get<2>(info.param); + return absl::StrFormat("%s_%s_Size_%zu", ToString(type), + FamilyToString(family), size); +} + +class ChannelTest : public ::testing::TestWithParam { + protected: + ChannelTest() + : size_(std::get<2>(GetParam())), + part_(size_ / 4), + src_(size_), + dst_(size_, 0) { + util::RandomNonZero(absl::MakeSpan(src_)); + DCHECK_NE(src_.data(), dst_.data()); + + CHECK_EQ(size_, 4 * part_); + src_iov_ = {src_.data() + 0 * part_, part_}; + src_iovs_ = {{src_.data() + 1 * part_, part_}, + {src_.data() + 2 * part_, part_}, + {src_.data() + 3 * part_, part_}}; + dst_iov_ = {dst_.data() + 0 * part_, part_}; + dst_iovs_ = {{dst_.data() + 1 * part_, part_}, + {dst_.data() + 2 * part_, 2 * part_}}; + } + + protected: + const size_t size_; + const size_t part_; + std::vector src_; + std::vector dst_; + IoVec src_iov_; + IoVec dst_iov_; + std::vector src_iovs_; + std::vector dst_iovs_; +}; + +INSTANTIATE_TEST_SUITE_P( + , ChannelTest, + Values(Param{TestChannelType::kTcp, AF_INET, /*size=*/1UL << 20}, + Param{TestChannelType::kTcp, AF_INET6, /*size=*/1UL << 20}, + Param{TestChannelType::kMemStream, 0, /*size=*/1UL << 20}, + Param{TestChannelType::kUdp, AF_INET, /*size=*/1UL << 10}, + Param{TestChannelType::kUdp, AF_INET6, /*size=*/1UL << 10}, + Param{TestChannelType::kMemMsg, 0, /*size=*/1UL << 10}), + ToString); + +TEST_P(ChannelTest, ReadWrite) { + const auto param = GetParam(); + const TestChannelType type = std::get<0>(param); + const int family = std::get<1>(param); + const auto chs = CreateTestChannelPair(type, family, /*error_rate=*/0); + Channel* sndr = chs.sndr.get(); + Channel* rcvr = chs.rcvr.get(); + + // Precondition: dst is different from src_. + ASSERT_THAT(dst_, Pointwise(Ne(), src_)); + + // Send to one channel a number of times. + EXPECT_EQ(sndr->Write((Byte*)src_iov_.iov_base, src_iov_.iov_len), part_); + EXPECT_EQ(sndr->WriteV(src_iovs_), size_ - part_); + + // Receive from the other channel in a different way. + EXPECT_EQ(rcvr->Read((Byte*)dst_iov_.iov_base, dst_iov_.iov_len), part_); + EXPECT_EQ(rcvr->ReadV(absl::MakeSpan(dst_iovs_)), size_ - part_); + + // Shutdown the channels and verify post-shutdown behavior. + sndr->Shutdown(); + rcvr->Shutdown(); + constexpr size_t kLen = 1; + EXPECT_EQ(sndr->Write(src_.data(), kLen), -1); + EXPECT_EQ(rcvr->Read(dst_.data(), kLen), 0); + + // Check that the data read is the same as written. + EXPECT_THAT(dst_, Pointwise(Eq(), src_)); + + LOG(INFO) << *sndr; + LOG(INFO) << *rcvr; +} + +TEST(UnreliableMessageChannelTest, ErrorRate) { + constexpr int kErrorRate = 30; // percentage + ConnectedChannelPair mem = CreateMemMsgChannelPair(kErrorRate); + Channel* sndr = mem.sndr.get(); + Channel* rcvr = mem.rcvr.get(); + + constexpr int kNumMessages = 1000; + constexpr size_t kMsgSize = 128; + std::vector src(kMsgSize, 1); + std::vector sink(kMsgSize, 0); + + int errors = 0; + for (int i = 0; i < kNumMessages; ++i) { + ASSERT_EQ(sndr->Write(src.data(), kMsgSize), kMsgSize); + if (rcvr->Read(sink.data(), kMsgSize) != kMsgSize) ++errors; + } + for (int i = 0; i < kNumMessages; ++i) { + ASSERT_EQ(sndr->WriteV({{src.data(), kMsgSize}}), kMsgSize); + IoVec iovs[] = {{sink.data(), kMsgSize}}; + if (rcvr->ReadV(iovs) != kMsgSize) ++errors; + } + + const double actual = 100.0 * errors / (2 * kNumMessages); + LOG(INFO) << "Actual error rate: " << actual << "%"; + EXPECT_NEAR(actual, kErrorRate, 10.0); +} + +} // namespace +} // namespace peregrine::internal::testing diff --git a/tpu_raiden/transport/peregrine/src/internal/channel/channel_test_util.cc b/tpu_raiden/transport/peregrine/src/internal/channel/channel_test_util.cc new file mode 100644 index 00000000..745121fc --- /dev/null +++ b/tpu_raiden/transport/peregrine/src/internal/channel/channel_test_util.cc @@ -0,0 +1,88 @@ +// Copyright 2026 Google LLC. +// +// 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 +// +// 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. + +#include "tpu_raiden/transport/peregrine/src/internal/channel/channel_test_util.h" + +#include + +#include +#include +#include + +#include "absl/log/check.h" +#include "tpu_raiden/transport/peregrine/src/internal/channel/channel_msg.h" +#include "tpu_raiden/transport/peregrine/src/internal/channel/channel_stream.h" +#include "tpu_raiden/transport/peregrine/src/internal/channel/channel_util.h" +#include "tpu_raiden/transport/peregrine/src/internal/channel/pipe.h" +#include "tpu_raiden/transport/peregrine/src/internal/socket/socket_test_util.h" + +namespace peregrine::internal::testing { + +std::string ToString(const TestChannelType t) { + switch (t) { + case TestChannelType::kTcp: + return "TcpChannel"; + case TestChannelType::kUdp: + return "UdpChannel"; + case TestChannelType::kMemStream: + return "MemStreamChannel"; + case TestChannelType::kMemMsg: + return "MemMessageChannel"; + } +} + +ConnectedChannelPair CreateTcpChannelPair(const int family) { + auto [sa, sb] = CreateTcpSocketPair(family); + DCHECK_NE(sa, nullptr); + DCHECK_NE(sb, nullptr); + return {CreateTcpChannel(std::move(sa)), CreateTcpChannel(std::move(sb))}; +} + +ConnectedChannelPair CreateUdpChannelPair(const int family) { + auto [sa, sb] = CreateUdpSocketPair(family); + DCHECK_NE(sa, nullptr); + DCHECK_NE(sb, nullptr); + return {CreateUdpChannel(std::move(sa)), CreateUdpChannel(std::move(sb))}; +} + +ConnectedChannelPair CreateMemStreamChannelPair(const int error_rate) { + auto [pipe_a, pipe_b] = BidiPipe::Create(); + auto a = std::make_unique(pipe_a, error_rate); + auto b = std::make_unique(pipe_b, error_rate); + return {std::move(a), std::move(b)}; +} + +ConnectedChannelPair CreateMemMsgChannelPair(const int error_rate) { + auto [pipe_a, pipe_b] = BidiPipe::Create(); + auto a = std::make_unique(pipe_a, error_rate); + auto b = std::make_unique(pipe_b, error_rate); + return {std::move(a), std::move(b)}; +} + +ConnectedChannelPair CreateTestChannelPair(const TestChannelType type, + const int family, + const int error_rate) { + switch (type) { + case TestChannelType::kTcp: + return CreateTcpChannelPair(family); + case TestChannelType::kUdp: + return CreateUdpChannelPair(family); + case TestChannelType::kMemStream: + return CreateMemStreamChannelPair(error_rate); + case TestChannelType::kMemMsg: + return CreateMemMsgChannelPair(error_rate); + } +} + +} // namespace peregrine::internal::testing diff --git a/tpu_raiden/transport/peregrine/src/internal/channel/channel_test_util.h b/tpu_raiden/transport/peregrine/src/internal/channel/channel_test_util.h new file mode 100644 index 00000000..f5fe4356 --- /dev/null +++ b/tpu_raiden/transport/peregrine/src/internal/channel/channel_test_util.h @@ -0,0 +1,62 @@ +// Copyright 2026 Google LLC. +// +// 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 +// +// 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. + +#ifndef THIRD_PARTY_PEREGRINE_SRC_INTERNAL_CHANNEL_CHANNEL_TEST_UTIL_H_ +#define THIRD_PARTY_PEREGRINE_SRC_INTERNAL_CHANNEL_CHANNEL_TEST_UTIL_H_ + +#include + +#include +#include + +#include "tpu_raiden/transport/peregrine/src/internal/channel/channel.h" + +namespace peregrine::internal::testing { + +enum class TestChannelType { + kTcp, + kUdp, + kMemStream, + kMemMsg, +}; + +// Returns a string representation for the test channel type. +std::string ToString(TestChannelType t); + +// A pair of connected channels. +struct ConnectedChannelPair final { + std::unique_ptr sndr; + std::unique_ptr rcvr; +}; + +// Creates a tcp channel pair in the given address family. +ConnectedChannelPair CreateTcpChannelPair(int family); + +// Creates a udp channel pair in the given address family. +ConnectedChannelPair CreateUdpChannelPair(int family); + +// Creates a memory stream channel pair. +ConnectedChannelPair CreateMemStreamChannelPair(int error_rate); + +// Creates a memory message channel pair. +ConnectedChannelPair CreateMemMsgChannelPair(int error_rate); + +// Creates a test channel pair with the given type and error rate. +ConnectedChannelPair CreateTestChannelPair(TestChannelType type, + int family = AF_INET, + int error_rate = 0); + +} // namespace peregrine::internal::testing + +#endif // THIRD_PARTY_PEREGRINE_SRC_INTERNAL_CHANNEL_CHANNEL_TEST_UTIL_H_ diff --git a/tpu_raiden/transport/peregrine/src/internal/channel/channel_types.h b/tpu_raiden/transport/peregrine/src/internal/channel/channel_types.h new file mode 100644 index 00000000..bd1a93a0 --- /dev/null +++ b/tpu_raiden/transport/peregrine/src/internal/channel/channel_types.h @@ -0,0 +1,51 @@ +// Copyright 2026 Google LLC. +// +// 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 +// +// 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. + +#ifndef THIRD_PARTY_PEREGRINE_SRC_INTERNAL_CHANNEL_CHANNEL_TYPES_H_ +#define THIRD_PARTY_PEREGRINE_SRC_INTERNAL_CHANNEL_CHANNEL_TYPES_H_ + +#include + +namespace peregrine::internal { + +// An enum that specifies the type of a channel. +enum class ChannelType : uint8_t { + kReliableStream, // e.g., tcp + kReliableMessage, // e.g., rdma + kUnreliableMessage, // e.g., udp +}; + +// Returns true iff the channel type is reliable stream. +constexpr bool IsReliableStream(ChannelType t) { + return t == ChannelType::kReliableStream; +} + +// Returns true iff the channel type is reliable message. +constexpr bool IsReliableMessage(ChannelType t) { + return t == ChannelType::kReliableMessage; +} + +// Returns true iff the channel type is unreliable message. +constexpr bool IsUnreliableMessage(ChannelType t) { + return t == ChannelType::kUnreliableMessage; +} + +// Returns true iff the channel type is reliable or unreliable message. +constexpr bool IsMessageChannel(ChannelType t) { + return IsReliableMessage(t) || IsUnreliableMessage(t); +} + +} // namespace peregrine::internal + +#endif // THIRD_PARTY_PEREGRINE_SRC_INTERNAL_CHANNEL_CHANNEL_TYPES_H_ diff --git a/tpu_raiden/transport/peregrine/src/internal/channel/channel_udp.cc b/tpu_raiden/transport/peregrine/src/internal/channel/channel_udp.cc new file mode 100644 index 00000000..ed7b9971 --- /dev/null +++ b/tpu_raiden/transport/peregrine/src/internal/channel/channel_udp.cc @@ -0,0 +1,62 @@ +// Copyright 2026 Google LLC. +// +// 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 +// +// 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. + +#include "tpu_raiden/transport/peregrine/src/internal/channel/channel_udp.h" + +#include + +#include "absl/base/optimization.h" +#include "absl/log/check.h" +#include "absl/strings/str_cat.h" +#include "absl/types/span.h" +#include "tpu_raiden/transport/peregrine/src/internal/base/types.h" +#include "tpu_raiden/transport/peregrine/src/internal/util/util.h" + +namespace peregrine::internal { + +ssize_t UdpChannel::WriteV(const absl::Span iovecs) { + DCHECK(IsValid(iovecs)); + DCHECK_GE(iovecs.size(), 1); + DCHECK_LE(iovecs.size(), IOV_MAX); + DCHECK_GE(TotalLength(iovecs), 1); + + if ABSL_PREDICT_FALSE (iovecs.size() == 1) { + const auto [buf, len] = BufLen(iovecs[0]); + return socket_->Send(buf, len); + } else { + DCHECK_GE(iovecs.size(), 2); + return socket_->SendV(iovecs); + } +} + +ssize_t UdpChannel::ReadV(absl::Span iovecs) { + DCHECK(IsValid(iovecs)); + DCHECK_GE(iovecs.size(), 1); + DCHECK_LE(iovecs.size(), IOV_MAX); + DCHECK_GE(TotalLength(iovecs), 1); + + if ABSL_PREDICT_FALSE (iovecs.size() == 1) { + const auto [buf, len] = BufLen(iovecs[0]); + return socket_->Recv(buf, len); + } else { + DCHECK_GE(iovecs.size(), 2); + return socket_->RecvV(iovecs); + } +} + +std::string UdpChannel::ToString() const { + return absl::StrCat("UdpChannel: ", socket_->ToString()); +} + +} // namespace peregrine::internal diff --git a/tpu_raiden/transport/peregrine/src/internal/channel/channel_udp.h b/tpu_raiden/transport/peregrine/src/internal/channel/channel_udp.h new file mode 100644 index 00000000..09f576f5 --- /dev/null +++ b/tpu_raiden/transport/peregrine/src/internal/channel/channel_udp.h @@ -0,0 +1,91 @@ +// Copyright 2026 Google LLC. +// +// 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 +// +// 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. + +#ifndef THIRD_PARTY_PEREGRINE_SRC_INTERNAL_CHANNEL_CHANNEL_UDP_H_ +#define THIRD_PARTY_PEREGRINE_SRC_INTERNAL_CHANNEL_CHANNEL_UDP_H_ + +#include + +#include +#include +#include +#include + +#include "absl/log/check.h" +#include "absl/types/span.h" +#include "tpu_raiden/transport/peregrine/src/internal/base/types.h" +#include "tpu_raiden/transport/peregrine/src/internal/channel/channel.h" +#include "tpu_raiden/transport/peregrine/src/internal/channel/channel_types.h" +#include "tpu_raiden/transport/peregrine/src/internal/socket/socket_udp.h" + +namespace peregrine::internal { + +// A udp socket based channel: unreliable, message. +// It is thread-compatible but not thread-safe. +class UdpChannel final : public Channel { + public: + // Constructor. + explicit UdpChannel(std::unique_ptr socket) + : socket_(std::move(socket)) { + DCHECK_NE(socket_, nullptr); + } + + // Returns the channel type. + constexpr ChannelType Type() const override { + return ChannelType::kUnreliableMessage; + } + + // Writes a buffer of `len` bytes to the channel. + // Returns the number of bytes actually written if successful. Zero byte means + // no data has been written due to non-error reasons. Returns -1 on error. + ssize_t Write(const Byte* buf, size_t len) override { + DCHECK_NE(buf, nullptr); + DCHECK_GE(len, 1); + return socket_->Send(buf, len); + } + + // Writes data from the `iovecs` buffers to the channel. + // Returns the number of bytes actually written if successful. Zero byte means + // no data has been written due to non-error reasons. Returns -1 on error. + ssize_t WriteV(absl::Span iovecs) override; + + // Reads one message of up to `len` bytes from the channel into the `buf`. + // Returns the number of bytes actually read if successful. Returns 0 if the + // received packet has no payload. Returns -1 on error. + ssize_t Read(Byte* buf, size_t len) override { + DCHECK_NE(buf, nullptr); + DCHECK_GE(len, 1); + return socket_->Recv(buf, len); + } + + // Reads one message of up to `length(iovecs)` bytes from the channel into + // the buffers. Returns the number of bytes actually read if successful. + // Returns 0 if the received packet has no payload. Returns -1 on error. + ssize_t ReadV(absl::Span iovecs) override; + + // Shuts down the channel. After the channel is shutdown, write calls will + // return -1, so no more data can be injected into the channel. Read calls + // will continue to read the remaining data in the channel, if any. + void Shutdown() override { socket_->Shutdown(); } + + // Returns a string representation for the channel. + std::string ToString() const override; + + private: + std::unique_ptr socket_; +}; + +} // namespace peregrine::internal + +#endif // THIRD_PARTY_PEREGRINE_SRC_INTERNAL_CHANNEL_CHANNEL_UDP_H_ diff --git a/tpu_raiden/transport/peregrine/src/internal/channel/channel_util.cc b/tpu_raiden/transport/peregrine/src/internal/channel/channel_util.cc new file mode 100644 index 00000000..3aec00ea --- /dev/null +++ b/tpu_raiden/transport/peregrine/src/internal/channel/channel_util.cc @@ -0,0 +1,48 @@ +// Copyright 2026 Google LLC. +// +// 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 +// +// 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. + +#include "tpu_raiden/transport/peregrine/src/internal/channel/channel_util.h" + +#include +#include + +#include "absl/log/check.h" +#include "tpu_raiden/transport/peregrine/src/internal/base/endpoint.h" +#include "tpu_raiden/transport/peregrine/src/internal/channel/channel.h" +#include "tpu_raiden/transport/peregrine/src/internal/socket/connector.h" +#include "tpu_raiden/transport/peregrine/src/internal/socket/socket_tcp.h" + +namespace peregrine::internal { + +Channels Create(const Endpoint& peer, const int n) { + DCHECK(peer.IsValid()); + DCHECK_GE(n, 1); + + Channels chs; + chs.reserve(n); + for (int i = 0; i < 2 * n; ++i) { + std::unique_ptr socket = TcpConnector::Create(peer); + if (socket == nullptr) continue; + DCHECK(socket->IsBlocking()); + + std::unique_ptr ch = CreateTcpChannel(std::move(socket)); + DCHECK_NE(ch, nullptr); + + chs.emplace_back(std::move(ch)); + if (chs.size() >= n) break; + } + return chs; +} + +} // namespace peregrine::internal diff --git a/tpu_raiden/transport/peregrine/src/internal/channel/channel_util.h b/tpu_raiden/transport/peregrine/src/internal/channel/channel_util.h new file mode 100644 index 00000000..6a9bfa12 --- /dev/null +++ b/tpu_raiden/transport/peregrine/src/internal/channel/channel_util.h @@ -0,0 +1,50 @@ +// Copyright 2026 Google LLC. +// +// 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 +// +// 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. + +#ifndef THIRD_PARTY_PEREGRINE_SRC_INTERNAL_CHANNEL_CHANNEL_UTIL_H_ +#define THIRD_PARTY_PEREGRINE_SRC_INTERNAL_CHANNEL_CHANNEL_UTIL_H_ + +#include +#include +#include + +#include "tpu_raiden/transport/peregrine/src/internal/base/endpoint.h" +#include "tpu_raiden/transport/peregrine/src/internal/channel/channel.h" +#include "tpu_raiden/transport/peregrine/src/internal/channel/channel_tcp.h" +#include "tpu_raiden/transport/peregrine/src/internal/channel/channel_udp.h" +#include "tpu_raiden/transport/peregrine/src/internal/socket/socket_tcp.h" +#include "tpu_raiden/transport/peregrine/src/internal/socket/socket_udp.h" + +namespace peregrine::internal { + +// Creates a tcp channel. +inline std::unique_ptr CreateTcpChannel( + std::unique_ptr socket) { + return std::make_unique(std::move(socket)); +} + +// Creates a udp channel. +inline std::unique_ptr CreateUdpChannel( + std::unique_ptr socket) { + return std::make_unique(std::move(socket)); +} + +using Channels = std::vector>; + +// Creates `n` channels connected to the `peer`. +Channels Create(const Endpoint& peer, int n); + +} // namespace peregrine::internal + +#endif // THIRD_PARTY_PEREGRINE_SRC_INTERNAL_CHANNEL_CHANNEL_UTIL_H_ diff --git a/tpu_raiden/transport/peregrine/src/internal/channel/channel_util_test.cc b/tpu_raiden/transport/peregrine/src/internal/channel/channel_util_test.cc new file mode 100644 index 00000000..63b14c2c --- /dev/null +++ b/tpu_raiden/transport/peregrine/src/internal/channel/channel_util_test.cc @@ -0,0 +1,85 @@ +// Copyright 2026 Google LLC. +// +// 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 +// +// 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. + +#include "tpu_raiden/transport/peregrine/src/internal/channel/channel_util.h" + +#include +#include +#include // NOLINT +#include +#include + +#include +#include "absl/log/check.h" +#include "absl/strings/str_format.h" +#include "tpu_raiden/transport/peregrine/src/internal/base/endpoint.h" +#include "tpu_raiden/transport/peregrine/src/internal/socket/acceptor.h" +#include "tpu_raiden/transport/peregrine/src/internal/socket/socket_tcp.h" +#include "tpu_raiden/transport/peregrine/src/internal/util/test_util.h" + +namespace peregrine::internal::testing { +namespace { + +using ::testing::TestParamInfo; +using ::testing::Values; + +using Param = std::tuple; + +std::string ToString(const TestParamInfo& info) { + const int family = std::get<0>(info.param); + DCHECK(family == AF_INET || family == AF_INET6); + return absl::StrFormat("IPv%d", family == AF_INET ? 4 : 6); +} + +class ChannelUtilTest : public ::testing::TestWithParam { + protected: + ChannelUtilTest() + : family_(std::get<0>(GetParam())), + local_(TestOnly_LocalEndpoint(family_, /*tcp=*/true)), + peer_(local_), + acceptor_(TcpAcceptor::Create(local_)) { + CHECK_EQ(local_, peer_); + CHECK_NE(acceptor_, nullptr); + } + + static void Accept(std::unique_ptr socket) { + auto x = std::move(socket); + CHECK_NE(x, nullptr); + } + + protected: + const int family_; + const Endpoint local_; + const Endpoint peer_; + std::unique_ptr acceptor_; +}; + +INSTANTIATE_TEST_SUITE_P(, ChannelUtilTest, + /*family=*/Values(AF_INET, AF_INET6), ToString); + +TEST_P(ChannelUtilTest, Create) { + std::jthread ta([&]() { + DCHECK(acceptor_->Socket().IsBlocking()); + acceptor_->Start(Accept); + }); + + constexpr int kNumChannels = 8; + Channels chs = Create(peer_, kNumChannels); + EXPECT_EQ(chs.size(), kNumChannels); + + acceptor_->Stop(); +} + +} // namespace +} // namespace peregrine::internal::testing diff --git a/tpu_raiden/transport/peregrine/src/internal/channel/pipe.cc b/tpu_raiden/transport/peregrine/src/internal/channel/pipe.cc new file mode 100644 index 00000000..79ac937b --- /dev/null +++ b/tpu_raiden/transport/peregrine/src/internal/channel/pipe.cc @@ -0,0 +1,46 @@ +// Copyright 2026 Google LLC. +// +// 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 +// +// 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. + +#include "tpu_raiden/transport/peregrine/src/internal/channel/pipe.h" + +#include +#include + +#include "absl/synchronization/mutex.h" + +namespace peregrine::internal::testing { + +MemPipe::MemPipe() : mu(), shutdown(false), queue() {} + +MemPipe::~MemPipe() { + absl::MutexLock lock(mu); + shutdown = true; + queue.clear(); +} + +bool MemPipe::HasData() const { return !queue.empty() || shutdown; } + +void MemPipe::Shutdown() { + absl::MutexLock lock(mu); + shutdown = true; +} + +std::pair BidiPipe::Create() { + auto p1 = std::make_shared(); + auto p2 = std::make_shared(); + return {BidiPipe{/*in_pipe=*/p1, /*out_pipe=*/p2}, + BidiPipe{/*in_pipe=*/p2, /*out_pipe=*/p1}}; +} + +} // namespace peregrine::internal::testing diff --git a/tpu_raiden/transport/peregrine/src/internal/channel/pipe.h b/tpu_raiden/transport/peregrine/src/internal/channel/pipe.h new file mode 100644 index 00000000..fc2e921c --- /dev/null +++ b/tpu_raiden/transport/peregrine/src/internal/channel/pipe.h @@ -0,0 +1,73 @@ +// Copyright 2026 Google LLC. +// +// 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 +// +// 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. + +#ifndef THIRD_PARTY_PEREGRINE_SRC_INTERNAL_CHANNEL_PIPE_H_ +#define THIRD_PARTY_PEREGRINE_SRC_INTERNAL_CHANNEL_PIPE_H_ + +#include +#include +#include + +#include "absl/base/thread_annotations.h" +#include "absl/synchronization/mutex.h" +#include "tpu_raiden/transport/peregrine/src/internal/util/test_iov.h" + +namespace peregrine::internal::testing { + +// A unidirectional pipe for in-process communication. +// It is thread-safe. +struct MemPipe final { + mutable absl::Mutex mu; + bool shutdown ABSL_GUARDED_BY(mu); + std::deque queue ABSL_GUARDED_BY(mu); + + // Constructor. + MemPipe(); + + // Destructor. + ~MemPipe(); + + // Returns true iff the pipe has data or is shutdown. + bool HasData() const ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu); + + // Shuts down the pipe so no more send/recv calls will be taken. + void Shutdown(); +}; + +// A bidirectional pipe for in-process communication. +// It is thread-safe. +class BidiPipe final { + public: + // Creates a pair of crossed bidirectional pipes. + static std::pair Create(); + + // Returns the input pipe. + std::shared_ptr InputPipe() const { return in_; } + + // Returns the output pipe. + std::shared_ptr OutputPipe() const { return out_; } + + private: + // Constructor. + BidiPipe(std::shared_ptr in, std::shared_ptr out) + : in_(std::move(in)), out_(std::move(out)) {} + + private: + std::shared_ptr in_; + std::shared_ptr out_; +}; + +} // namespace peregrine::internal::testing + +#endif // THIRD_PARTY_PEREGRINE_SRC_INTERNAL_CHANNEL_PIPE_H_ diff --git a/tpu_raiden/transport/peregrine/src/internal/util/BUILD b/tpu_raiden/transport/peregrine/src/internal/util/BUILD index 3b246989..4e081e84 100644 --- a/tpu_raiden/transport/peregrine/src/internal/util/BUILD +++ b/tpu_raiden/transport/peregrine/src/internal/util/BUILD @@ -37,6 +37,33 @@ cc_test( ], ) +cc_library( + name = "test_iov", + testonly = True, + srcs = ["test_iov.cc"], + hdrs = ["test_iov.h"], + deps = [ + ":util", + "//tpu_raiden/transport/peregrine/src/internal/base:types", + "@com_google_absl//absl/base:core_headers", + "@com_google_absl//absl/log:check", + "@com_google_absl//absl/types:span", + ], +) + +cc_test( + name = "test_iov_test", + srcs = ["test_iov_test.cc"], + deps = [ + ":test_iov", + "//tpu_raiden/transport/peregrine/src/internal/base:types", + "@com_google_absl//absl/log", + "@com_google_absl//absl/log:check", + "@com_google_absl//absl/types:span", + "@com_google_googletest//:gtest_main", + ], +) + cc_library( name = "test_util", testonly = True, diff --git a/tpu_raiden/transport/peregrine/src/internal/util/test_iov.cc b/tpu_raiden/transport/peregrine/src/internal/util/test_iov.cc new file mode 100644 index 00000000..f032d934 --- /dev/null +++ b/tpu_raiden/transport/peregrine/src/internal/util/test_iov.cc @@ -0,0 +1,59 @@ +// Copyright 2026 Google LLC. +// +// 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 +// +// 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. + +#include "tpu_raiden/transport/peregrine/src/internal/util/test_iov.h" + +#include +#include +#include +#include + +#include "absl/base/optimization.h" +#include "absl/log/check.h" +#include "absl/types/span.h" +#include "tpu_raiden/transport/peregrine/src/internal/base/types.h" +#include "tpu_raiden/transport/peregrine/src/internal/util/util.h" + +namespace peregrine::internal::testing { + +OwnedIoVec TestOnly_Linearize(const absl::Span iovecs) { + // Calculate the total length of all the iovecs. + const size_t size = TotalLength(iovecs); + if ABSL_PREDICT_FALSE (size <= 0) { + return OwnedIoVec{.data = nullptr, .size = 0}; + } + + // Allocate a buffer to hold all the data. + auto buf = std::make_unique_for_overwrite(size); + DCHECK_NE(buf, nullptr); + + // Copy the data from the iovecs into the buffer. + size_t offset = 0; + for (const auto& v : iovecs) { + void* __restrict const dst = buf.get() + offset; + const void* __restrict const src = v.iov_base; + const size_t n = v.iov_len; + DCHECK(IsValid(v)); + if (src != nullptr && n > 0) { + std::memcpy(dst, src, n); + offset += n; + } + } + DCHECK_EQ(offset, size); + + // Return the buffer, together with its ownership. + return OwnedIoVec{.data = std::move(buf), .size = size}; +} + +} // namespace peregrine::internal::testing diff --git a/tpu_raiden/transport/peregrine/src/internal/util/test_iov.h b/tpu_raiden/transport/peregrine/src/internal/util/test_iov.h new file mode 100644 index 00000000..88f928d1 --- /dev/null +++ b/tpu_raiden/transport/peregrine/src/internal/util/test_iov.h @@ -0,0 +1,40 @@ +// Copyright 2026 Google LLC. +// +// 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 +// +// 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. + +#ifndef THIRD_PARTY_PEREGRINE_SRC_INTERNAL_UTIL_TEST_IOV_H_ +#define THIRD_PARTY_PEREGRINE_SRC_INTERNAL_UTIL_TEST_IOV_H_ + +#include + +#include +#include + +#include "absl/types/span.h" +#include "tpu_raiden/transport/peregrine/src/internal/base/types.h" + +namespace peregrine::internal::testing { + +// `OwnedIoVec` describes a continuous buffer of data owned by the caller. +struct OwnedIoVec { + std::unique_ptr data; + size_t size; +}; + +// Linearizes a sequence of `IoVecs` into a newly created contiguous buffer. +// Returns the buffer with its ownership moved to the caller. +OwnedIoVec TestOnly_Linearize(absl::Span iovecs); + +} // namespace peregrine::internal::testing + +#endif // THIRD_PARTY_PEREGRINE_SRC_INTERNAL_UTIL_TEST_IOV_H_ diff --git a/tpu_raiden/transport/peregrine/src/internal/util/test_iov_test.cc b/tpu_raiden/transport/peregrine/src/internal/util/test_iov_test.cc new file mode 100644 index 00000000..5dca30f1 --- /dev/null +++ b/tpu_raiden/transport/peregrine/src/internal/util/test_iov_test.cc @@ -0,0 +1,63 @@ +// Copyright 2026 Google LLC. +// +// 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 +// +// 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. + +#include "tpu_raiden/transport/peregrine/src/internal/util/test_iov.h" + +#include +#include + +#include +#include +#include "absl/log/check.h" +#include "absl/log/log.h" +#include "absl/types/span.h" +#include "tpu_raiden/transport/peregrine/src/internal/base/types.h" + +namespace peregrine::internal::testing { +namespace { + +using ::testing::ElementsAreArray; + +class IoVecTest : public ::testing::Test { + protected: + IoVecTest() + : bytes1_(kSize, 11), + bytes2_(kSize, 22), + iov1_(bytes1_.data(), bytes1_.size()), + iov2_(bytes2_.data(), bytes2_.size()) { + CHECK_NE(bytes1_, bytes2_); + } + + protected: + static constexpr size_t kSize = 8; + std::vector bytes1_; + std::vector bytes2_; + const IoVec iov1_; + const IoVec iov2_; +}; + +TEST_F(IoVecTest, Linearize) { + const std::vector iovecs = {iov1_, iov2_}; + const OwnedIoVec owned = TestOnly_Linearize(iovecs); + EXPECT_NE(owned.data, nullptr); + EXPECT_EQ(owned.size, 2 * kSize); + + const absl::Span owned1(owned.data.get(), kSize); + const absl::Span owned2(owned.data.get() + kSize, kSize); + EXPECT_THAT(owned1, ElementsAreArray(bytes1_)); + EXPECT_THAT(owned2, ElementsAreArray(bytes2_)); +} + +} // namespace +} // namespace peregrine::internal::testing diff --git a/tpu_raiden/transport/peregrine/src/internal/util/util.h b/tpu_raiden/transport/peregrine/src/internal/util/util.h index c25ae119..24fb721a 100644 --- a/tpu_raiden/transport/peregrine/src/internal/util/util.h +++ b/tpu_raiden/transport/peregrine/src/internal/util/util.h @@ -20,8 +20,8 @@ #include #include #include +#include -#include "absl/log/check.h" #include "absl/types/span.h" #include "tpu_raiden/transport/peregrine/src/internal/base/types.h" @@ -38,10 +38,14 @@ constexpr bool IsPowerOfTwo(T n) { inline constexpr IoVec kEoF = {}; static_assert(kEoF.iov_base == nullptr && kEoF.iov_len == 0); +// Returns the `IoVec`'s buffer pointer and length as a pair. +inline std::pair BufLen(const IoVec& iov) { + return {reinterpret_cast(iov.iov_base), iov.iov_len}; +} + // Returns true iff the `IoVec` is valid. inline bool IsValid(const IoVec& v) { - DCHECK_GE(v.iov_len, 0); - return v.iov_len == 0 || v.iov_base != nullptr; + return v.iov_base != nullptr && v.iov_len > 0; } // Returns true iff all the `iovecs` are valid. diff --git a/tpu_raiden/transport/peregrine/src/internal/util/util_test.cc b/tpu_raiden/transport/peregrine/src/internal/util/util_test.cc index acdc5287..6d02905e 100644 --- a/tpu_raiden/transport/peregrine/src/internal/util/util_test.cc +++ b/tpu_raiden/transport/peregrine/src/internal/util/util_test.cc @@ -35,9 +35,8 @@ TEST(UtilTest, IsValid) { ASSERT_NE(nonnull, nullptr); EXPECT_TRUE(IsValid({.iov_base = nonnull, .iov_len = 1})); - EXPECT_TRUE(IsValid({.iov_base = nonnull, .iov_len = 0})); - - EXPECT_TRUE(IsValid({.iov_base = nullptr, .iov_len = 0})); + EXPECT_FALSE(IsValid({.iov_base = nonnull, .iov_len = 0})); + EXPECT_FALSE(IsValid({.iov_base = nullptr, .iov_len = 0})); EXPECT_FALSE(IsValid({.iov_base = nullptr, .iov_len = 1})); }