Add method to update Group ID of BER (#347)

ber: add method to update group id
faasm · Aug 4, 2023 · 957ae10 · 957ae10
1 parent ccf99e8
commit 957ae10
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 0 deletions.
diff --git a/include/faabric/util/batch.h b/include/faabric/util/batch.h
@@ -17,6 +17,9 @@ std::shared_ptr<faabric::BatchExecuteRequest> batchExecFactory(
 
 bool isBatchExecRequestValid(std::shared_ptr<faabric::BatchExecuteRequest> ber);
 
+void updateBatchExecGroupId(std::shared_ptr<faabric::BatchExecuteRequest> ber,
+                            int newGroupId);
+
 // ----------
 // Batch Execute Requests' Status
 // ----------

diff --git a/src/util/batch.cpp b/src/util/batch.cpp
@@ -61,6 +61,18 @@ bool isBatchExecRequestValid(std::shared_ptr<faabric::BatchExecuteRequest> ber)
     return true;
 }
 
+void updateBatchExecGroupId(std::shared_ptr<faabric::BatchExecuteRequest> ber,
+                            int newGroupId)
+{
+    ber->set_groupid(newGroupId);
+    for (auto msg : *ber->mutable_messages()) {
+        msg.set_groupid(newGroupId);
+    }
+
+    // Sanity-check in debug mode
+    assert(isBatchExecRequestValid(ber));
+}
+
 std::shared_ptr<faabric::BatchExecuteRequestStatus> batchExecStatusFactory(
   int32_t appId)
 {

diff --git a/tests/test/util/test_batch.cpp b/tests/test/util/test_batch.cpp
@@ -88,6 +88,21 @@ TEST_CASE("Test batch. exec request sanity checks")
     REQUIRE(isBerValid == isBatchExecRequestValid(ber));
 }
 
+TEST_CASE("Test updating the group ID of a BER")
+{
+    int nMessages = 4;
+    std::shared_ptr<faabric::BatchExecuteRequest> ber =
+      batchExecFactory("demo", "echo", nMessages);
+
+    // By default the BER is valid
+    REQUIRE(isBatchExecRequestValid(ber));
+
+    int newGroupId = 1337;
+    updateBatchExecGroupId(ber, newGroupId);
+    REQUIRE(isBatchExecRequestValid(ber));
+    REQUIRE(ber->groupid() == newGroupId);
+}
+
 TEST_CASE("Test BER status factory")
 {
     int appId;