diff --git a/rosbag2_cpp/include/rosbag2_cpp/compression_options.hpp b/rosbag2_cpp/include/rosbag2_cpp/compression_options.hpp index 0f8d223f60..e1e4b65ef1 100644 --- a/rosbag2_cpp/include/rosbag2_cpp/compression_options.hpp +++ b/rosbag2_cpp/include/rosbag2_cpp/compression_options.hpp @@ -26,7 +26,7 @@ namespace rosbag2_cpp * Modes are used to specify whether to compress by individual serialized bag messages or by file. * rosbag2_cpp defaults to NONE. */ -enum class ROSBAG2_CPP_PUBLIC CompressionMode : uint32_t +enum class ROSBAG2_CPP_PUBLIC CompressionMode: uint32_t { NONE = 0, FILE, @@ -40,7 +40,8 @@ enum class ROSBAG2_CPP_PUBLIC CompressionMode : uint32_t * \param compression_mode A case insensitive string that is either "FILE" or "MESSAGE". * \return CompressionMode NONE if compression_mode is invalid. FILE or MESSAGE otherwise. */ -ROSBAG2_CPP_PUBLIC CompressionMode compression_mode_from_string(const std::string & compression_mode); +ROSBAG2_CPP_PUBLIC CompressionMode compression_mode_from_string( + const std::string & compression_mode); /** * Converts a rosbag2_cpp::CompressionMode enum into a string. @@ -50,5 +51,14 @@ ROSBAG2_CPP_PUBLIC CompressionMode compression_mode_from_string(const std::strin */ ROSBAG2_CPP_PUBLIC std::string compression_mode_to_string(CompressionMode compression_mode); +/** + * Compression options used in the writer which are passed down from the CLI in rosbag2_transport. + */ +struct CompressionOptions +{ + std::string compression_format; + CompressionMode compression_mode; +}; + } // namespace rosbag2_cpp #endif // ROSBAG2_CPP__COMPRESSION_OPTIONS_HPP_ diff --git a/rosbag2_cpp/include/rosbag2_cpp/writer.hpp b/rosbag2_cpp/include/rosbag2_cpp/writer.hpp index e4e0be926f..53d2a2bbae 100644 --- a/rosbag2_cpp/include/rosbag2_cpp/writer.hpp +++ b/rosbag2_cpp/include/rosbag2_cpp/writer.hpp @@ -20,6 +20,7 @@ #include #include +#include "rosbag2_cpp/compression_options.hpp" #include "rosbag2_cpp/converter_options.hpp" #include "rosbag2_cpp/storage_options.hpp" #include "rosbag2_cpp/visibility_control.hpp" @@ -54,13 +55,19 @@ class ROSBAG2_CPP_PUBLIC Writer final ~Writer(); /** - * Opens a new bagfile and prepare it for writing messages. The bagfile must not exist. + * Opens a new bagfile and prepares it for writing messages. + * + * The bagfile must not exist. * This must be called before any other function is used. * - * \param storage_options Options to configure the storage - * \param converter_options options to define in which format incoming messages are stored - **/ - void open(const StorageOptions & storage_options, const ConverterOptions & converter_options); + * \param storage_options Options to configure the storage. + * \param converter_options Options to define in which format incoming messages are stored. + * \param compression_options Options to configure message or bagfile compression. + */ + void open( + const StorageOptions & storage_options, + const ConverterOptions & converter_options, + const CompressionOptions & compression_options = {"", CompressionMode::NONE}); /** * Create a new topic in the underlying storage. Needs to be called for every topic used within diff --git a/rosbag2_cpp/include/rosbag2_cpp/writer_interfaces/base_writer_interface.hpp b/rosbag2_cpp/include/rosbag2_cpp/writer_interfaces/base_writer_interface.hpp index 37958930f3..7ba64f6e9f 100644 --- a/rosbag2_cpp/include/rosbag2_cpp/writer_interfaces/base_writer_interface.hpp +++ b/rosbag2_cpp/include/rosbag2_cpp/writer_interfaces/base_writer_interface.hpp @@ -30,10 +30,12 @@ namespace writer_interfaces class ROSBAG2_CPP_PUBLIC BaseWriterInterface { public: - virtual ~BaseWriterInterface() {} + virtual ~BaseWriterInterface() = default; virtual void open( - const StorageOptions & storage_options, const ConverterOptions & converter_options) = 0; + const StorageOptions & storage_options, + const ConverterOptions & converter_options, + const CompressionOptions & compression_options) = 0; virtual void reset() = 0; diff --git a/rosbag2_cpp/include/rosbag2_cpp/writers/sequential_writer.hpp b/rosbag2_cpp/include/rosbag2_cpp/writers/sequential_writer.hpp index 5b112b1a40..c89973848b 100644 --- a/rosbag2_cpp/include/rosbag2_cpp/writers/sequential_writer.hpp +++ b/rosbag2_cpp/include/rosbag2_cpp/writers/sequential_writer.hpp @@ -20,12 +20,14 @@ #include #include +#include "rosbag2_cpp/compression_options.hpp" #include "rosbag2_cpp/converter.hpp" #include "rosbag2_cpp/serialization_format_converter_factory.hpp" #include "rosbag2_cpp/storage_options.hpp" #include "rosbag2_cpp/writer_interfaces/base_writer_interface.hpp" #include "rosbag2_cpp/visibility_control.hpp" +#include "rosbag2_compression/base_compressor_interface.hpp" #include "rosbag2_storage/metadata_io.hpp" #include "rosbag2_storage/storage_factory.hpp" #include "rosbag2_storage/storage_factory_interface.hpp" @@ -71,7 +73,9 @@ class ROSBAG2_CPP_PUBLIC SequentialWriter * \param converter_options options to define in which format incoming messages are stored **/ void open( - const StorageOptions & storage_options, const ConverterOptions & converter_options) override; + const StorageOptions & storage_options, + const ConverterOptions & converter_options, + const CompressionOptions & compression_options) override; void reset() override; @@ -102,6 +106,29 @@ class ROSBAG2_CPP_PUBLIC SequentialWriter */ void write(std::shared_ptr message) override; +protected: + /** + * Initialize the compressor. + * + * \throws runtime_error If the compression implementation does not exist. + */ + virtual void init_compression(const CompressionOptions & compression_options); + + /** + * Compress the most recent file and update the metadata file path. + */ + virtual void compress_last_file(); + + /** + * Checks if the compression by message option is specified and a compressor exists. + * + * If the above conditions are satisfied, compresses the serialized bag message. + * + * \param message The message to compress. + * \return True if compression occurred, false otherwise. + */ + virtual void compress_message(std::shared_ptr message); + private: std::string base_folder_; std::unique_ptr storage_factory_; @@ -109,6 +136,7 @@ class ROSBAG2_CPP_PUBLIC SequentialWriter std::shared_ptr storage_; std::unique_ptr metadata_io_; std::unique_ptr converter_; + std::unique_ptr compressor_; // Used in bagfile splitting; specifies the best-effort maximum sub-section of a bagfile in bytes. uint64_t max_bagfile_size_; @@ -118,6 +146,11 @@ class ROSBAG2_CPP_PUBLIC SequentialWriter rosbag2_storage::BagMetadata metadata_; + // Used in invoking compression + rosbag2_cpp::CompressionMode compression_mode_; + + bool should_compress_last_file_; + // Closes the current backed storage and opens the next bagfile. void split_bagfile(); diff --git a/rosbag2_cpp/src/rosbag2_cpp/writer.cpp b/rosbag2_cpp/src/rosbag2_cpp/writer.cpp index 639fae6d09..882f5ded9c 100644 --- a/rosbag2_cpp/src/rosbag2_cpp/writer.cpp +++ b/rosbag2_cpp/src/rosbag2_cpp/writer.cpp @@ -41,9 +41,11 @@ Writer::~Writer() } void Writer::open( - const StorageOptions & storage_options, const ConverterOptions & converter_options) + const StorageOptions & storage_options, + const ConverterOptions & converter_options, + const CompressionOptions & compression_options) { - writer_impl_->open(storage_options, converter_options); + writer_impl_->open(storage_options, converter_options, compression_options); } void Writer::create_topic(const rosbag2_storage::TopicMetadata & topic_with_type) diff --git a/rosbag2_cpp/src/rosbag2_cpp/writers/sequential_writer.cpp b/rosbag2_cpp/src/rosbag2_cpp/writers/sequential_writer.cpp index 6ca5f1a591..37dda911cb 100644 --- a/rosbag2_cpp/src/rosbag2_cpp/writers/sequential_writer.cpp +++ b/rosbag2_cpp/src/rosbag2_cpp/writers/sequential_writer.cpp @@ -15,6 +15,7 @@ #include "rosbag2_cpp/writers/sequential_writer.hpp" #include +#include #include #include #include @@ -22,10 +23,11 @@ #include #include "rcpputils/filesystem_helper.hpp" - #include "rcutils/filesystem.h" +#include "rosbag2_compression/zstd_compressor.hpp" #include "rosbag2_cpp/info.hpp" +#include "rosbag2_cpp/logging.hpp" #include "rosbag2_cpp/storage_options.hpp" namespace rosbag2_cpp @@ -47,6 +49,7 @@ std::string format_storage_uri(const std::string & base_folder, uint64_t storage } } // namespace +// TODO(piraka9011) Initialize defaults in header file instead. SequentialWriter::SequentialWriter( std::unique_ptr storage_factory, std::shared_ptr converter_factory, @@ -56,16 +59,35 @@ SequentialWriter::SequentialWriter( storage_(nullptr), metadata_io_(std::move(metadata_io)), converter_(nullptr), + compressor_{nullptr}, max_bagfile_size_(rosbag2_storage::storage_interfaces::MAX_BAGFILE_SIZE_NO_SPLIT), topics_names_to_info_(), - metadata_() -{} + metadata_(), + compression_mode_{CompressionMode::NONE}, + should_compress_last_file_{true} {} + SequentialWriter::~SequentialWriter() { reset(); } +void SequentialWriter::init_compression(const CompressionOptions & compression_options) +{ + if (compression_options.compression_mode != rosbag2_cpp::CompressionMode::NONE) { + if (compression_options.compression_format == "zstd") { + compressor_ = std::make_unique(); + } else { + std::stringstream err; + err << "Unsupported compression format " << compression_options.compression_format; + throw std::invalid_argument{err.str()}; + } + } + metadata_.compression_format = compression_options.compression_format; + metadata_.compression_mode = + rosbag2_cpp::compression_mode_to_string(compression_options.compression_mode); +} + void SequentialWriter::init_metadata() { metadata_ = rosbag2_storage::BagMetadata{}; @@ -77,10 +99,12 @@ void SequentialWriter::init_metadata() void SequentialWriter::open( const StorageOptions & storage_options, - const ConverterOptions & converter_options) + const ConverterOptions & converter_options, + const CompressionOptions & compression_options) { max_bagfile_size_ = storage_options.max_bagfile_size; base_folder_ = storage_options.uri; + compression_mode_ = compression_options.compression_mode; if (converter_options.output_serialization_format != converter_options.input_serialization_format) @@ -89,25 +113,34 @@ void SequentialWriter::open( } const auto storage_uri = format_storage_uri(base_folder_, 0); - storage_ = storage_factory_->open_read_write(storage_uri, storage_options.storage_id); if (!storage_) { throw std::runtime_error("No storage could be initialized. Abort"); } - if (max_bagfile_size_ != 0 && - max_bagfile_size_ < storage_->get_minimum_split_file_size()) + if (storage_options.max_bagfile_size != 0 && + storage_options.max_bagfile_size < storage_->get_minimum_split_file_size()) { - throw std::runtime_error( - "Invalid bag splitting size given. Please provide a different value."); + throw std::invalid_argument{ + "Invalid bag splitting size given. Please provide a different value."}; } init_metadata(); + init_compression(compression_options); } void SequentialWriter::reset() { if (!base_folder_.empty()) { + // Reset may be called before initializing the compressor (ex. bad options). + // We compress the last file only if it hasn't been compressed earlier (ex. in split_bagfile()). + if (compressor_ && compression_mode_ == CompressionMode::FILE && should_compress_last_file_) { + try { + compress_last_file(); + } catch(const std::runtime_error & e) { + ROSBAG2_CPP_LOG_WARN_STREAM("Could not compress the last bag file.\n" << e.what()); + } + } finalize_metadata(); metadata_io_->write_metadata(base_folder_, metadata_); } @@ -163,19 +196,32 @@ void SequentialWriter::remove_topic(const rosbag2_storage::TopicMetadata & topic } } +void SequentialWriter::compress_last_file() +{ + assert(compressor_ != nullptr); + metadata_.relative_file_paths.back() = + compressor_->compress_uri(metadata_.relative_file_paths.back()); +} + void SequentialWriter::split_bagfile() { + if (compression_mode_ == CompressionMode::FILE) { + compress_last_file(); + } + + // Add a check to make sure reset() does not compress the file again if we couldn't load the + // storage plugin. + should_compress_last_file_ = false; const auto storage_uri = format_storage_uri( base_folder_, metadata_.relative_file_paths.size()); storage_ = storage_factory_->open_read_write(storage_uri, metadata_.storage_identifier); - if (!storage_) { std::stringstream errmsg; errmsg << "Failed to rollover bagfile to new file: \"" << storage_uri << "\"!"; - throw std::runtime_error(errmsg.str()); } + should_compress_last_file_ = true; metadata_.relative_file_paths.push_back(storage_->get_relative_file_path()); @@ -185,6 +231,13 @@ void SequentialWriter::split_bagfile() } } +void SequentialWriter::compress_message( + std::shared_ptr message) +{ + assert(compressor_ != nullptr); + compressor_->compress_serialized_bag_message(converted_message.get()); +} + void SequentialWriter::write(std::shared_ptr message) { if (!storage_) { @@ -205,7 +258,12 @@ void SequentialWriter::write(std::shared_ptrwrite(converter_ ? converter_->convert(message) : message); + auto converted_message = converter_ ? converter_->convert(message) : message; + if (compression_mode_ == CompressionMode::MESSAGE) { + compress_message(converted_message); + } + + storage_->write(converted_message); } bool SequentialWriter::should_split_bagfile() const diff --git a/rosbag2_cpp/test/rosbag2_cpp/test_multifile_reader.cpp b/rosbag2_cpp/test/rosbag2_cpp/test_multifile_reader.cpp index 5d144ee9c0..fb1ecbd3aa 100644 --- a/rosbag2_cpp/test/rosbag2_cpp/test_multifile_reader.cpp +++ b/rosbag2_cpp/test/rosbag2_cpp/test_multifile_reader.cpp @@ -153,6 +153,7 @@ class ReaderCompressionTest : public Test rosbag2_storage::TopicMetadata{"test", "test_msgs/BasicTypes", serialization_format_}; auto message = std::make_shared(); message->topic_name = topic_metadata_.name; + ON_CALL(*storage_, read_next()).WillByDefault(Return(message)); EXPECT_CALL(*storage_factory_, open_read_only(_, _)).WillRepeatedly(Return(storage_)); EXPECT_CALL(*metadata_io_, metadata_file_exists(_)).WillRepeatedly(Return(true)); diff --git a/rosbag2_cpp/test/rosbag2_cpp/test_sequential_writer.cpp b/rosbag2_cpp/test/rosbag2_cpp/test_sequential_writer.cpp index 5b61f74467..dac0cbf31a 100644 --- a/rosbag2_cpp/test/rosbag2_cpp/test_sequential_writer.cpp +++ b/rosbag2_cpp/test/rosbag2_cpp/test_sequential_writer.cpp @@ -250,3 +250,118 @@ TEST_F(SequentialWriterTest, writer_splits_when_storage_bagfile_size_gt_max_bagf EXPECT_EQ(expected_path, path); } } + +TEST_F(SequentialWriterTest, open_throws_error_on_invalid_compression_format) { + auto sequential_writer = std::make_unique( + std::move(storage_factory_), converter_factory_, std::move(metadata_io_)); + writer_ = std::make_unique(std::move(sequential_writer)); + const auto compression_options = + rosbag2_cpp::CompressionOptions{"bad_format", rosbag2_cpp::CompressionMode::FILE}; + const auto converter_options = rosbag2_cpp::ConverterOptions{"format1", "format1"}; + EXPECT_ANY_THROW(writer_->open(storage_options_, converter_options, compression_options)); +} + +class FakeSequentialWriter : public rosbag2_cpp::writers::SequentialWriter +{ +public: + FakeSequentialWriter( + std::unique_ptr storage_factory, + std::shared_ptr converter_factory, + std::unique_ptr metadata_io) + : SequentialWriter(std::move(storage_factory), converter_factory, std::move(metadata_io)) {} + + ~FakeSequentialWriter() override = default; + + void init_compression(const rosbag2_cpp::CompressionOptions &) override + { + } + + void compress_last_file() override + { + compress_file_call_counter++; + } + + void compress_message(std::shared_ptr) override + { + compress_message_call_counter++; + } + + int compress_file_call_counter = 0; + int compress_message_call_counter = 0; +}; + +TEST_F(SequentialWriterTest, writer_compresses_file_with_compression_mode_file_flag) +{ + const int message_count = 10; + const int max_bagfile_size = 5; + fake_storage_size_ = 0; + + ON_CALL(*storage_, write).WillByDefault( + [this](const std::shared_ptr &) { + fake_storage_size_ += 1; + }); + + ON_CALL(*storage_, get_bagfile_size).WillByDefault( + [this]() { + return fake_storage_size_; + }); + + ON_CALL(*storage_, get_relative_file_path).WillByDefault( + [this]() { + return fake_storage_uri_; + }); + + EXPECT_CALL(*metadata_io_, write_metadata).Times(1); + + // Intercept the metadata write so we can analyze it. + ON_CALL(*metadata_io_, write_metadata).WillByDefault( + [this](const std::string &, const rosbag2_storage::BagMetadata & metadata) { + fake_metadata_ = metadata; + }); + + auto sequential_writer = std::make_unique( + std::move(storage_factory_), converter_factory_, std::move(metadata_io_)); + + storage_options_.max_bagfile_size = max_bagfile_size; + const auto compression_options = + rosbag2_cpp::CompressionOptions{"zstd", rosbag2_cpp::CompressionMode::FILE}; + const auto converter_options = rosbag2_cpp::ConverterOptions{"format1", "format1"}; + + sequential_writer->open(storage_options_, converter_options, compression_options); + sequential_writer->create_topic({"test_topic", "test_msgs/BasicTypes", ""}); + + auto message = std::make_shared(); + message->topic_name = "test_topic"; + for (auto i = 0; i < message_count; ++i) { + sequential_writer->write(message); + } + + EXPECT_GT(sequential_writer->compress_file_call_counter, 0); +} + +TEST_F(SequentialWriterTest, writer_compresses_message_with_compression_mode_message_flag) +{ + EXPECT_CALL(*metadata_io_, write_metadata).Times(1); + + // Intercept the metadata write so we can analyze it. + ON_CALL(*metadata_io_, write_metadata).WillByDefault( + [this](const std::string &, const rosbag2_storage::BagMetadata & metadata) { + fake_metadata_ = metadata; + }); + + auto sequential_writer = std::make_unique( + std::move(storage_factory_), converter_factory_, std::move(metadata_io_)); + + const auto compression_options = + rosbag2_cpp::CompressionOptions{"zstd", rosbag2_cpp::CompressionMode::MESSAGE}; + const auto converter_options = rosbag2_cpp::ConverterOptions{"format1", "format1"}; + + sequential_writer->open(storage_options_, converter_options, compression_options); + sequential_writer->create_topic({"test_topic", "test_msgs/BasicTypes", ""}); + + auto message = std::make_shared(); + message->topic_name = "test_topic"; + sequential_writer->write(message); + + EXPECT_GT(sequential_writer->compress_message_call_counter, 0); +} diff --git a/rosbag2_transport/test/rosbag2_transport/mock_sequential_writer.hpp b/rosbag2_transport/test/rosbag2_transport/mock_sequential_writer.hpp index 5840dcf17c..17d54f7020 100644 --- a/rosbag2_transport/test/rosbag2_transport/mock_sequential_writer.hpp +++ b/rosbag2_transport/test/rosbag2_transport/mock_sequential_writer.hpp @@ -27,10 +27,12 @@ class MockSequentialWriter : public rosbag2_cpp::writer_interfaces::BaseWriterIn public: void open( const rosbag2_cpp::StorageOptions & storage_options, - const rosbag2_cpp::ConverterOptions & converter_options) override + const rosbag2_cpp::ConverterOptions & converter_options, + const rosbag2_cpp::CompressionOptions & compressionOptions) override { (void) storage_options; (void) converter_options; + (void) compressionOptions; } void reset() override {}