feat(edp-simulator)!: Specify EventGroup metadata via command-line op…

…tion (#1886)

EventGroup metadata is now supplied by the `--event-group-metadata` option instead of using the specs passed to the `--event-group-spec` option. The `--event-group-metadata-type-url` option is used to indicate the message type of the EventGroup metadata, which must be a known metadata type.

BREAKING CHANGE: The EDP simulator has a new required `--event-group-metadata` option.
BREAKING CHANGE: The EDP simulator now takes the `--known-event-group-metadata-type` option.

Closes #1880

docs/gke/simulator-deployment.md

@@ -82,12 +82,29 @@ bazel build //src/main/k8s/dev:synthetic_generator_edp_simulators.tar \
 --define image_repo_prefix=halo-cmm-demo --define image_tag=build-0001
 ```
 
-The resulting archive will contain `SyntheticEventGroupSpec` messages in text
-format under `src/main/k8s/dev/synthetic_generator_config_files/`. These can be
-replaced in order to customize the synthetic generator.
-
 Extract the generated archive to some directory.
 
+## Customize Behavior
+
+The extracted Kustomization directory will contain `SyntheticEventGroupSpec`
+messages in text format under
+`src/main/k8s/dev/synthetic_generator_config_files/`. These can be replaced in
+order to customize the synthetic generator. This directory will also contain a
+`known_event_group_metadata_type_set.pb` file which is a serialized
+FileDescriptorSet containing known types used in EventGroup metadata
+descriptors. If not using the default, replace this file with the same one used
+for the Kingdom and Reporting deployments.
+
+By default, the simulator will create EventGroups using the configured
+`SyntheticEventGroupSpec` message as the metadata. This is required for running
+the [correctness test](correctness-test.md). If the simulator is not being used
+for that test, you can opt to use your own EventGroup metadata. To do this, add
+files containing metadata messages in protobuf text format to the ConfigMap
+generator. You can then update
+`src/main/k8s/dev/synthetic_generator_edp_simulators/synthetic_generator_edp_simulator_gke.yaml`
+to modify the `--event-group-metadata` and `--event-group-metadata-type-url`
+options passed to the simulator Deployment.
+
 ## Apply K8s Kustomization
 
 From the Kustomization directory, run

src/main/k8s/dev/BUILD.bazel

@@ -397,6 +397,7 @@ kustomization_dir(
     srcs = [
         "synthetic_generator_config_files_kustomization.yaml",
         "//src/main/k8s/testing/data:synthetic_generation_specs_large",
+        "//src/main/k8s/testing/secretfiles:known_event_group_metadata_type_set",
     ],
     renames = {
         "synthetic_generator_config_files_kustomization.yaml": "kustomization.yaml",

src/main/k8s/dev/synthetic_generator_config_files_kustomization.yaml

@@ -19,3 +19,4 @@ configMapGenerator:
   - synthetic_event_group_spec_large_1.textproto
   - synthetic_event_group_spec_large_2.textproto
   - synthetic_event_group_spec_large_3.textproto
+  - known_event_group_metadata_type_set.pb

src/main/k8s/dev/synthetic_generator_edp_simulator_gke.cue

@@ -42,6 +42,7 @@ edp_simulators: {
 			_additional_args: [
 				"--population-spec=\(_populationSpec)",
 				"--event-group-spec==\(EventGroupSpec)",
+				"--event-group-metadata==\(EventGroupSpec)",
 				if (edp.supportHmss) {"--support-hmss"},
 			]
 			deployment: {

src/main/k8s/edp_simulator.cue

@@ -73,6 +73,7 @@ import "list"
 				"--kingdom-public-api-cert-host=localhost",
 				"--log-sketch-details=\(_logSketchDetails)",
 				"--health-file=\(HealthFile)",
+				"--known-event-group-metadata-type=/etc/\(#AppName)/config-files/known_event_group_metadata_type_set.pb",
 			] + _requisitionFulfillmentServiceFlags + _additional_args
 		}
 		spec: template: spec: {

src/main/k8s/local/edp_simulators.cue

@@ -89,6 +89,7 @@ edpSimulators: {
 			_additional_args: [
 				"--population-spec=\(_populationSpec)",
 				"--event-group-spec==\(edpConfig.eventGroupSpec)",
+				"--event-group-metadata==\(edpConfig.eventGroupSpec)",
 				if (edpConfig.supportHmss) {"--support-hmss"},
 			]
 

src/main/k8s/testing/secretfiles/BUILD.bazel

@@ -131,7 +131,6 @@ proto_descriptor_set(
     name = "known_event_group_metadata_type_set",
     deps = [
         "//src/main/proto/wfa/measurement/api/v2alpha/event_group_metadata/testing:simulator_synthetic_data_spec_proto",
-        "//src/main/proto/wfa/measurement/api/v2alpha/event_group_metadata/testing:test_metadata_message_proto",
     ],
 )
 

src/main/kotlin/org/wfanet/measurement/integration/common/InProcessEdpSimulator.kt

@@ -43,7 +43,6 @@ import org.wfanet.measurement.api.v2alpha.RequisitionsGrpcKt.RequisitionsCorouti
 import org.wfanet.measurement.api.v2alpha.event_group_metadata.testing.SyntheticEventGroupSpec
 import org.wfanet.measurement.api.v2alpha.event_group_metadata.testing.SyntheticPopulationSpec
 import org.wfanet.measurement.api.v2alpha.event_templates.testing.TestEvent
-import org.wfanet.measurement.api.v2alpha.populationSpec
 import org.wfanet.measurement.common.Health
 import org.wfanet.measurement.common.identity.withPrincipalName
 import org.wfanet.measurement.common.throttler.MinimumIntervalThrottler

src/main/kotlin/org/wfanet/measurement/loadtest/dataprovider/EdpSimulator.kt

@@ -69,7 +69,6 @@ import org.wfanet.measurement.api.v2alpha.EventGroupMetadataDescriptor
 import org.wfanet.measurement.api.v2alpha.EventGroupMetadataDescriptorsGrpcKt.EventGroupMetadataDescriptorsCoroutineStub
 import org.wfanet.measurement.api.v2alpha.EventGroupsGrpcKt.EventGroupsCoroutineStub
 import org.wfanet.measurement.api.v2alpha.FulfillRequisitionRequest
-import org.wfanet.measurement.api.v2alpha.FulfillRequisitionRequestKt.HeaderKt.honestMajorityShareShuffle
 import org.wfanet.measurement.api.v2alpha.FulfillRequisitionRequestKt.bodyChunk
 import org.wfanet.measurement.api.v2alpha.FulfillRequisitionRequestKt.header
 import org.wfanet.measurement.api.v2alpha.ListEventGroupsRequestKt

src/main/kotlin/org/wfanet/measurement/loadtest/dataprovider/SyntheticGeneratorEdpSimulatorRunner.kt

@@ -16,6 +16,8 @@ package org.wfanet.measurement.loadtest.dataprovider
 
 import com.google.protobuf.DescriptorProtos
 import com.google.protobuf.Descriptors
+import com.google.protobuf.DynamicMessage
+import com.google.protobuf.Message
 import com.google.protobuf.TypeRegistry
 import java.io.File
 import org.wfanet.measurement.api.v2alpha.EventGroup
@@ -48,6 +50,27 @@ class SyntheticGeneratorEdpSimulatorRunner : EdpSimulatorRunner() {
   )
   private lateinit var eventMessageDescriptorSetFiles: List<File>
 
+  @CommandLine.Option(
+    names = ["--known-event-group-metadata-type"],
+    description =
+      [
+        "File path to FileDescriptorSet containing known EventGroup metadata types.",
+        "This is in addition to standard protobuf well-known types.",
+        "Can be specified multiple times.",
+      ],
+    required = false,
+    defaultValue = "",
+  )
+  private fun setKnownEventGroupMetadataTypes(files: List<File>) {
+    knownEventGroupMetadataTypes =
+      ProtoReflection.buildFileDescriptors(
+        loadFileDescriptorSets(files),
+        WELL_KNOWN_METADATA_TYPES.asIterable(),
+      )
+  }
+
+  private lateinit var knownEventGroupMetadataTypes: List<Descriptors.FileDescriptor>
+
   @CommandLine.Option(
     names = ["--population-spec"],
     description = ["Path to SyntheticPopulationSpec message in text format."],
@@ -66,25 +89,55 @@ class SyntheticGeneratorEdpSimulatorRunner : EdpSimulatorRunner() {
   )
   private lateinit var eventGroupSpecFileByReferenceIdSuffix: Map<String, File>
 
+  @CommandLine.Option(
+    names = ["--event-group-metadata"],
+    description =
+      [
+        "Key-value pair of EventGroup reference ID suffix and file path of " +
+          "EventGroup metadata message in protobuf text format.",
+        "This can be specified multiple times.",
+      ],
+    required = true,
+  )
+  private lateinit var eventGroupMetadataFileByReferenceIdSuffix: Map<String, File>
+
+  @CommandLine.Option(
+    names = ["--event-group-metadata-type-url"],
+    description =
+      [
+        "Type URL of the EventGroup metadata message type.",
+        "This must be a known EventGroup metadata type.",
+      ],
+    required = false,
+    defaultValue = DEFAULT_METADATA_TYPE_URL,
+  )
+  private lateinit var eventGroupMetadataTypeUrl: String
+
   @CommandLine.Option(
     names = ["--support-hmss"],
     description = ["Whether to support HMSS protocol requisitions."],
   )
   private var supportHmss: Boolean = false
 
   override fun run() {
+    val typeRegistry: TypeRegistry = buildTypeRegistry()
     val syntheticPopulationSpec =
       parseTextProto(populationSpecFile, SyntheticPopulationSpec.getDefaultInstance())
     val eventGroupSpecByReferenceIdSuffix =
       eventGroupSpecFileByReferenceIdSuffix.mapValues {
         parseTextProto(it.value, SyntheticEventGroupSpec.getDefaultInstance())
       }
-    val eventMessageRegistry: TypeRegistry = buildEventMessageRegistry()
+    val metadataByReferenceIdSuffix: Map<String, Message> =
+      eventGroupMetadataFileByReferenceIdSuffix.mapValues {
+        val descriptor: Descriptors.Descriptor =
+          typeRegistry.getNonNullDescriptorForTypeUrl(eventGroupMetadataTypeUrl)
+        parseTextProto(it.value, DynamicMessage.getDefaultInstance(descriptor), typeRegistry)
+      }
     val eventMessageDescriptor: Descriptors.Descriptor =
-      eventMessageRegistry.getDescriptorForTypeUrl(syntheticPopulationSpec.eventMessageTypeUrl)
+      typeRegistry.getNonNullDescriptorForTypeUrl(syntheticPopulationSpec.eventMessageTypeUrl)
 
     val eventQuery =
-      object : SyntheticGeneratorEventQuery(syntheticPopulationSpec, eventMessageRegistry) {
+      object : SyntheticGeneratorEventQuery(syntheticPopulationSpec, typeRegistry) {
         override fun getSyntheticDataSpec(eventGroup: EventGroup): SyntheticEventGroupSpec {
           val suffix =
             EdpSimulator.getEventGroupReferenceIdSuffix(eventGroup, flags.dataProviderDisplayName)
@@ -97,28 +150,55 @@ class SyntheticGeneratorEdpSimulatorRunner : EdpSimulatorRunner() {
     run(
       eventQuery,
       EdpSimulator.buildEventTemplates(eventMessageDescriptor),
-      eventGroupSpecByReferenceIdSuffix,
-      listOf(SyntheticEventGroupSpec.getDescriptor().file),
+      metadataByReferenceIdSuffix,
+      knownEventGroupMetadataTypes,
       hmssVidIndexMap,
     )
   }
 
-  private fun buildEventMessageRegistry(): TypeRegistry {
-    val builder = TypeRegistry.newBuilder().add(TestEvent.getDescriptor())
-    if (::eventMessageDescriptorSetFiles.isInitialized) {
-      val fileDescriptorSets: List<DescriptorProtos.FileDescriptorSet> =
-        eventMessageDescriptorSetFiles.map {
-          it.inputStream().use { input -> DescriptorProtos.FileDescriptorSet.parseFrom(input) }
+  private fun buildTypeRegistry(): TypeRegistry {
+    return TypeRegistry.newBuilder()
+      .apply {
+        add(ProtoReflection.WELL_KNOWN_TYPES.flatMap { it.messageTypes })
+        add(knownEventGroupMetadataTypes.flatMap { it.messageTypes })
+        add(TestEvent.getDescriptor())
+        if (::eventMessageDescriptorSetFiles.isInitialized) {
+          add(
+            ProtoReflection.buildDescriptors(loadFileDescriptorSets(eventMessageDescriptorSetFiles))
+          )
         }
-      builder.add(ProtoReflection.buildDescriptors(fileDescriptorSets))
+      }
+      .build()
+  }
+
+  private fun loadFileDescriptorSets(
+    files: Iterable<File>
+  ): List<DescriptorProtos.FileDescriptorSet> {
+    return files.map { file ->
+      file.inputStream().use { input -> DescriptorProtos.FileDescriptorSet.parseFrom(input) }
     }
-    return builder.build()
   }
 
   companion object {
     private const val TEST_EVENT_MESSAGE_TYPE =
       "wfa.measurement.api.v2alpha.event_templates.testing.TestEvent"
+    private const val SYNTHETIC_EVENT_GROUP_SPEC_MESSAGE_TYPE =
+      "wfa.measurement.api.v2alpha.event_group_metadata.testing.SyntheticEventGroupSpec"
+    private const val DEFAULT_METADATA_TYPE_URL =
+      "${ProtoReflection.DEFAULT_TYPE_URL_PREFIX}/$SYNTHETIC_EVENT_GROUP_SPEC_MESSAGE_TYPE"
+    private val WELL_KNOWN_METADATA_TYPES =
+      ProtoReflection.WELL_KNOWN_TYPES.asSequence() + SyntheticEventGroupSpec.getDescriptor().file
+
+    init {
+      check(TestEvent.getDescriptor().fullName == TEST_EVENT_MESSAGE_TYPE)
+      check(
+        SyntheticEventGroupSpec.getDescriptor().fullName == SYNTHETIC_EVENT_GROUP_SPEC_MESSAGE_TYPE
+      )
+    }
   }
 }
 
 fun main(args: Array<String>) = commandLineMain(SyntheticGeneratorEdpSimulatorRunner(), args)
+
+private fun TypeRegistry.getNonNullDescriptorForTypeUrl(typeUrl: String): Descriptors.Descriptor =
+  checkNotNull(getDescriptorForTypeUrl(typeUrl)) { "Descriptor not found for type URL $typeUrl" }