added seek interface (ros2#836)

- added seek(time) interface to read_only_interface
- default sqlite implementation
- added seek implementation to sequential reader (assumes multi-file is in time order)
- updated set_filter/reset_filter to be able to update in between reads without resetting to beginning
- added tests for seeking/filtering both single and multi-file bags

Signed-off-by: Sonia Jin <diegothemuich@gmail.com>

rosbag2_compression/test/rosbag2_compression/mock_storage.hpp

@@ -45,6 +45,7 @@ class MockStorage : public rosbag2_storage::storage_interfaces::ReadWriteInterfa
   MOCK_METHOD0(get_metadata, rosbag2_storage::BagMetadata());
   MOCK_METHOD0(reset_filter, void());
   MOCK_METHOD1(set_filter, void(const rosbag2_storage::StorageFilter &));
+  MOCK_METHOD1(seek, void(const rcutils_time_point_value_t &));
   MOCK_CONST_METHOD0(get_bagfile_size, uint64_t());
   MOCK_CONST_METHOD0(get_relative_file_path, std::string());
   MOCK_CONST_METHOD0(get_storage_identifier, std::string());

rosbag2_cpp/include/rosbag2_cpp/reader.hpp

@@ -169,6 +169,11 @@ class ROSBAG2_CPP_PUBLIC Reader final
    */
   void reset_filter();
 
+  /**
+   * Skip to a specific timestamp for reading.
+   */
+  void seek(const rcutils_time_point_value_t & timestamp);
+
   reader_interfaces::BaseReaderInterface & get_implementation_handle() const
   {
     return *reader_impl_;

rosbag2_cpp/include/rosbag2_cpp/reader_interfaces/base_reader_interface.hpp

@@ -18,6 +18,7 @@
 #include <memory>
 #include <vector>
 
+#include "rcutils/types.h"
 #include "rosbag2_cpp/converter_options.hpp"
 #include "rosbag2_cpp/visibility_control.hpp"
 
@@ -54,6 +55,8 @@ class ROSBAG2_CPP_PUBLIC BaseReaderInterface
   virtual void set_filter(const rosbag2_storage::StorageFilter & storage_filter) = 0;
 
   virtual void reset_filter() = 0;
+
+  virtual void seek(const rcutils_time_point_value_t & timestamp) = 0;
 };
 
 }  // namespace reader_interfaces

rosbag2_cpp/include/rosbag2_cpp/readers/sequential_reader.hpp

@@ -76,6 +76,12 @@ class ROSBAG2_CPP_PUBLIC SequentialReader
 
   void reset_filter() override;
 
