record: Add limit option

Update #3665.

3 files changed, 24 insertions, 9 deletions

diff --git a/trace/record/client.h b/trace/record/client.h
index 1bf8d37..274db20 100644
--- a/trace/record/client.h
+++ b/trace/record/client.h
@@ -87,6 +87,8 @@ class Client {
 
   void Destroy();
 
+  void set_limit(uint64_t limit) { limit_ = limit; }
+
  protected:
   void Initialize(rtems_record_client_handler handler) {
     rtems_record_client_init(&base_, handler, this);
@@ -98,6 +100,7 @@ class Client {
   rtems_record_client_context base_;
   FileDescriptor input_;
   sig_atomic_t stop_ = 0;
+  uint64_t limit_ = 0;
 };
 
 #endif  // RTEMS_TOOLS_TRACE_RECORD_CLIENT_H_
diff --git a/trace/record/record-client-base.cc b/trace/record/record-client-base.cc
index d9e1eea..607a250 100644
--- a/trace/record/record-client-base.cc
+++ b/trace/record/record-client-base.cc
@@ -89,15 +89,22 @@ void FileDescriptor::Destroy() {
 }
 
 void Client::Run() {
-  while (stop_ == 0) {
-    int buf[8192];
-    ssize_t n = input_.Read(buf, sizeof(buf));
+  uint64_t todo = UINT64_MAX;
+
+  if (limit_ != 0) {
+    todo = limit_;
+  }
 
-    if (n > 0) {
-      rtems_record_client_run(&base_, buf, static_cast<size_t>(n));
-    } else {
+  while (stop_ == 0 && todo > 0) {
+    int buf[8192];
+    size_t m = std::min(sizeof(buf), todo);
+    ssize_t n = input_.Read(buf, m);
+    if (n <= 0) {
       break;
     }
+
+    rtems_record_client_run(&base_, buf, static_cast<size_t>(n));
+    todo -= static_cast<size_t>(n);
   }
 }
 
diff --git a/trace/record/record-main-lttng.cc b/trace/record/record-main-lttng.cc
index add84e1..e0e7cc5 100644
--- a/trace/record/record-main-lttng.cc
+++ b/trace/record/record-main-lttng.cc
@@ -548,7 +548,7 @@ static const struct option kLongOpts[] = {{"help", 0, NULL, 'h'},
                                           {NULL, 0, NULL, 0}};
 
 static void Usage(char** argv) {
-  std::cout << argv[0] << " [--host=HOST] [--port=PORT] [INPUT-FILE]"
+  std::cout << argv[0] << " [--host=HOST] [--port=PORT] [--limit=LIMIT] [INPUT-FILE]"
             << std::endl
             << std::endl
             << "Mandatory arguments to long options are mandatory for short "
@@ -560,6 +560,8 @@ static void Usage(char** argv) {
             << std::endl
             << "  -p, --port=PORT            the TCP port of the record server"
             << std::endl
+            << "  -l, --limit=LIMIT          limit in bytes to process"
+            << std::endl
             << "  INPUT-FILE                 the input file" << std::endl;
 }
 
@@ -570,7 +572,7 @@ int main(int argc, char** argv) {
   int opt;
   int longindex;
 
-  while ((opt = getopt_long(argc, argv, "hH:p:", &kLongOpts[0],
+  while ((opt = getopt_long(argc, argv, "hH:l:p:", &kLongOpts[0],
                             &longindex)) != -1) {
     switch (opt) {
       case 'h':
@@ -580,7 +582,10 @@ int main(int argc, char** argv) {
         host = optarg;
         break;
       case 'p':
-        port = (uint16_t)strtoul(optarg, NULL, 10);
+        port = (uint16_t)strtoul(optarg, NULL, 0);
+        break;
+      case 'l':
+        client.set_limit(strtoull(optarg, NULL, 0));
         break;
       default:
         return 1;