feat: Adjust prime sieve udf to calculate primes for a specified duration

Bug: b/409344921
Change-Id: I4bf716fe954d43a9fc61aba79c0d7cf137a11263
GitOrigin-RevId: aac782eb2564bcaace0ed1d7dcb51e4aac6c8149

Author: Privacy Sandbox Team

src/roma/byob/benchmark/roma_byob_rpc_factory.h

@@ -26,6 +26,7 @@
 #include "absl/status/statusor.h"
 #include "absl/synchronization/blocking_counter.h"
 #include "absl/time/time.h"
+#include "google/protobuf/util/time_util.h"
 #include "src/roma/byob/config/config.h"
 #include "src/roma/byob/interface/roma_service.h"
 #include "src/roma/byob/sample_udf/sample_udf_interface.pb.h"
@@ -57,7 +58,7 @@ std::pair<ExecutionFunc, CleanupFunc> CreateByobRpcFunc(
     Mode mode, std::atomic<std::int64_t>& completions,
     SyscallFiltering syscall_filtering, bool disable_ipc_namespace,
     absl::Duration connection_timeout, std::string_view function_name,
-    int prime_count) {
+    const std::vector<std::string>& input_args) {
   std::unique_ptr<AppService> roma_service = std::make_unique<AppService>();
   CHECK_OK(roma_service->Init(
       /*config=*/{.lib_mounts = std::string(lib_mounts),
@@ -72,8 +73,19 @@ std::pair<ExecutionFunc, CleanupFunc> CreateByobRpcFunc(
   ::privacy_sandbox::roma_byob::example::SampleRequest request;
   if (function_name == "PrimeSieve") {
     request.set_function(FUNCTION_PRIME_SIEVE);
-    if (prime_count > 0) {
-      request.set_prime_count(prime_count);
+
+    if (!input_args.empty()) {
+      int prime_count = 0;
+      absl::Duration prime_duration;
+      if (absl::SimpleAtoi(input_args[0], &prime_count)) {
+        LOG(INFO) << "Running PrimeSieve with prime count: " << prime_count;
+        request.set_prime_count(prime_count);
+      } else if (absl::ParseDuration(input_args[0], &prime_duration)) {
+        LOG(INFO) << "Running PrimeSieve for duration: " << prime_duration;
+        *request.mutable_duration() =
+            google::protobuf::util::TimeUtil::NanosecondsToDuration(
+                absl::ToInt64Nanoseconds(prime_duration));
+      }
     }
   } else if (function_name == "Clone") {
     request.set_function(FUNCTION_CLONE);

src/roma/byob/sample_udf/BUILD.bazel

@@ -36,7 +36,10 @@ proto_library(
     name = "sample_proto",
     srcs = ["sample.proto"],
     visibility = ["//src/roma/byob:__subpackages__"],
-    deps = ["//apis/privacysandbox/apis/roma/app_api/v1:options_proto"],
+    deps = [
+        "//apis/privacysandbox/apis/roma/app_api/v1:options_proto",
+        "@com_google_protobuf//:duration_proto",
+    ],
 )
 
 cc_proto_library(
@@ -106,6 +109,7 @@ sample_api = declare_roma_api(
 roma_byob_sdk(
     name = "sample_byob_sdk",
     generated_proto_path = "generated/sample.proto",
+    proto_deps = ["@com_google_protobuf//:duration_proto"],
     roma_app_api = sample_api,
     visibility = ["//visibility:public"],
 )
@@ -266,7 +270,9 @@ cc_binary(
     visibility = ["//visibility:public"],
     deps = [
         ":sample_byob_sdk_cc_proto",
+        "//src/util:duration",
         "@com_google_absl//absl/log",
+        "@com_google_absl//absl/time",
         "@com_google_protobuf//:protobuf",
     ],
 )

src/roma/byob/sample_udf/generated/sample.proto

@@ -28,6 +28,7 @@ package privacy_sandbox.roma_byob.example;
 
 
 
+import "google/protobuf/duration.proto";
 
 //
 // UDF rpc: Sample
@@ -42,6 +43,8 @@ message SampleRequest {
   FunctionType function = 1;
   // The count of prime numbers to generate.
   int32 prime_count = 2;
+  // The duration of RunPrimeSieve.
+  google.protobuf.Duration duration = 3;
 }
 
 // SampleResponse returns the function's output.

src/roma/byob/sample_udf/sample.proto

@@ -17,13 +17,15 @@ syntax = "proto3";
 package privacy_sandbox.server_common.byob;
 
 import "apis/privacysandbox/apis/roma/app_api/v1/options.proto";
+import "google/protobuf/duration.proto";
 
 service SampleService {
   option (privacysandbox.apis.roma.app_api.v1.roma_svc_annotation) = {
     name: 'Sample Server',
     roma_app_name: 'SampleServer',
     description: 'A sample server showcasing declarative APIs for BYOB',
-    cpp_namespace: 'privacy_sandbox::roma_byob::example'
+    cpp_namespace: 'privacy_sandbox::roma_byob::example',
+    proto_imports: ['google/protobuf/duration.proto'],
   };
 
   rpc Sample(SampleRequest) returns (SampleResponse) {
@@ -114,6 +116,10 @@ message SampleRequest {
   int32 prime_count = 2 [
     (privacysandbox.apis.roma.app_api.v1.roma_field_annotation) = {description: 'The count of prime numbers to generate.'}
   ];
+  google.protobuf.Duration duration = 3
+      [(privacysandbox.apis.roma.app_api.v1.roma_field_annotation) = {
+        description: 'The duration of RunPrimeSieve.'
+      }];
 }
 
 message SampleResponse {

src/roma/byob/sample_udf/sample_udf.cc

@@ -22,8 +22,11 @@
 #include <vector>
 
 #include "absl/log/log.h"
+#include "absl/time/time.h"
 #include "google/protobuf/util/delimited_message_util.h"
+#include "google/protobuf/util/time_util.h"
 #include "src/roma/byob/sample_udf/sample_udf_interface.pb.h"
+#include "src/util/duration.h"
 
 using ::google::protobuf::io::FileInputStream;
 using ::privacy_sandbox::roma_byob::example::FUNCTION_CLONE;
@@ -43,20 +46,35 @@ void RunHelloWorld(SampleResponse& bin_response) {
 void RunPrimeSieve(SampleRequest& bin_request, SampleResponse& bin_response) {
   int prime_count =
       bin_request.prime_count() > 0 ? bin_request.prime_count() : kPrimeCount;
+  int duration_ns = google::protobuf::util::TimeUtil::DurationToNanoseconds(
+      bin_request.duration());
+  absl::Duration duration = absl::Nanoseconds(duration_ns);
+  bool use_duration = duration_ns > 0;
+  double sqrt_prime_count = sqrt(prime_count);
 
   // Create a boolean array of size n+1
   std::vector<bool> primes(prime_count + 1, true);
   // Set first two values to false
   primes[0] = false;
   primes[1] = false;
   // Loop through the elements
-  for (int i = 2; i <= sqrt(prime_count); i++) {
+  privacy_sandbox::server_common::Stopwatch stopwatch;
+  for (int i = 2; i <= sqrt_prime_count; i++) {
     if (primes[i]) {
       for (int j = i * i; j <= prime_count; j += i) {
         primes[j] = false;
       }
     }
+
+    if (use_duration) {
+      if (stopwatch.GetElapsedTime() >= duration) {
+        break;
+      } else if (i + 1 > sqrt_prime_count) {
+        i = 1;  // Reset to restart the sieve if time remains
+      }
+    }
   }
+
   // Loop through the array from 2 to n
   for (int i = 2; i <= prime_count; i++) {
     if (primes[i]) {

src/roma/traffic_generator/traffic_generator.cc

@@ -250,14 +250,7 @@ TrafficGenerator::RpcFuncCreator TrafficGenerator::GetDefaultRpcFuncCreator() {
   const std::string function_name = absl::GetFlag(FLAGS_function_name);
   const std::vector<std::string> input_args = absl::GetFlag(FLAGS_input_args);
 
-  int prime_count = 0;
-  if (function_name == "PrimeSieve" && !input_args.empty() &&
-      absl::SimpleAtoi(input_args[0], &prime_count)) {
-    LOG(INFO) << "Running PrimeSieve with prime count: " << prime_count;
-  }
-
-  return [function_name, input_args,
-          prime_count](std::atomic<std::int64_t>& completions)
+  return [function_name, input_args](std::atomic<std::int64_t>& completions)
              -> std::pair<ExecutionFunc, CleanupFunc> {
     return (absl::GetFlag(FLAGS_mode) == "byob")
                ? CreateByobRpcFunc(absl::GetFlag(FLAGS_num_workers),
@@ -267,7 +260,7 @@ TrafficGenerator::RpcFuncCreator TrafficGenerator::GetDefaultRpcFuncCreator() {
                                    absl::GetFlag(FLAGS_syscall_filtering),
                                    absl::GetFlag(FLAGS_disable_ipc_namespace),
                                    absl::GetFlag(FLAGS_byob_connection_timeout),
-                                   function_name, prime_count)
+                                   function_name, input_args)
                : CreateV8RpcFunc(absl::GetFlag(FLAGS_num_workers),
                                  absl::GetFlag(FLAGS_udf_path), function_name,
                                  input_args, completions);