+  /**
+   * seek(t) will cause subsequent reads to return messages that satisfy
+   * timestamp >= time t.
+   */
+  void seek(const rcutils_time_point_value_t & timestamp) override;
+
   /**
    * Ask whether there is another database file to read from the list of relative
    * file paths.
@@ -95,9 +101,15 @@ class ROSBAG2_CPP_PUBLIC SequentialReader
   virtual std::string get_current_uri() const;
 
 protected:
+  /**
+  * Opens the current file and sets up the filters in the new storage.
+  *
+  */
+  virtual void load_current_file();
+
   /**
   * Increment the current file iterator to point to the next file in the list of relative file
-  * paths.
+  * paths, and opens the next file by calling open_current_file()
   *
   * Expected usage:
   * if (has_next_file()) load_next_file();
@@ -144,6 +156,7 @@ class ROSBAG2_CPP_PUBLIC SequentialReader
   std::unique_ptr<Converter> converter_{};
   std::unique_ptr<rosbag2_storage::MetadataIo> metadata_io_{};
   rosbag2_storage::BagMetadata metadata_{};
+  rcutils_time_point_value_t seek_time_ = 0;
   rosbag2_storage::StorageFilter topics_filter_{};
   std::vector<rosbag2_storage::TopicMetadata> topics_metadata_{};
   std::vector<std::string> file_paths_{};  // List of database files.

rosbag2_cpp/src/rosbag2_cpp/reader.cpp

@@ -87,4 +87,9 @@ void Reader::reset_filter()
   reader_impl_->reset_filter();
 }
 
+void Reader::seek(const rcutils_time_point_value_t & timestamp)
+{
+  reader_impl_->seek(timestamp);
+}
+
 }  // namespace rosbag2_cpp

rosbag2_cpp/src/rosbag2_cpp/readers/sequential_reader.cpp

@@ -142,10 +142,6 @@ bool SequentialReader::has_next()
     bool current_storage_has_next = storage_->has_next();
     if (!current_storage_has_next && has_next_file()) {
       load_next_file();
-      storage_options_.uri = get_current_file();
-
-      storage_ = storage_factory_->open_read_only(storage_options_);
-      storage_->set_filter(topics_filter_);
       // recursively call has_next again after rollover
       return has_next();
     }
@@ -154,13 +150,18 @@ bool SequentialReader::has_next()
   throw std::runtime_error("Bag is not open. Call open() before reading.");
 }
 
+// Note: if files in the bag are not time
+// normalized, it's possible to read messages that have a timestamp
+// before the timestamp of the last read upon a file roll-over.
 std::shared_ptr<rosbag2_storage::SerializedBagMessage> SequentialReader::read_next()
 {
   if (storage_) {
     // performs rollover if necessary
-    has_next();
-    auto message = storage_->read_next();
-    return converter_ ? converter_->convert(message) : message;
+    if (has_next()) {
+      auto message = storage_->read_next();
+      return converter_ ? converter_->convert(message) : message;
+    }
+    throw std::runtime_error("Bag is at end. No next message.");
   }
   throw std::runtime_error("Bag is not open. Call open() before reading.");
 }
@@ -191,25 +192,42 @@ void SequentialReader::set_filter(
 
 void SequentialReader::reset_filter()
 {
-  topics_filter_ = rosbag2_storage::StorageFilter();
+  set_filter(rosbag2_storage::StorageFilter());
+}
+
+void SequentialReader::seek(const rcutils_time_point_value_t & timestamp)
+{
+  seek_time_ = timestamp;
   if (storage_) {
-    storage_->reset_filter();
+    // reset to the first file
+    current_file_iterator_ = file_paths_.begin();
+    load_current_file();
     return;
   }
   throw std::runtime_error(
-          "Bag is not open. Call open() before resetting filter.");
+          "Bag is not open. Call open() before seeking time.");
 }
 
 bool SequentialReader::has_next_file() const
 {
   return current_file_iterator_ + 1 != file_paths_.end();
 }
 
+void SequentialReader::load_current_file()
+{
+  preprocess_current_file();
+  storage_options_.uri = get_current_file();
+  // open and set filters
+  storage_ = storage_factory_->open_read_only(storage_options_);
+  storage_->seek(seek_time_);
+  set_filter(topics_filter_);
+}
+
 void SequentialReader::load_next_file()
 {
   assert(current_file_iterator_ != file_paths_.end());
   current_file_iterator_++;
-  preprocess_current_file();
+  load_current_file();
 }
 
 std::string SequentialReader::get_current_file() const

rosbag2_cpp/test/rosbag2_cpp/mock_storage.hpp

@@ -46,6 +46,7 @@ class MockStorage : public rosbag2_storage::storage_interfaces::ReadWriteInterfa
   MOCK_METHOD0(get_metadata, rosbag2_storage::BagMetadata());
   MOCK_METHOD0(reset_filter, void());
   MOCK_METHOD1(set_filter, void(const rosbag2_storage::StorageFilter &));
+  MOCK_METHOD1(seek, void(const rcutils_time_point_value_t &));
   MOCK_CONST_METHOD0(get_bagfile_size, uint64_t());
   MOCK_CONST_METHOD0(get_relative_file_path, std::string());
   MOCK_CONST_METHOD0(get_storage_identifier, std::string());

rosbag2_cpp/test/rosbag2_cpp/test_multifile_reader.cpp

@@ -222,3 +222,14 @@ TEST_F(MultifileReaderTest, get_all_topics_and_types_returns_from_io_metadata)
   const auto all_topics_and_types = reader_->get_all_topics_and_types();
   EXPECT_FALSE(all_topics_and_types.empty());
 }
+
+TEST_F(MultifileReaderTest, seek_bag)
+{
+  init();
+  reader_->open(default_storage_options_, {"", storage_serialization_format_});
+  EXPECT_CALL(*storage_, has_next()).Times(3).WillRepeatedly(Return(false));
+  EXPECT_CALL(*storage_, seek(_)).Times(3);
+  EXPECT_CALL(*storage_, set_filter(_)).Times(3);
+  reader_->seek(9999999999999);
+  reader_->has_next();
+}

rosbag2_cpp/test/rosbag2_cpp/test_sequential_reader.cpp

@@ -68,6 +68,7 @@ class SequentialReaderTest : public Test
 
     EXPECT_CALL(*storage_, get_all_topics_and_types())
     .Times(AtMost(1)).WillRepeatedly(Return(topics_and_types));
+    EXPECT_CALL(*storage_, has_next()).WillRepeatedly(Return(true));
     EXPECT_CALL(*storage_, read_next()).WillRepeatedly(Return(message));
 
     EXPECT_CALL(*storage_factory, open_read_only(_));
@@ -138,16 +139,14 @@ TEST_F(SequentialReaderTest, set_filter_calls_storage) {
   EXPECT_ANY_THROW(reader_->get_implementation_handle().set_filter(storage_filter));
   EXPECT_ANY_THROW(reader_->get_implementation_handle().reset_filter());
 
-  EXPECT_CALL(*storage_, set_filter(_)).Times(2);
+  EXPECT_CALL(*storage_, set_filter(_)).Times(3);
   reader_->open(default_storage_options_, {"", storage_serialization_format_});
   reader_->get_implementation_handle().set_filter(storage_filter);
   reader_->read_next();
   storage_filter.topics.clear();
   storage_filter.topics.push_back("topic2");
   reader_->get_implementation_handle().set_filter(storage_filter);
   reader_->read_next();
-
-  EXPECT_CALL(*storage_, reset_filter()).Times(1);
   reader_->get_implementation_handle().reset_filter();
   reader_->read_next();
 }

rosbag2_py/CMakeLists.txt

@@ -160,6 +160,12 @@ if(BUILD_TESTING)
       PYTHONPATH=${CMAKE_CURRENT_SOURCE_DIR}
       ${other_environment_vars}
   )
+  ament_add_pytest_test(test_sequential_reader_multiple_files_py "test/test_sequential_reader_multiple_files.py"
+    PYTHON_EXECUTABLE "${_PYTHON_EXECUTABLE}"
+    APPEND_ENV
+      PYTHONPATH=${CMAKE_CURRENT_SOURCE_DIR}
+      ${other_environment_vars}
+  )
   ament_add_pytest_test(test_sequential_writer_py
     "test/test_sequential_writer.py"
     PYTHON_EXECUTABLE "${_PYTHON_EXECUTABLE}"

rosbag2_py/resources/wbag/metadata.yaml

@@ -0,0 +1,65 @@
+rosbag2_bagfile_information:
+  version: 4
+  storage_identifier: sqlite3
+  relative_file_paths:
+    - wbag_0.db3
+    - wbag_1.db3
+    - wbag_2.db3
+    - wbag_3.db3
+    - wbag_4.db3
+  duration:
+    nanoseconds: 411
+  starting_time:
+    nanoseconds_since_epoch: 2623
+  message_count: 6074
+  topics_with_message_count:
+    - topic_metadata:
+        name: HHH
+        type: std_msgs/msg/String
+        serialization_format: cdr
+        offered_qos_profiles: ""
+      message_count: 726
+    - topic_metadata:
+        name: GGG
+        type: std_msgs/msg/String
+        serialization_format: cdr
+        offered_qos_profiles: ""
+      message_count: 731
+    - topic_metadata:
+        name: FFF
+        type: std_msgs/msg/String
+        serialization_format: cdr
+        offered_qos_profiles: ""
+      message_count: 772
+    - topic_metadata:
+        name: EEE
+        type: std_msgs/msg/String
+        serialization_format: cdr
+        offered_qos_profiles: ""
+      message_count: 804
+    - topic_metadata:
+        name: CCC
+        type: std_msgs/msg/String
+        serialization_format: cdr
+        offered_qos_profiles: ""
+      message_count: 742
+    - topic_metadata:
+        name: DDD
+        type: std_msgs/msg/String
+        serialization_format: cdr
+        offered_qos_profiles: ""
+      message_count: 753
+    - topic_metadata:
+        name: BBB
+        type: std_msgs/msg/String
+        serialization_format: cdr
+        offered_qos_profiles: ""
+      message_count: 742
+    - topic_metadata:
+        name: AAA
+        type: std_msgs/msg/String
+        serialization_format: cdr
+        offered_qos_profiles: ""
+      message_count: 804
+  compression_format: ""
+  compression_mode: ""

rosbag2_py/src/rosbag2_py/_reader.cpp

@@ -83,6 +83,11 @@ class Reader
     reader_->reset_filter();
   }
 
+  void seek(const rcutils_time_point_value_t & timestamp)
+  {
+    reader_->seek(timestamp);
+  }
+
 protected:
   std::unique_ptr<rosbag2_cpp::Reader> reader_;
 };
@@ -116,7 +121,8 @@ PYBIND11_MODULE(_reader, m) {
     "get_all_topics_and_types",
     &rosbag2_py::Reader<rosbag2_cpp::readers::SequentialReader>::get_all_topics_and_types)
   .def("set_filter", &rosbag2_py::Reader<rosbag2_cpp::readers::SequentialReader>::set_filter)
-  .def("reset_filter", &rosbag2_py::Reader<rosbag2_cpp::readers::SequentialReader>::reset_filter);
+  .def("reset_filter", &rosbag2_py::Reader<rosbag2_cpp::readers::SequentialReader>::reset_filter)
+  .def("seek", &rosbag2_py::Reader<rosbag2_cpp::readers::SequentialReader>::seek);
 
   pybind11::class_<rosbag2_py::Reader<rosbag2_compression::SequentialCompressionReader>>(
     m, "SequentialCompressionReader")
@@ -135,7 +141,10 @@ PYBIND11_MODULE(_reader, m) {
     &rosbag2_py::Reader<rosbag2_compression::SequentialCompressionReader>::set_filter)
   .def(
     "reset_filter",
-    &rosbag2_py::Reader<rosbag2_compression::SequentialCompressionReader>::reset_filter);
+    &rosbag2_py::Reader<rosbag2_compression::SequentialCompressionReader>::reset_filter)
+  .def(
+    "seek",
+    &rosbag2_py::Reader<rosbag2_compression::SequentialCompressionReader>::seek);
 
   m.def(
     "get_registered_readers",

rosbag2_py/test/test_sequential_reader.py

@@ -80,6 +80,68 @@ def test_sequential_reader():
             msg_counter += 1
 
 
+def test_sequential_reader_seek():
+    bag_path = str(Path(__file__).parent.parent / 'resources' / 'talker')
+    storage_options, converter_options = get_rosbag_options(bag_path)
+
+    reader = rosbag2_py.SequentialReader()
+    reader.open(storage_options, converter_options)
+
+    topic_types = reader.get_all_topics_and_types()
+
+    # Create a map for quicker lookup
+    type_map = {topic_types[i].name: topic_types[i].type for i in range(len(topic_types))}
+
+    # Seek No Filter
+    reader = rosbag2_py.SequentialReader()
+    reader.open(storage_options, converter_options)
+    reader.seek(1585866237113147888)
+
+    msg_counter = 5
+
+    (topic, data, t) = reader.read_next()
+    msg_type = get_message(type_map[topic])
+    msg = deserialize_message(data, msg_type)
+
+    assert isinstance(msg, Log)
+
+    (topic, data, t) = reader.read_next()
+    msg_type = get_message(type_map[topic])
+    msg = deserialize_message(data, msg_type)
+
+    isinstance(msg, String)
+    assert msg.data == f'Hello, world! {msg_counter}'
+    msg_counter += 1
+
+    # Set Filter will continue
+    storage_filter = rosbag2_py.StorageFilter(topics=['/topic'])
+    reader.set_filter(storage_filter)
+
+    (topic, data, t) = reader.read_next()
+    msg_type = get_message(type_map[topic])
+    msg = deserialize_message(data, msg_type)
+    isinstance(msg, String)
+    assert msg.data == f'Hello, world! {msg_counter}'
+
+    # Seek will keep filter
+    reader.seek(1585866239113147888)
+
+    msg_counter = 8
+
+    (topic, data, t) = reader.read_next()
+    msg_type = get_message(type_map[topic])
+    msg = deserialize_message(data, msg_type)
+    isinstance(msg, String)
+    assert msg.data == f'Hello, world! {msg_counter}'
+    msg_counter += 1
+
+    (topic, data, t) = reader.read_next()
+    msg_type = get_message(type_map[topic])
+    msg = deserialize_message(data, msg_type)
+    isinstance(msg, String)
+    assert msg.data == f'Hello, world! {msg_counter}'
+
+
 def test_plugin_list():
     reader_plugins = rosbag2_py.get_registered_readers()
     assert 'my_read_only_test_plugin' in reader_plugins