diff --git a/src/main/docker/images.bzl b/src/main/docker/images.bzl index 1d3ee0553cb..76112bef8d4 100644 --- a/src/main/docker/images.bzl +++ b/src/main/docker/images.bzl @@ -41,6 +41,11 @@ COMMON_IMAGES = [ image = "//src/main/kotlin/org/wfanet/measurement/duchy/deploy/common/job:computations_cleaner_image", repository = _PREFIX + "/duchy/computations-cleaner", ), + struct( + name = "duchy_mill_job_scheduler_image", + image = "//src/main/kotlin/org/wfanet/measurement/duchy/deploy/common/daemon/mill:job_scheduler_image", + repository = _PREFIX + "/duchy/mill-job-scheduler", + ), struct( name = "kingdom_data_server_image", image = "//src/main/kotlin/org/wfanet/measurement/kingdom/deploy/gcloud/server:gcp_kingdom_data_server_image", @@ -127,8 +132,8 @@ GKE_IMAGES = [ repository = _PREFIX + "/duchy/requisition-fulfillment", ), struct( - name = "duchy_liquid_legions_v2_mill_daemon_image", - image = "//src/main/kotlin/org/wfanet/measurement/duchy/deploy/gcloud/daemon/mill/liquidlegionsv2:gcs_liquid_legions_v2_mill_daemon_image", + name = "duchy_liquid_legions_v2_mill_job_image", + image = "//src/main/kotlin/org/wfanet/measurement/duchy/deploy/gcloud/job/mill/liquidlegionsv2:gcs_liquid_legions_v2_mill_job_image", repository = _PREFIX + "/duchy/liquid-legions-v2-mill", ), struct( @@ -172,8 +177,8 @@ EKS_IMAGES = [ repository = _PREFIX + "/duchy/aws-requisition-fulfillment", ), struct( - name = "duchy_s3_liquid_legions_v2_mill_daemon_image", - image = "//src/main/kotlin/org/wfanet/measurement/duchy/deploy/aws/daemon/mill/liquidlegionsv2:s3_liquid_legions_v2_mill_daemon_image", + name = "duchy_s3_liquid_legions_v2_mill_job_image", + image = "//src/main/kotlin/org/wfanet/measurement/duchy/deploy/aws/job/mill/liquidlegionsv2:s3_liquid_legions_v2_mill_job_image", repository = _PREFIX + "/duchy/aws-liquid-legions-v2-mill", ), struct( @@ -196,8 +201,8 @@ LOCAL_IMAGES = [ repository = _PREFIX + "/duchy/local-herald", ), struct( - name = "forwarded_storage_liquid_legions_v2_mill_daemon_image", - image = "//src/main/kotlin/org/wfanet/measurement/duchy/deploy/common/daemon/mill/liquidlegionsv2:forwarded_storage_liquid_legions_v2_mill_daemon_image", + name = "forwarded_storage_liquid_legions_v2_mill_job_image", + image = "//src/main/kotlin/org/wfanet/measurement/duchy/deploy/common/job/mill/liquidlegionsv2:forwarded_storage_liquid_legions_v2_mill_job_image", repository = _PREFIX + "/duchy/local-liquid-legions-v2-mill", ), struct( diff --git a/src/main/k8s/base.cue b/src/main/k8s/base.cue index 9661edf2ab9..75bd92a5c2a 100644 --- a/src/main/k8s/base.cue +++ b/src/main/k8s/base.cue @@ -73,6 +73,37 @@ objects: [ for objectSet in objectSets for object in objectSet {object}] metadata: #ObjectMeta } +// K8s Role. +#Role: { + apiVersion: "rbac.authorization.k8s.io/v1" + kind: "Role" + metadata: #ObjectMeta + rules: [...{ + apiGroups?: [...string] + resources?: [...string] + verbs: [...string] + resourceNames?: [...string] + }] +} + +// K8s RoleBinding. +#RoleBinding: { + apiVersion: "rbac.authorization.k8s.io/v1" + kind: "RoleBinding" + metadata: #ObjectMeta + roleRef: { + apiGroup: string + kind: string + name: string + } + subjects: [...{ + kind: string + name: string + apiGroup?: string + namespace?: string + }] +} + #ResourceQuantity: { cpu?: string memory?: string @@ -400,6 +431,49 @@ objects: [ for objectSet in objectSets for object in objectSet {object}] tolerations: [ for _, toleration in _tolerations {toleration}] } +// K8s PodTemplateSpec. +#PodTemplateSpec: { + metadata: #ObjectMeta & { + annotations: { + "instrumentation.opentelemetry.io/inject-java": string | *"true" + } + } + spec: #PodSpec +} + +// K8s PodTemplate. +#PodTemplate: { + let Name = metadata.name + + _secretName?: string + _container: #Container & { + _javaOptions: { + heapDumpOnOutOfMemory: true + heapDumpPath: "/run/heap-dumps" + } + } + + apiVersion: "v1" + kind: "PodTemplate" + metadata: #ObjectMeta + template: #PodTemplateSpec & { + metadata: labels: { + app: "\(Name)-app" + } + spec: { + _mounts: { + if _secretName != _|_ { + "\(Name)-files": { + volume: secret: secretName: _secretName + } + } + "heap-dumps": volume: emptyDir: {} + } + _containers: "\(Name)-container": _container + } + } +} + // K8s Pod. #Pod: { apiVersion: "v1" @@ -491,16 +565,13 @@ objects: [ for objectSet in objectSets for object in objectSet {object}] selector: #LabelSelector & { matchLabels: app: _name + "-app" } - template: { + template: #PodTemplateSpec & { metadata: { labels: { app: _name + "-app" } - annotations: { - "instrumentation.opentelemetry.io/inject-java": string | *"true" - } } - spec: #PodSpec & { + spec: { _mounts: { if _secretName != _|_ { "\(_name)-files": { @@ -542,20 +613,23 @@ objects: [ for objectSet in objectSets for object in objectSet {object}] name: _name + "-cronjob" } spec: { - schedule: string + schedule: string + concurrencyPolicy?: "Allow" | "Forbid" | "Replace" + startingDeadlineSeconds?: int64 + suspend?: bool + successfulJobsHistoryLimit?: int32 & >0 + failedJobsHistoryLimit?: int32 & >0 + jobTemplate: { spec: { backoffLimit: uint | *0 - template: { - metadata: #ObjectMeta & { + template: #PodTemplateSpec & { + metadata: { labels: { app: _name + "-app" } - annotations: { - "instrumentation.opentelemetry.io/inject-java": _ | *"true" - } } - spec: #PodSpec & { + spec: { if _secretName != _|_ { _mounts: "\(_name)-files": { volume: secret: secretName: _secretName diff --git a/src/main/k8s/dev/base_gke.cue b/src/main/k8s/dev/base_gke.cue index d036920bc5d..2c2ad275734 100644 --- a/src/main/k8s/dev/base_gke.cue +++ b/src/main/k8s/dev/base_gke.cue @@ -16,18 +16,22 @@ package k8s #NetworkPolicy: { _egresses: { + // See https://cloud.google.com/kubernetes-engine/docs/how-to/network-policy#network-policy-and-workload-identity gkeMetadataServer: { to: [{ipBlock: cidr: "169.254.169.252/32"}] - ports: [ - { - protocol: "TCP" - port: 988 - }, - { - protocol: "TCP" - port: 80 - }] + ports: [{ + protocol: "TCP" + port: 988 + }] + } + gkeDataplaneV2: { + to: [{ipBlock: cidr: "169.254.169.254/32"}] + ports: [{ + protocol: "TCP" + port: 80 + }] } + openTelemetryCollector: { to: [{podSelector: matchLabels: app: "opentelemetry-collector-app"}] ports: [{ diff --git a/src/main/k8s/dev/duchy_eks.cue b/src/main/k8s/dev/duchy_eks.cue index 18f15718f32..89f8d07e448 100644 --- a/src/main/k8s/dev/duchy_eks.cue +++ b/src/main/k8s/dev/duchy_eks.cue @@ -56,7 +56,7 @@ _duchyCertName: "duchies/\(_duchyName)/certificates/\(_certificateId)" } } #Llv2MillMaxHeapSize: "1G" -#Llv2MillReplicas: 1 +#Llv2MillMaxConcurrency: 10 #HmssMillResourceRequirements: ResourceRequirements=#ResourceRequirements & { requests: { cpu: "2" @@ -89,19 +89,13 @@ _duchyCertName: "duchies/\(_duchyName)/certificates/\(_certificateId)" } #ControlServiceMaxHeapSize: "320M" -objectSets: [ - default_deny_ingress_and_egress, - duchy.deployments, - duchy.services, - duchy.networkPolicies, - duchy.cronjobs, -] +objectSets: [default_deny_ingress_and_egress] + [ for objectSet in duchy {objectSet}] duchy: #PostgresDuchy & { _imageSuffixes: { "herald-daemon": "duchy/aws-herald" "computation-control-server": "duchy/aws-computation-control" - "liquid-legions-v2-mill-daemon": "duchy/aws-liquid-legions-v2-mill" + "llv2-mill": "duchy/aws-liquid-legions-v2-mill" "hmss-mill-daemon": "duchy/aws-honest-majority-share-shuffle-mill" "requisition-fulfillment-server": "duchy/aws-requisition-fulfillment" "internal-api-server": "duchy/aws-postgres-internal-server" @@ -119,11 +113,14 @@ duchy: #PostgresDuchy & { "worker1": _worker1SystemApiTarget "worker2": _worker2SystemApiTarget } - _kingdom_system_api_target: #KingdomSystemApiTarget - _kingdom_public_api_target: #KingdomPublicApiTarget - _blob_storage_flags: #AwsS3Config.flags - _verbose_grpc_logging: "false" - _postgresConfig: #AwsPostgresConfig + _kingdom_system_api_target: #KingdomSystemApiTarget + _kingdom_public_api_target: #KingdomPublicApiTarget + _blob_storage_flags: #AwsS3Config.flags + _verbose_grpc_logging: "false" + _duchyMillParallelism: 4 + _liquidLegionsV2WorkLockDuration: "10m" + _postgresConfig: #AwsPostgresConfig + services: { "requisition-fulfillment-server": _eipAllocations: _publicApiEipAllocs "computation-control-server": _eipAllocations: _systemApiEipAllocs @@ -138,18 +135,8 @@ duchy: #PostgresDuchy & { serviceAccountName: #StorageServiceAccount } } - "liquid-legions-v2-mill-daemon-deployment": { - _workLockDuration: "10m" - _container: { - _javaOptions: maxHeapSize: #Llv2MillMaxHeapSize - resources: #Llv2MillResourceRequirements - } - spec: { - replicas: #Llv2MillReplicas - template: spec: #ServiceAccountPodSpec & #SpotVmPodSpec & { - serviceAccountName: #StorageServiceAccount - } - } + "mill-job-scheduler-deployment": { + _liquidLegionsV2MaxConcurrency: #Llv2MillMaxConcurrency } "hmss-mill-daemon-deployment": { _workLockDuration: "5m" @@ -191,4 +178,15 @@ duchy: #PostgresDuchy & { } } } + podTemplates: { + "llv2-mill": { + _container: { + _javaOptions: maxHeapSize: #Llv2MillMaxHeapSize + resources: #Llv2MillResourceRequirements + } + template: spec: #ServiceAccountPodSpec & #SpotVmPodSpec & { + serviceAccountName: #StorageServiceAccount + } + } + } } diff --git a/src/main/k8s/dev/duchy_gke.cue b/src/main/k8s/dev/duchy_gke.cue index 141c67cceb1..5c80935f4d6 100644 --- a/src/main/k8s/dev/duchy_gke.cue +++ b/src/main/k8s/dev/duchy_gke.cue @@ -57,7 +57,7 @@ _duchy_cert_name: "duchies/\(_duchy_name)/certificates/\(_certificateId)" } } #Llv2MillMaxHeapSize: "1G" -#Llv2MillReplicas: 1 +#Llv2MillMaxConcurrency: 10 #HmssMillResourceRequirements: ResourceRequirements=#ResourceRequirements & { requests: { cpu: "2" @@ -90,15 +90,7 @@ _duchy_cert_name: "duchies/\(_duchy_name)/certificates/\(_certificateId)" } #ControlServiceMaxHeapSize: "320M" -objectSets: [ - default_deny_ingress_and_egress, - duchy.serviceAccounts, - duchy.configMaps, - duchy.deployments, - duchy.services, - duchy.networkPolicies, - duchy.cronjobs, -] +objectSets: [default_deny_ingress_and_egress] + [ for objectSet in duchy {objectSet}] _cloudStorageConfig: #CloudStorageConfig & { bucket: _cloudStorageBucket @@ -119,18 +111,18 @@ duchy: #SpannerDuchy & { "worker1": _worker1SystemApiTarget "worker2": _worker2SystemApiTarget } - _kingdom_system_api_target: #KingdomSystemApiTarget - _kingdom_public_api_target: #KingdomPublicApiTarget - _blob_storage_flags: _cloudStorageConfig.flags - _verbose_grpc_logging: "false" - _duchyMillParallelism: 4 + _kingdom_system_api_target: #KingdomSystemApiTarget + _kingdom_public_api_target: #KingdomPublicApiTarget + _blob_storage_flags: _cloudStorageConfig.flags + _verbose_grpc_logging: "false" + _duchyMillParallelism: 4 + _liquidLegionsV2WorkLockDuration: "10m" - serviceAccounts: [string]: #WorkloadIdentityServiceAccount serviceAccounts: { - "\(#InternalServerServiceAccount)": { + "\(#InternalServerServiceAccount)": #WorkloadIdentityServiceAccount & { _iamServiceAccountName: "\(_duchy_name)-duchy-internal" } - "\(#StorageServiceAccount)": { + "\(#StorageServiceAccount)": #WorkloadIdentityServiceAccount & { _iamServiceAccountName: "\(_duchy_name)-duchy-storage" } } @@ -155,18 +147,8 @@ duchy: #SpannerDuchy & { serviceAccountName: #StorageServiceAccount } } - "liquid-legions-v2-mill-daemon-deployment": { - _workLockDuration: "10m" - _container: { - _javaOptions: maxHeapSize: #Llv2MillMaxHeapSize - resources: #Llv2MillResourceRequirements - } - spec: { - replicas: #Llv2MillReplicas - template: spec: #ServiceAccountPodSpec & #SpotVmPodSpec & { - serviceAccountName: #StorageServiceAccount - } - } + "mill-job-scheduler-deployment": { + _liquidLegionsV2MaxConcurrency: #Llv2MillMaxConcurrency } "hmss-mill-daemon-deployment": { _workLockDuration: "5m" @@ -205,4 +187,16 @@ duchy: #SpannerDuchy & { "requisition-fulfillment-server": _ipAddressName: _publicApiAddressName "computation-control-server": _ipAddressName: _systemApiAddressName } + + podTemplates: { + "llv2-mill": { + _container: { + _javaOptions: maxHeapSize: #Llv2MillMaxHeapSize + resources: #Llv2MillResourceRequirements + } + template: spec: #ServiceAccountPodSpec & #SpotVmPodSpec & { + serviceAccountName: #StorageServiceAccount + } + } + } } diff --git a/src/main/k8s/dev/kingdom_gke.cue b/src/main/k8s/dev/kingdom_gke.cue index dc6a2eda0b0..ec96d7c7d4f 100644 --- a/src/main/k8s/dev/kingdom_gke.cue +++ b/src/main/k8s/dev/kingdom_gke.cue @@ -40,15 +40,7 @@ _systemApiAddressName: string @tag("system_api_address_name") } } -objectSets: [ - default_deny_ingress_and_egress, - kingdom.serviceAccounts, - kingdom.configMaps, - kingdom.deployments, - kingdom.services, - kingdom.networkPolicies, - kingdom.cronjobs, -] +objectSets: [default_deny_ingress_and_egress] + [ for objectSet in kingdom {objectSet}] kingdom: #Kingdom & { _kingdom_secret_name: _secret_name diff --git a/src/main/k8s/duchy.cue b/src/main/k8s/duchy.cue index 93ce7e81dce..2a725d07df9 100644 --- a/src/main/k8s/duchy.cue +++ b/src/main/k8s/duchy.cue @@ -27,11 +27,12 @@ import ("strings") } _duchy_secret_name: string _computation_control_targets: [Name=_]: string - _deletableComputationStates: [...#TerminalComputationState] | *[] - _computationsTimeToLive: string | *"180d" - _duchyMillParallelism: uint | *2 - _kingdom_system_api_target: string - _kingdom_public_api_target: string + _deletableComputationStates: [...#TerminalComputationState] | *[] + _computationsTimeToLive: string | *"180d" + _duchyMillParallelism: uint | *2 + _liquidLegionsV2WorkLockDuration?: string + _kingdom_system_api_target: string + _kingdom_public_api_target: string _blob_storage_flags: [...string] _verbose_grpc_logging: "true" | "false" @@ -48,7 +49,8 @@ import ("strings") "async-computation-control-server": string | *"duchy/async-computation-control" "computation-control-server": string | *"duchy/computation-control" "herald-daemon": string | *"duchy/herald" - "liquid-legions-v2-mill-daemon": string | *"duchy/liquid-legions-v2-mill" + "mill-job-scheduler": string | *"duchy/mill-job-scheduler" + "llv2-mill": string | *"duchy/liquid-legions-v2-mill" "hmss-mill-daemon": string | *"duchy/honest-majority-share-shuffle-mill" "requisition-fulfillment-server": string | *"duchy/requisition-fulfillment" "computations-cleaner": string | *"duchy/computations-cleaner" @@ -138,28 +140,40 @@ import ("strings") "\(_name)-internal-api-server", ] } - "liquid-legions-v2-mill-daemon-deployment": { - _workLockDuration?: string - _container: args: [ - _duchyInternalApiTargetFlag, - _duchyInternalApiCertHostFlag, - _duchy_name_flag, - _duchy_info_config_flag, - _duchy_tls_cert_file_flag, - _duchy_tls_key_file_flag, - _duchyMillParallelismFlag, - _duchy_cert_collection_file_flag, - _duchy_cs_cert_file_flag, - _duchy_cs_key_file_flag, - _duchy_cs_cert_rename_name_flag, - _kingdom_system_api_target_flag, - _kingdom_system_api_cert_host_flag, - if (_millPollingInterval != _|_) {"--polling-interval=\(_millPollingInterval)"}, - if (_workLockDuration != _|_) {"--work-lock-duration=\(_workLockDuration)"}, - ] + _blob_storage_flags + _computation_control_target_flags - spec: template: spec: _dependencies: [ - "\(_name)-internal-api-server", "\(_name)-computation-control-server", - ] + "mill-job-scheduler-deployment": Deployment={ + let DeploymentName = Deployment.metadata.name + _liquidLegionsV2MaxConcurrency?: int32 & >0 + _container: { + _javaOptions: maxHeapSize: "40M" + resources: Resources={ + requests: { + cpu: "25m" + memory: _ | *"224Mi" + } + limits: { + memory: _ | *Resources.requests.memory + } + } + args: [ + "--deployment-name=\(DeploymentName)", + "--llv2-pod-template-name=\(_object_prefix)llv2-mill", + _duchyInternalApiTargetFlag, + _duchyInternalApiCertHostFlag, + _duchy_name_flag, + _duchy_tls_cert_file_flag, + _duchy_tls_key_file_flag, + _duchy_cert_collection_file_flag, + if (_millPollingInterval != _|_) {"--polling-delay=\(_millPollingInterval)"}, + if (_liquidLegionsV2WorkLockDuration != _|_) {"--llv2-work-lock-duration=\(_liquidLegionsV2WorkLockDuration)"}, + if (_liquidLegionsV2MaxConcurrency != _|_) {"--llv2-maximum-concurrency=\(_liquidLegionsV2MaxConcurrency)"}, + ] + } + spec: template: spec: { + _dependencies: [ + "\(_name)-internal-api-server", "\(_name)-computation-control-server", + ] + serviceAccountName: _object_prefix + "mill-job-scheduler" + } } "hmss-mill-daemon-deployment": { _workLockDuration?: string @@ -262,7 +276,7 @@ import ("strings") } } - cronjobs: [Name=_]: #CronJob & { + cronJobs: [Name=_]: #CronJob & { _unprefixed_name: strings.TrimSuffix(Name, "-cronjob") _name: _object_prefix + _unprefixed_name _secretName: _duchy_secret_name @@ -272,7 +286,7 @@ import ("strings") } } - cronjobs: { + cronJobs: { "computations-cleaner": { _container: args: [ _duchyInternalApiTargetFlag, @@ -288,6 +302,35 @@ import ("strings") } } + podTemplates: [Name=string]: #PodTemplate & { + _container: image: _images[Name] + metadata: name: _object_prefix + Name + } + podTemplates: { + "llv2-mill": { + _secretName: _duchy_secret_name + _container: args: [ + _duchyInternalApiTargetFlag, + _duchyInternalApiCertHostFlag, + _duchy_name_flag, + _duchy_info_config_flag, + _duchy_tls_cert_file_flag, + _duchy_tls_key_file_flag, + _duchyMillParallelismFlag, + _duchy_cert_collection_file_flag, + _duchy_cs_cert_file_flag, + _duchy_cs_key_file_flag, + _duchy_cs_cert_rename_name_flag, + _kingdom_system_api_target_flag, + _kingdom_system_api_cert_host_flag, + if (_liquidLegionsV2WorkLockDuration != _|_) {"--work-lock-duration=\(_liquidLegionsV2WorkLockDuration)"}, + ] + _blob_storage_flags + _computation_control_target_flags + template: spec: { + restartPolicy: "Never" + } + } + } + networkPolicies: [Name=_]: #NetworkPolicy & { _name: _object_prefix + Name } @@ -297,7 +340,8 @@ import ("strings") _app_label: _object_prefix + "internal-api-server-app" _sourceMatchLabels: [ _object_prefix + "herald-daemon-app", - _object_prefix + "liquid-legions-v2-mill-daemon-app", + _object_prefix + "mill-job-scheduler-app", + _object_prefix + "llv2-mill-app", _object_prefix + "hmss-mill-daemon-app", _object_prefix + "async-computation-control-server-app", _object_prefix + "requisition-fulfillment-server-app", @@ -348,8 +392,8 @@ import ("strings") any: {} } } - "liquid-legions-v2-mill-daemon": { - _app_label: _object_prefix + "liquid-legions-v2-mill-daemon-app" + "llv2-mill": { + _app_label: _object_prefix + "llv2-mill-app" _egresses: { // Need to send external traffic. any: {} @@ -384,4 +428,49 @@ import ("strings") serviceAccounts: [Name=string]: #ServiceAccount & { metadata: name: Name } + serviceAccounts: { + "\(_object_prefix)mill-job-scheduler": {} + } + + roles: [Name=string]: #Role & { + metadata: name: Name + } + roles: { + "\(_object_prefix)mill-job-scheduler": { + rules: [ + { + apiGroups: ["batch"] + resources: ["jobs"] + verbs: ["get", "list", "create", "delete"] + }, + { + apiGroups: [""] + resources: ["podtemplates"] + verbs: ["get"] + }, + { + apiGroups: ["apps"] + resources: ["deployments"] + verbs: ["get"] + }, + ] + } + } + + roleBindings: [Name=string]: #RoleBinding & { + metadata: name: Name + } + roleBindings: { + "\(_object_prefix)mill-job-scheduler-binding": { + roleRef: { + apiGroup: "rbac.authorization.k8s.io" + kind: "Role" + name: _object_prefix + "mill-job-scheduler" + } + subjects: [{ + kind: "ServiceAccount" + name: _object_prefix + "mill-job-scheduler" + }] + } + } } diff --git a/src/main/k8s/kingdom.cue b/src/main/k8s/kingdom.cue index 65567b72c90..f6088f5ca26 100644 --- a/src/main/k8s/kingdom.cue +++ b/src/main/k8s/kingdom.cue @@ -192,7 +192,7 @@ import ("strings") } } - cronjobs: [Name=_]: #CronJob & { + cronJobs: [Name=_]: #CronJob & { _name: strings.TrimSuffix(Name, "-cronjob") _secretName: _kingdom_secret_name _system: "kingdom" @@ -201,7 +201,7 @@ import ("strings") } } - cronjobs: { + cronJobs: { "completed-measurements-deletion": { _container: args: [ _internal_api_target_flag, diff --git a/src/main/k8s/local/duchies.cue b/src/main/k8s/local/duchies.cue index 0edf7f42146..b47fa704eb2 100644 --- a/src/main/k8s/local/duchies.cue +++ b/src/main/k8s/local/duchies.cue @@ -55,23 +55,6 @@ _duchyConfigs: { } } -let EnvVars = #EnvVarMap & { - "POSTGRES_USER": { - valueFrom: - secretKeyRef: { - name: _duchyDbSecretName - key: "username" - } - } - "POSTGRES_PASSWORD": { - valueFrom: - secretKeyRef: { - name: _duchyDbSecretName - key: "password" - } - } -} - objectSets: [ for duchy in duchies for objectSet in duchy {objectSet}] _computationControlTargets: { @@ -80,11 +63,11 @@ _computationControlTargets: { } } -_baseDuchyConfig: { +#Duchy: { _imageSuffixes: { "computation-control-server": "duchy/local-computation-control" "herald-daemon": "duchy/local-herald" - "liquid-legions-v2-mill-daemon": "duchy/local-liquid-legions-v2-mill" + "llv2-mill": "duchy/local-liquid-legions-v2-mill" "hmss-mill-daemon": "duchy/local-honest-majority-share-shuffle-mill" "requisition-fulfillment-server": "duchy/local-requisition-fulfillment" } @@ -99,43 +82,61 @@ _baseDuchyConfig: { _verbose_grpc_logging: "true" } -duchies: [ - for duchyConfig in _duchyConfigs { - if (duchyConfig.databaseType == "spanner") { - #SpannerDuchy & _baseDuchyConfig & { - _imageSuffixes: { - "internal-api-server": "duchy/local-spanner-computations" +#SpannerDuchy: { + _imageSuffixes: { + "internal-api-server": "duchy/local-spanner-computations" + } +} + +#PostgresDuchy: { + _imageSuffixes: { + "internal-api-server": "duchy/local-postgres-internal-server" + } + _postgresConfig: { + serviceName: "postgres" + password: "$(POSTGRES_PASSWORD)" + user: "$(POSTGRES_USER)" + } + deployments: { + "internal-api-server-deployment": { + let EnvVars = #EnvVarMap & { + "POSTGRES_USER": { + valueFrom: + secretKeyRef: { + name: _duchyDbSecretName + key: "username" + } } - _duchy: { - name: duchyConfig.name - protocols_setup_config: duchyConfig.protocolsSetupConfig - cs_cert_resource_name: duchyConfig.certificateResourceName - duchyKeyEncryptionKeyFile: duchyConfig.duchyKeyEncryptionKeyFile + "POSTGRES_PASSWORD": { + valueFrom: + secretKeyRef: { + name: _duchyDbSecretName + key: "password" + } } } + + _container: _envVars: EnvVars + _updateSchemaContainer: _envVars: EnvVars } - if (duchyConfig.databaseType == "postgres") { - #PostgresDuchy & _baseDuchyConfig & { - _imageSuffixes: { - "internal-api-server": "duchy/local-postgres-internal-server" - } - _duchy: { - name: duchyConfig.name - protocols_setup_config: duchyConfig.protocolsSetupConfig - cs_cert_resource_name: duchyConfig.certificateResourceName - duchyKeyEncryptionKeyFile: duchyConfig.duchyKeyEncryptionKeyFile - } - _postgresConfig: { - serviceName: "postgres" - password: "$(POSTGRES_PASSWORD)" - user: "$(POSTGRES_USER)" - } - deployments: { - "internal-api-server-deployment": { - _container: _envVars: EnvVars - _updateSchemaContainer: _envVars: EnvVars - } - } + } +} + +duchies: [ + for duchyConfig in _duchyConfigs { + { + _duchy: { + name: duchyConfig.name + protocols_setup_config: duchyConfig.protocolsSetupConfig + cs_cert_resource_name: duchyConfig.certificateResourceName + duchyKeyEncryptionKeyFile: duchyConfig.duchyKeyEncryptionKeyFile + } + + if (duchyConfig.databaseType == "spanner") { + #SpannerDuchy + } + if (duchyConfig.databaseType == "postgres") { + #PostgresDuchy } } }, diff --git a/src/main/kotlin/org/wfanet/measurement/common/k8s/BUILD.bazel b/src/main/kotlin/org/wfanet/measurement/common/k8s/BUILD.bazel index 1d8927275aa..a9686d1132f 100644 --- a/src/main/kotlin/org/wfanet/measurement/common/k8s/BUILD.bazel +++ b/src/main/kotlin/org/wfanet/measurement/common/k8s/BUILD.bazel @@ -5,8 +5,12 @@ package( ) kt_jvm_library( - name = "k8s", - srcs = glob(["*.kt"]), + name = "client", + srcs = [ + "JsonPatchOperation.kt", + "KubernetesClient.kt", + "Names.kt", + ], deps = [ "//imports/java/io/kubernetes/client", "//imports/java/io/kubernetes/client/extended", diff --git a/src/main/kotlin/org/wfanet/measurement/common/k8s/KubernetesClient.kt b/src/main/kotlin/org/wfanet/measurement/common/k8s/KubernetesClient.kt index 682e025cca6..ca1426b43f0 100644 --- a/src/main/kotlin/org/wfanet/measurement/common/k8s/KubernetesClient.kt +++ b/src/main/kotlin/org/wfanet/measurement/common/k8s/KubernetesClient.kt @@ -25,12 +25,16 @@ import io.kubernetes.client.openapi.ApiException import io.kubernetes.client.openapi.Configuration import io.kubernetes.client.openapi.JSON import io.kubernetes.client.openapi.apis.AppsV1Api +import io.kubernetes.client.openapi.apis.BatchV1Api import io.kubernetes.client.openapi.apis.CoreV1Api import io.kubernetes.client.openapi.models.V1Deployment -import io.kubernetes.client.openapi.models.V1DeploymentList +import io.kubernetes.client.openapi.models.V1Job +import io.kubernetes.client.openapi.models.V1JobList import io.kubernetes.client.openapi.models.V1LabelSelector import io.kubernetes.client.openapi.models.V1Pod import io.kubernetes.client.openapi.models.V1PodList +import io.kubernetes.client.openapi.models.V1PodTemplate +import io.kubernetes.client.openapi.models.V1PodTemplateSpec import io.kubernetes.client.openapi.models.V1ReplicaSet import io.kubernetes.client.openapi.models.V1ServiceAccount import io.kubernetes.client.openapi.models.V1Status @@ -41,6 +45,7 @@ import java.io.File import java.time.Duration import java.util.concurrent.atomic.AtomicBoolean import kotlin.coroutines.CoroutineContext +import kotlin.random.Random import kotlinx.coroutines.CompletableDeferred import kotlinx.coroutines.Deferred import kotlinx.coroutines.Dispatchers @@ -58,32 +63,87 @@ import kotlinx.coroutines.launch import org.jetbrains.annotations.Blocking import org.jetbrains.annotations.BlockingExecutor -class KubernetesClient( +/** Client interface for the Kubernetes API. */ +interface KubernetesClient { + /** Gets a single [V1Deployment] by [name]. */ + suspend fun getDeployment( + name: String, + namespace: String = Namespaces.NAMESPACE_DEFAULT, + ): V1Deployment? + + suspend fun getPodTemplate( + name: String, + namespace: String = Namespaces.NAMESPACE_DEFAULT, + ): V1PodTemplate? + + /** Gets the [V1ReplicaSet] for the current revision of [deployment]. */ + suspend fun getNewReplicaSet(deployment: V1Deployment): V1ReplicaSet? + + /** Lists Pods for the specified [replicaSet]. */ + suspend fun listPods(replicaSet: V1ReplicaSet): V1PodList + + /** Lists Jobs using the specified [matchLabelsSelector]. */ + suspend fun listJobs( + matchLabelsSelector: String, + namespace: String = Namespaces.NAMESPACE_DEFAULT, + ): V1JobList + + suspend fun createJob(job: V1Job): V1Job + + suspend fun deleteJob( + name: String, + namespace: String = Namespaces.NAMESPACE_DEFAULT, + propagationPolicy: PropagationPolicy = PropagationPolicy.BACKGROUND, + ): V1Status + + /** Suspends until the [V1Deployment] is complete. */ + suspend fun waitUntilDeploymentComplete( + name: String, + namespace: String = Namespaces.NAMESPACE_DEFAULT, + timeout: Duration, + ): V1Deployment + + /** Suspends until the [V1ServiceAccount] exists. */ + suspend fun waitForServiceAccount( + name: String, + namespace: String = Namespaces.NAMESPACE_DEFAULT, + timeout: Duration, + ): V1ServiceAccount + + @Blocking fun kubectlApply(config: File): Sequence + + @Blocking fun kubectlApply(config: String): Sequence + + @Blocking fun kubectlApply(k8sObjects: Iterable): Sequence + + companion object { + /** Generates a random suffix for a Kubernetes object name. */ + fun generateNameSuffix(random: Random) = Names.generateNameSuffix(random) + } +} + +/** Default implementation of [KubernetesClient]. */ +class KubernetesClientImpl( val apiClient: ApiClient = Configuration.getDefaultApiClient(), private val coroutineContext: @BlockingExecutor CoroutineContext = Dispatchers.IO, -) { +) : KubernetesClient { private val coreApi = CoreV1Api(apiClient) private val appsApi = AppsV1Api(apiClient) + private val batchApi = BatchV1Api(apiClient) - /** Gets a single [V1Deployment] by [name]. */ - suspend fun getDeployment( - name: String, - namespace: String = Namespaces.NAMESPACE_DEFAULT, - ): V1Deployment? { - val deployments: List = - apiCall { callback -> - appsApi - .listNamespacedDeployment(namespace) - .fieldSelector("metadata.name=$name") - .executeAsync(callback) - } - .items - check(deployments.size <= 1) - return deployments.singleOrNull() + override suspend fun getDeployment(name: String, namespace: String): V1Deployment? { + return apiCall { callback -> + appsApi.readNamespacedDeployment(name, namespace).executeAsync(callback) + } } - /** Gets the [V1ReplicaSet] for the current revision of [deployment]. */ - suspend fun getNewReplicaSet(deployment: V1Deployment): V1ReplicaSet? { + override suspend fun getPodTemplate(name: String, namespace: String): V1PodTemplate? { + return apiCall { callback -> + coreApi.readNamespacedPodTemplate(name, namespace).executeAsync(callback) + } + } + + override suspend fun getNewReplicaSet(deployment: V1Deployment): V1ReplicaSet? { val namespace: String = deployment.metadata?.namespace ?: Namespaces.NAMESPACE_DEFAULT val labelSelector = deployment.labelSelector val revision = deployment.metadata?.annotations?.get(REVISION_ANNOTATION) ?: return null @@ -98,8 +158,7 @@ class KubernetesClient( .find { it.metadata?.annotations?.get(REVISION_ANNOTATION) == revision } } - /** Lists Pods for the specified [replicaSet]. */ - suspend fun listPods(replicaSet: V1ReplicaSet): V1PodList { + override suspend fun listPods(replicaSet: V1ReplicaSet): V1PodList { val namespace: String = replicaSet.metadata?.namespace ?: Namespaces.NAMESPACE_DEFAULT val labelSelector: V1LabelSelector = checkNotNull(replicaSet.spec).selector @@ -111,6 +170,35 @@ class KubernetesClient( } } + override suspend fun listJobs(matchLabelsSelector: String, namespace: String): V1JobList { + return apiCall { callback -> + batchApi + .listNamespacedJob(namespace) + .labelSelector(matchLabelsSelector) + .executeAsync(callback) + } + } + + override suspend fun createJob(job: V1Job): V1Job { + val namespace: String = job.metadata.namespace ?: Namespaces.NAMESPACE_DEFAULT + return apiCall { callback -> + batchApi.createNamespacedJob(namespace, job).executeAsync(callback) + } + } + + override suspend fun deleteJob( + name: String, + namespace: String, + propagationPolicy: PropagationPolicy, + ): V1Status { + return apiCall { callback -> + batchApi + .deleteNamespacedJob(name, namespace) + .propagationPolicy(propagationPolicy.value) + .executeAsync(callback) + } + } + private inline fun watch( call: okhttp3.Call ): Flow> { @@ -146,10 +234,9 @@ class KubernetesClient( .flowOn(coroutineContext) } - /** Suspends until the [V1Deployment] is complete. */ - suspend fun waitUntilDeploymentComplete( + override suspend fun waitUntilDeploymentComplete( name: String, - namespace: String = Namespaces.NAMESPACE_DEFAULT, + namespace: String, timeout: Duration, ): V1Deployment { return watch( @@ -173,10 +260,9 @@ class KubernetesClient( .first { it.complete } } - /** Suspends until the [V1ServiceAccount] exists. */ - suspend fun waitForServiceAccount( + override suspend fun waitForServiceAccount( name: String, - namespace: String = Namespaces.NAMESPACE_DEFAULT, + namespace: String, timeout: Duration, ): V1ServiceAccount { return watch( @@ -201,29 +287,30 @@ class KubernetesClient( } @Blocking - fun kubectlApply(config: File): Sequence { + override fun kubectlApply(config: File): Sequence { @Suppress("UNCHECKED_CAST") val k8sObjects = Yaml.loadAll(config) as List return kubectlApply(k8sObjects) } @Blocking - fun kubectlApply(config: String): Sequence { + override fun kubectlApply(config: String): Sequence { @Suppress("UNCHECKED_CAST") val k8sObjects = Yaml.loadAll(config) as List return kubectlApply(k8sObjects) } @Blocking - fun kubectlApply(k8sObjects: Iterable): Sequence = sequence { - k8sObjects.map { k8sObject -> - // TODO(kubernetes-client/java#3076): Remove when fixed. - if (k8sObject is V1Pod) { - val podSpec = k8sObject.spec - if (podSpec.overhead != null && podSpec.overhead.isEmpty()) podSpec.overhead(null) - } + override fun kubectlApply(k8sObjects: Iterable): Sequence = + sequence { + k8sObjects.map { k8sObject -> + // TODO(kubernetes-client/java#3076): Remove when fixed. + if (k8sObject is V1Pod) { + val podSpec = k8sObject.spec + if (podSpec.overhead != null && podSpec.overhead.isEmpty()) podSpec.overhead(null) + } - yield(Kubectl.apply(k8sObject.javaClass).apiClient(apiClient).resource(k8sObject).execute()) + yield(Kubectl.apply(k8sObject.javaClass).apiClient(apiClient).resource(k8sObject).execute()) + } } - } companion object { private const val REVISION_ANNOTATION = "deployment.kubernetes.io/revision" @@ -253,11 +340,27 @@ private val V1Deployment.complete: Boolean private val V1Deployment.labelSelector: String get() = checkNotNull(spec?.selector).matchLabelsSelector -private val V1LabelSelector.matchLabelsSelector: String +val V1LabelSelector.matchLabelsSelector: String get() { return matchLabels.map { (key, value) -> "$key=$value" }.joinToString(",") } +fun V1PodTemplateSpec.clone(): V1PodTemplateSpec { + return V1PodTemplateSpec.fromJson(toJson()) +} + +val V1Job.failed: Boolean + get() { + val conditions = status.conditions ?: emptyList() + return conditions.any { it.type == "Failed" && it.status == "True" } + } + +val V1Job.complete: Boolean + get() { + val conditions = status.conditions ?: emptyList() + return conditions.any { it.type == "Complete" && it.status == "True" } + } + private class DeferredApiCallback private constructor(private val delegate: CompletableDeferred) : ApiCallback, Deferred by delegate { @@ -291,6 +394,18 @@ private suspend inline fun apiCall( executeAsync: (callback: ApiCallback) -> okhttp3.Call ): T = apiCallAsync(executeAsync).await() +/** + * Policy for whether and how garbage collection will be performed. + * + * See + * https://kubernetes.io/docs/reference/kubernetes-api/common-parameters/common-parameters/#propagationPolicy + */ +enum class PropagationPolicy(val value: String) { + ORPHAN("Orphan"), + BACKGROUND("Background"), + FOREGROUND("Foreground") +} + /** * Type of Watch event. * diff --git a/src/main/kotlin/org/wfanet/measurement/common/k8s/Names.kt b/src/main/kotlin/org/wfanet/measurement/common/k8s/Names.kt new file mode 100644 index 00000000000..39dddaf95ce --- /dev/null +++ b/src/main/kotlin/org/wfanet/measurement/common/k8s/Names.kt @@ -0,0 +1,49 @@ +/* + * Copyright 2024 The Cross-Media Measurement Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.wfanet.measurement.common.k8s + +import kotlin.random.Random +import kotlin.random.nextULong + +/** Utilities for Kubernetes object names. */ +internal object Names { + /** Set of characters for name generation. */ + private const val NAME_SUFFIX_ALPHABET = "bcdfghjklmnpqrstvwxz2456789" + /** Length of the name suffix. */ + private const val NAME_SUFFIX_LENGTH = 5 + /** Number of bits needed to index into [NAME_SUFFIX_ALPHABET] */ + private const val NAME_SUFFIX_INDEX_BITS = 5 + /** Bit mask for getting index from the rightmost bits of a [ULong]. */ + private val NAME_SUFFIX_INDEX_MASK = ULong.MAX_VALUE.shl(NAME_SUFFIX_INDEX_BITS).inv() + + /** + * Generates a random suffix for a Kubernetes object name. + * + * This is similar to the Kubernetes `generateName` functionality. See + * https://github.com/kubernetes/apimachinery/blob/master/pkg/util/rand/rand.go + */ + fun generateNameSuffix(random: Random): String { + var suffix = "" + var source = random.nextULong() + while (suffix.length < NAME_SUFFIX_LENGTH) { + val index = source.and(NAME_SUFFIX_INDEX_MASK).toInt().mod(NAME_SUFFIX_ALPHABET.length) + suffix += NAME_SUFFIX_ALPHABET[index] + source = source.shr(NAME_SUFFIX_INDEX_BITS) + } + return suffix + } +} diff --git a/src/main/kotlin/org/wfanet/measurement/duchy/deploy/aws/daemon/mill/liquidlegionsv2/BUILD.bazel b/src/main/kotlin/org/wfanet/measurement/duchy/deploy/aws/job/mill/liquidlegionsv2/BUILD.bazel similarity index 65% rename from src/main/kotlin/org/wfanet/measurement/duchy/deploy/aws/daemon/mill/liquidlegionsv2/BUILD.bazel rename to src/main/kotlin/org/wfanet/measurement/duchy/deploy/aws/job/mill/liquidlegionsv2/BUILD.bazel index b91aec0388f..19ecd051110 100644 --- a/src/main/kotlin/org/wfanet/measurement/duchy/deploy/aws/daemon/mill/liquidlegionsv2/BUILD.bazel +++ b/src/main/kotlin/org/wfanet/measurement/duchy/deploy/aws/job/mill/liquidlegionsv2/BUILD.bazel @@ -14,10 +14,10 @@ package( ) kt_jvm_library( - name = "s3_liquid_legions_v2_mill_daemon", - srcs = ["S3LiquidLegionsV2MillDaemon.kt"], + name = "s3_liquid_legions_v2_mill_job", + srcs = ["S3LiquidLegionsV2MillJob.kt"], deps = [ - "//src/main/kotlin/org/wfanet/measurement/duchy/deploy/common/daemon/mill/liquidlegionsv2:liquid_legions_v2_mill_daemon", + "//src/main/kotlin/org/wfanet/measurement/duchy/deploy/common/job/mill/liquidlegionsv2:liquid_legions_v2_mill_job", "@wfa_common_jvm//imports/java/picocli", "@wfa_common_jvm//src/main/kotlin/org/wfanet/measurement/aws/s3", "@wfa_common_jvm//src/main/kotlin/org/wfanet/measurement/common", @@ -25,18 +25,18 @@ kt_jvm_library( ) java_binary( - name = "S3LiquidLegionsV2MillDaemon", - main_class = "org.wfanet.measurement.duchy.deploy.aws.daemon.mill.liquidlegionsv2.S3LiquidLegionsV2MillDaemonKt", + name = "S3LiquidLegionsV2MillJob", + main_class = "org.wfanet.measurement.duchy.deploy.aws.job.mill.liquidlegionsv2.S3LiquidLegionsV2MillJobKt", runtime_deps = [ - ":s3_liquid_legions_v2_mill_daemon", + ":s3_liquid_legions_v2_mill_job", "//imports/java/software/amazon/awssdk/sts", ], ) java_image( - name = "s3_liquid_legions_v2_mill_daemon_image", - binary = ":S3LiquidLegionsV2MillDaemon", - main_class = "org.wfanet.measurement.duchy.deploy.aws.daemon.mill.liquidlegionsv2.S3LiquidLegionsV2MillDaemonKt", + name = "s3_liquid_legions_v2_mill_job_image", + binary = ":S3LiquidLegionsV2MillJob", + main_class = "org.wfanet.measurement.duchy.deploy.aws.job.mill.liquidlegionsv2.S3LiquidLegionsV2MillJobKt", target_compatible_with = DISTROLESS_JAVA, visibility = ["//src:docker_image_deployment"], ) diff --git a/src/main/kotlin/org/wfanet/measurement/duchy/deploy/aws/daemon/mill/liquidlegionsv2/S3LiquidLegionsV2MillDaemon.kt b/src/main/kotlin/org/wfanet/measurement/duchy/deploy/aws/job/mill/liquidlegionsv2/S3LiquidLegionsV2MillJob.kt similarity index 75% rename from src/main/kotlin/org/wfanet/measurement/duchy/deploy/aws/daemon/mill/liquidlegionsv2/S3LiquidLegionsV2MillDaemon.kt rename to src/main/kotlin/org/wfanet/measurement/duchy/deploy/aws/job/mill/liquidlegionsv2/S3LiquidLegionsV2MillJob.kt index 317ba5172bc..7c1500b4417 100644 --- a/src/main/kotlin/org/wfanet/measurement/duchy/deploy/aws/daemon/mill/liquidlegionsv2/S3LiquidLegionsV2MillDaemon.kt +++ b/src/main/kotlin/org/wfanet/measurement/duchy/deploy/aws/job/mill/liquidlegionsv2/S3LiquidLegionsV2MillJob.kt @@ -12,21 +12,21 @@ // See the License for the specific language governing permissions and // limitations under the License. -package org.wfanet.measurement.duchy.deploy.aws.daemon.mill.liquidlegionsv2 +package org.wfanet.measurement.duchy.deploy.aws.job.mill.liquidlegionsv2 import org.wfanet.measurement.aws.s3.S3Flags import org.wfanet.measurement.aws.s3.S3StorageClient import org.wfanet.measurement.common.commandLineMain -import org.wfanet.measurement.duchy.deploy.common.daemon.mill.liquidlegionsv2.LiquidLegionsV2MillDaemon +import org.wfanet.measurement.duchy.deploy.common.job.mill.liquidlegionsv2.LiquidLegionsV2MillJob import picocli.CommandLine @CommandLine.Command( - name = "S3LiquidLegionsV2MillDaemon", - description = ["Liquid Legions V2 Mill daemon."], + name = "S3LiquidLegionsV2MillJob", + description = ["Liquid Legions V2 Mill job."], mixinStandardHelpOptions = true, showDefaultValues = true, ) -class S3LiquidLegionsV2MillDaemon : LiquidLegionsV2MillDaemon() { +class S3LiquidLegionsV2MillJob : LiquidLegionsV2MillJob() { @CommandLine.Mixin private lateinit var s3Flags: S3Flags override fun run() { @@ -35,4 +35,4 @@ class S3LiquidLegionsV2MillDaemon : LiquidLegionsV2MillDaemon() { } } -fun main(args: Array) = commandLineMain(S3LiquidLegionsV2MillDaemon(), args) +fun main(args: Array) = commandLineMain(S3LiquidLegionsV2MillJob(), args) diff --git a/src/main/kotlin/org/wfanet/measurement/duchy/deploy/common/daemon/herald/BUILD.bazel b/src/main/kotlin/org/wfanet/measurement/duchy/deploy/common/daemon/herald/BUILD.bazel index 2e31574aefc..d3839aa860a 100644 --- a/src/main/kotlin/org/wfanet/measurement/duchy/deploy/common/daemon/herald/BUILD.bazel +++ b/src/main/kotlin/org/wfanet/measurement/duchy/deploy/common/daemon/herald/BUILD.bazel @@ -10,8 +10,8 @@ kt_jvm_library( visibility = ["//src/main/kotlin/org/wfanet/measurement/duchy/deploy:__subpackages__"], deps = [ "//src/main/kotlin/org/wfanet/measurement/common/identity", - "//src/main/kotlin/org/wfanet/measurement/duchy/daemon/herald", "//src/main/kotlin/org/wfanet/measurement/duchy/deploy/common:flags", + "//src/main/kotlin/org/wfanet/measurement/duchy/herald", "//src/main/kotlin/org/wfanet/measurement/duchy/storage:tink_key_store", "//src/main/proto/wfa/measurement/internal/duchy:computations_service_kt_jvm_grpc_proto", "//src/main/proto/wfa/measurement/internal/duchy:continuation_tokens_service_kt_jvm_grpc_proto", diff --git a/src/main/kotlin/org/wfanet/measurement/duchy/deploy/common/daemon/herald/HeraldDaemon.kt b/src/main/kotlin/org/wfanet/measurement/duchy/deploy/common/daemon/herald/HeraldDaemon.kt index b1b4e361b23..c77d780edbf 100644 --- a/src/main/kotlin/org/wfanet/measurement/duchy/deploy/common/daemon/herald/HeraldDaemon.kt +++ b/src/main/kotlin/org/wfanet/measurement/duchy/deploy/common/daemon/herald/HeraldDaemon.kt @@ -35,11 +35,11 @@ import org.wfanet.measurement.common.grpc.withShutdownTimeout import org.wfanet.measurement.common.grpc.withVerboseLogging import org.wfanet.measurement.common.identity.withDuchyId import org.wfanet.measurement.common.parseTextProto -import org.wfanet.measurement.duchy.daemon.herald.ContinuationTokenManager -import org.wfanet.measurement.duchy.daemon.herald.Herald import org.wfanet.measurement.duchy.deploy.common.CommonDuchyFlags import org.wfanet.measurement.duchy.deploy.common.ComputationsServiceFlags import org.wfanet.measurement.duchy.deploy.common.SystemApiFlags +import org.wfanet.measurement.duchy.herald.ContinuationTokenManager +import org.wfanet.measurement.duchy.herald.Herald import org.wfanet.measurement.duchy.storage.TinkKeyStore import org.wfanet.measurement.internal.duchy.ComputationsGrpcKt.ComputationsCoroutineStub import org.wfanet.measurement.internal.duchy.ContinuationTokensGrpcKt.ContinuationTokensCoroutineStub diff --git a/src/main/kotlin/org/wfanet/measurement/duchy/deploy/common/daemon/mill/BUILD.bazel b/src/main/kotlin/org/wfanet/measurement/duchy/deploy/common/daemon/mill/BUILD.bazel new file mode 100644 index 00000000000..10a2450f0f6 --- /dev/null +++ b/src/main/kotlin/org/wfanet/measurement/duchy/deploy/common/daemon/mill/BUILD.bazel @@ -0,0 +1,41 @@ +load("@rules_java//java:defs.bzl", "java_binary") +load("@wfa_common_jvm//build:defs.bzl", "test_target") +load("@wfa_rules_kotlin_jvm//kotlin:defs.bzl", "kt_jvm_library") +load("//src/main/docker:macros.bzl", "java_image") + +package( + default_visibility = [test_target(":__pkg__")], +) + +kt_jvm_library( + name = "job_scheduler", + srcs = ["MillJobScheduler.kt"], + deps = [ + "//imports/java/io/kubernetes/client", + "//imports/java/io/kubernetes/client/extended", + "//src/main/kotlin/org/wfanet/measurement/common/k8s:client", + "//src/main/kotlin/org/wfanet/measurement/duchy:computation_stage", + "//src/main/kotlin/org/wfanet/measurement/duchy/deploy/common:flags", + "//src/main/proto/wfa/measurement/internal/duchy:computations_service_kt_jvm_grpc_proto", + "@wfa_common_jvm//imports/java/org/jetbrains/annotations", + "@wfa_common_jvm//imports/kotlin/kotlinx/coroutines:core", + "@wfa_common_jvm//src/main/kotlin/org/wfanet/measurement/common", + "@wfa_common_jvm//src/main/kotlin/org/wfanet/measurement/common/grpc", + ], +) + +java_binary( + name = "MillJobScheduler", + main_class = "org.wfanet.measurement.duchy.deploy.common.daemon.mill.MillJobScheduler", + runtime_deps = [ + ":job_scheduler", + "@wfa_common_jvm//src/main/kotlin/org/wfanet/measurement/gcloud/logging", + ], +) + +java_image( + name = "job_scheduler_image", + binary = ":MillJobScheduler", + main_class = "org.wfanet.measurement.duchy.deploy.common.daemon.mill.MillJobScheduler", + visibility = ["//src:docker_image_deployment"], +) diff --git a/src/main/kotlin/org/wfanet/measurement/duchy/deploy/common/daemon/mill/MillJobScheduler.kt b/src/main/kotlin/org/wfanet/measurement/duchy/deploy/common/daemon/mill/MillJobScheduler.kt new file mode 100644 index 00000000000..be4f459f821 --- /dev/null +++ b/src/main/kotlin/org/wfanet/measurement/duchy/deploy/common/daemon/mill/MillJobScheduler.kt @@ -0,0 +1,389 @@ +/* + * Copyright 2024 The Cross-Media Measurement Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.wfanet.measurement.duchy.deploy.common.daemon.mill + +import io.grpc.StatusException +import io.kubernetes.client.openapi.Configuration +import io.kubernetes.client.openapi.models.V1Container +import io.kubernetes.client.openapi.models.V1Deployment +import io.kubernetes.client.openapi.models.V1Job +import io.kubernetes.client.openapi.models.V1JobSpec +import io.kubernetes.client.openapi.models.V1LabelSelector +import io.kubernetes.client.openapi.models.V1ObjectMeta +import io.kubernetes.client.openapi.models.V1OwnerReference +import io.kubernetes.client.openapi.models.V1PodTemplate +import io.kubernetes.client.openapi.models.V1PodTemplateSpec +import io.kubernetes.client.util.ClientBuilder +import java.time.Duration +import java.util.logging.Logger +import kotlin.properties.Delegates +import kotlin.random.Random +import kotlinx.coroutines.currentCoroutineContext +import kotlinx.coroutines.isActive +import kotlinx.coroutines.runBlocking +import kotlinx.coroutines.time.delay +import org.wfanet.measurement.common.commandLineMain +import org.wfanet.measurement.common.grpc.TlsFlags +import org.wfanet.measurement.common.grpc.buildMutualTlsChannel +import org.wfanet.measurement.common.grpc.withDefaultDeadline +import org.wfanet.measurement.common.grpc.withShutdownTimeout +import org.wfanet.measurement.common.k8s.KubernetesClient +import org.wfanet.measurement.common.k8s.KubernetesClientImpl +import org.wfanet.measurement.common.k8s.clone +import org.wfanet.measurement.common.k8s.complete +import org.wfanet.measurement.common.k8s.failed +import org.wfanet.measurement.common.k8s.matchLabelsSelector +import org.wfanet.measurement.common.toProtoDuration +import org.wfanet.measurement.duchy.deploy.common.CommonDuchyFlags +import org.wfanet.measurement.duchy.deploy.common.ComputationsServiceFlags +import org.wfanet.measurement.duchy.toProtocolStage +import org.wfanet.measurement.internal.duchy.ClaimWorkRequest +import org.wfanet.measurement.internal.duchy.ClaimWorkResponse +import org.wfanet.measurement.internal.duchy.ComputationStage +import org.wfanet.measurement.internal.duchy.ComputationToken +import org.wfanet.measurement.internal.duchy.ComputationTypeEnum.ComputationType +import org.wfanet.measurement.internal.duchy.ComputationsGrpcKt +import org.wfanet.measurement.internal.duchy.claimWorkRequest +import org.wfanet.measurement.internal.duchy.protocol.LiquidLegionsSketchAggregationV2 +import org.wfanet.measurement.internal.duchy.protocol.ReachOnlyLiquidLegionsSketchAggregationV2 +import picocli.CommandLine + +/** + * Scheduler for Mill Kubernetes Jobs. + * + * See [MillJobScheduler.Flags] for additional parameter documentation. + */ +class MillJobScheduler( + private val duchyId: String, + private val computationsStub: ComputationsGrpcKt.ComputationsCoroutineStub, + private val deploymentName: String, + private val pollingDelay: Duration, + private val successfulJobHistoryLimit: Int, + private val failedJobHistoryLimit: Int, + private val liquidLegionsV2PodTemplateName: String, + private val liquidLegionsV2MaximumConcurrency: Int, + private val liquidLegionsV2WorkLockDuration: Duration, + private val random: Random = Random.Default, + private val k8sClient: KubernetesClient = KubernetesClientImpl(), +) { + private lateinit var deployment: V1Deployment + private lateinit var liquidLegionsV2PodTemplate: V1PodTemplate + + suspend fun run() { + deployment = + k8sClient.getDeployment(deploymentName) ?: error("Deployment $deploymentName not found") + liquidLegionsV2PodTemplate = + k8sClient.getPodTemplate(liquidLegionsV2PodTemplateName) + ?: error("PodTemplate $liquidLegionsV2PodTemplateName not found") + + while (currentCoroutineContext().isActive) { + for (computationTypeConfig in COMPUTATION_TYPE_CONFIGS) { + val ownedJobs: List = getOwnedJobs(computationTypeConfig.millType) + process(computationTypeConfig, ownedJobs) + ownedJobs.cleanUp() + } + delay(pollingDelay) + } + } + + private suspend fun process( + computationTypeConfig: ComputationTypeConfig, + ownedJobs: List, + ) { + val millType: MillType = computationTypeConfig.millType + val maximumConcurrency = millType.maximumConcurrency + val computationType: ComputationType = computationTypeConfig.computationType + + val activeJobCount: Int = ownedJobs.count { (it.status.active ?: 0) > 0 } + if (activeJobCount >= maximumConcurrency) { + logger.fine { + "Not scheduling: Maximum concurrency limit $maximumConcurrency reached for $millType" + } + return + } + + logger.fine { + "Active job count $activeJobCount is below maximum concurrency $maximumConcurrency for " + + "mill type $millType. Checking for work..." + } + val jobName: String = generateJobName(millType) + val claimedToken: ComputationToken? = + claimWork( + claimWorkRequest { + this.computationType = computationType + owner = jobName + lockDuration = millType.workLockDuration.toProtoDuration() + prioritizedStages += computationTypeConfig.prioritizedStages + } + ) + if (claimedToken == null) { + logger.fine { "Not scheduling: No work available for computation type $computationType" } + return + } + + val claimedComputationId: String = claimedToken.globalComputationId + logger.info { "Claimed work item for Computation $claimedComputationId" } + val template = + millType.podTemplate.template.clone().apply { + val container: V1Container = spec.containers.first() + container.addArgsItem("--mill-id=$jobName") + container.addArgsItem("--computation-type=$computationType") + container.addArgsItem("--claimed-computation-id=$claimedComputationId") + } + createJob(jobName, millType, template) + logger.info { "Scheduled Job $jobName for Computation $claimedComputationId" } + if (activeJobCount + 1 >= maximumConcurrency) { + logger.info { "Mill type $millType is now at maximum concurrency limit $maximumConcurrency" } + } + } + + /** Cleans up completed Jobs by enforcing history limits. */ + private suspend fun Collection.cleanUp() { + deleteOverLimit(successfulJobHistoryLimit) { it.complete && !it.failed } + deleteOverLimit(failedJobHistoryLimit) { it.failed } + } + + /** Deletes completed Jobs matching [predicate] that are over the history [limit]. */ + private suspend fun Collection.deleteOverLimit(limit: Int, predicate: (V1Job) -> Boolean) { + val affectedJobs: List = filter(predicate) + val overLimitCount: Int = affectedJobs.size - limit + if (overLimitCount <= 0) { + return + } + + affectedJobs + .sortedBy { it.status.completionTime } + .take(overLimitCount) + .forEach { + k8sClient.deleteJob(it.metadata.name, it.metadata.namespace) + logger.info { "Deleted Job ${it.metadata.name}" } + } + } + + private fun generateJobName(millType: MillType): String { + val suffix = KubernetesClient.generateNameSuffix(random) + return "$duchyId-${millType.jobNamePrefix}-$suffix" + } + + private suspend fun claimWork(request: ClaimWorkRequest): ComputationToken? { + val response: ClaimWorkResponse = + try { + computationsStub.claimWork(request) + } catch (e: StatusException) { + throw Exception("Error claiming work", e) + } + return if (response.hasToken()) response.token else null + } + + private suspend fun getOwnedJobs(millType: MillType): List { + val labelSelector = V1LabelSelector().apply { matchLabels[MILL_TYPE_LABEL] = millType.name } + return k8sClient.listJobs(labelSelector.matchLabelsSelector).items.filter { + it.metadata.ownerReferences.any { ownerRef -> ownerRef.uid == deployment.metadata.uid } + } + } + + private suspend fun createJob( + name: String, + millType: MillType, + template: V1PodTemplateSpec, + ): V1Job { + val job = + V1Job().apply { + apiVersion = "batch/v1" + kind = "Job" + metadata = + V1ObjectMeta().apply { + this.name = name + putLabelsItem(MILL_TYPE_LABEL, millType.name) + ownerReferences = + listOf( + V1OwnerReference().apply { + apiVersion = deployment.apiVersion + kind = deployment.kind + this.name = deployment.metadata.name + uid = deployment.metadata.uid + } + ) + } + spec = V1JobSpec().apply { this.template = template } + } + + return k8sClient.createJob(job) + } + + private enum class MillType(val jobNamePrefix: String) { + LIQUID_LEGIONS_V2("llv2-mill-job"), + } + + private data class ComputationTypeConfig( + val computationType: ComputationType, + val millType: MillType, + val prioritizedStages: List, + ) + + private val MillType.workLockDuration: Duration + get() = + when (this) { + MillType.LIQUID_LEGIONS_V2 -> liquidLegionsV2WorkLockDuration + } + + private val MillType.maximumConcurrency: Int + get() = + when (this) { + MillType.LIQUID_LEGIONS_V2 -> liquidLegionsV2MaximumConcurrency + } + + private val MillType.podTemplate: V1PodTemplate + get() = + when (this) { + MillType.LIQUID_LEGIONS_V2 -> liquidLegionsV2PodTemplate + } + + private class Flags { + @CommandLine.Mixin + lateinit var tls: TlsFlags + private set + + @CommandLine.Mixin + lateinit var duchy: CommonDuchyFlags + private set + + @CommandLine.Mixin + lateinit var internalApi: ComputationsServiceFlags + private set + + @CommandLine.Option( + names = ["--channel-shutdown-timeout"], + defaultValue = "3s", + description = ["How long to allow for the gRPC channel to shutdown."], + ) + lateinit var channelShutdownTimeout: Duration + private set + + @CommandLine.Option( + names = ["--deployment-name"], + description = ["Name of this K8s Deployment"], + required = true, + ) + lateinit var deploymentName: String + private set + + @CommandLine.Option( + names = ["--polling-delay"], + description = ["Delay before subsequent polling attempts"], + defaultValue = "2s", + ) + lateinit var pollingDelay: Duration + private set + + @set:CommandLine.Option( + names = ["--successful-jobs-history-limit"], + description = ["Number of successful jobs to retain per mill type"], + defaultValue = "3", + ) + var successfulJobsHistoryLimit by Delegates.notNull() + private set + + @set:CommandLine.Option( + names = ["--failed-jobs-history-limit"], + description = ["Number of failed jobs to retain per mill type"], + defaultValue = "1", + ) + var failedJobsHistoryLimit by Delegates.notNull() + private set + + @CommandLine.Option( + names = ["--llv2-pod-template-name"], + description = ["Name of the K8s PodTemplate for Liquid Legions v2 Jobs"], + required = true, + ) + lateinit var liquidLegionsV2PodTemplateName: String + private set + + @set:CommandLine.Option( + names = ["--llv2-maximum-concurrency"], + description = ["Maximum number of concurrent Liquid Legions v2 Jobs"], + defaultValue = "1", + ) + var liquidLegionsV2MaximumConcurrency by Delegates.notNull() + private set + + @CommandLine.Option( + names = ["--llv2-work-lock-duration"], + defaultValue = "5m", + description = ["How long to hold work locks for Liquid Legions v2"], + ) + lateinit var liquidLegionsV2WorkLockDuration: Duration + private set + } + + companion object { + private const val MILL_TYPE_LABEL = "mill-type" + private val COMPUTATION_TYPE_CONFIGS = + listOf( + ComputationTypeConfig( + ComputationType.LIQUID_LEGIONS_SKETCH_AGGREGATION_V2, + MillType.LIQUID_LEGIONS_V2, + listOf(LiquidLegionsSketchAggregationV2.Stage.INITIALIZATION_PHASE.toProtocolStage()), + ), + ComputationTypeConfig( + ComputationType.REACH_ONLY_LIQUID_LEGIONS_SKETCH_AGGREGATION_V2, + MillType.LIQUID_LEGIONS_V2, + listOf( + ReachOnlyLiquidLegionsSketchAggregationV2.Stage.INITIALIZATION_PHASE.toProtocolStage() + ), + ), + ) + + private val logger: Logger = Logger.getLogger(this::class.java.name) + + @CommandLine.Command( + name = "MillJobScheduler", + description = ["Mill Job scheduler"], + mixinStandardHelpOptions = true, + showDefaultValues = true, + ) + private fun run(@CommandLine.Mixin flags: Flags) { + Configuration.setDefaultApiClient(ClientBuilder.cluster().build()) + val internalApiChannel = + buildMutualTlsChannel( + flags.internalApi.target, + flags.tls.signingCerts, + flags.internalApi.certHost, + ) + .withShutdownTimeout(flags.channelShutdownTimeout) + .withDefaultDeadline(flags.internalApi.defaultDeadlineDuration) + val computationsStub = + ComputationsGrpcKt.ComputationsCoroutineStub(internalApiChannel).withWaitForReady() + val app = + MillJobScheduler( + flags.duchy.duchyName, + computationsStub, + flags.deploymentName, + flags.pollingDelay, + flags.successfulJobsHistoryLimit, + flags.failedJobsHistoryLimit, + flags.liquidLegionsV2PodTemplateName, + flags.liquidLegionsV2MaximumConcurrency, + flags.liquidLegionsV2WorkLockDuration, + ) + runBlocking { app.run() } + } + + @JvmStatic fun main(args: Array) = commandLineMain(Companion::run, args) + } +} diff --git a/src/main/kotlin/org/wfanet/measurement/duchy/deploy/common/daemon/mill/shareshuffle/BUILD.bazel b/src/main/kotlin/org/wfanet/measurement/duchy/deploy/common/daemon/mill/shareshuffle/BUILD.bazel index e1ed7176989..ac38eb94dd4 100644 --- a/src/main/kotlin/org/wfanet/measurement/duchy/deploy/common/daemon/mill/shareshuffle/BUILD.bazel +++ b/src/main/kotlin/org/wfanet/measurement/duchy/deploy/common/daemon/mill/shareshuffle/BUILD.bazel @@ -14,9 +14,10 @@ kt_jvm_library( visibility = ["//src/main/kotlin/org/wfanet/measurement/duchy/deploy:__subpackages__"], deps = [ "//src/main/kotlin/org/wfanet/measurement/common/identity", - "//src/main/kotlin/org/wfanet/measurement/duchy/daemon/mill/shareshuffle:honest_majority_share_shuffle_mill", "//src/main/kotlin/org/wfanet/measurement/duchy/db/computation", "//src/main/kotlin/org/wfanet/measurement/duchy/deploy/common:flags", + "//src/main/kotlin/org/wfanet/measurement/duchy/mill:mill_flags", + "//src/main/kotlin/org/wfanet/measurement/duchy/mill/shareshuffle:honest_majority_share_shuffle_mill", "//src/main/kotlin/org/wfanet/measurement/duchy/storage:tink_key_store", "//src/main/proto/wfa/measurement/internal/duchy:computations_service_kt_jvm_grpc_proto", "//src/main/proto/wfa/measurement/internal/duchy/protocol:honest_majority_share_shuffle_kt_jvm_proto", diff --git a/src/main/kotlin/org/wfanet/measurement/duchy/deploy/common/daemon/mill/shareshuffle/HonestMajorityShareShuffleMillDaemon.kt b/src/main/kotlin/org/wfanet/measurement/duchy/deploy/common/daemon/mill/shareshuffle/HonestMajorityShareShuffleMillDaemon.kt index beb5b93d84d..acce1680000 100644 --- a/src/main/kotlin/org/wfanet/measurement/duchy/deploy/common/daemon/mill/shareshuffle/HonestMajorityShareShuffleMillDaemon.kt +++ b/src/main/kotlin/org/wfanet/measurement/duchy/deploy/common/daemon/mill/shareshuffle/HonestMajorityShareShuffleMillDaemon.kt @@ -42,10 +42,10 @@ import org.wfanet.measurement.common.identity.withDuchyId import org.wfanet.measurement.common.logAndSuppressExceptionSuspend import org.wfanet.measurement.common.parseTextProto import org.wfanet.measurement.common.throttler.MinimumIntervalThrottler -import org.wfanet.measurement.duchy.daemon.mill.Certificate -import org.wfanet.measurement.duchy.daemon.mill.shareshuffle.HonestMajorityShareShuffleMill -import org.wfanet.measurement.duchy.daemon.mill.shareshuffle.crypto.JniHonestMajorityShareShuffleCryptor import org.wfanet.measurement.duchy.db.computation.ComputationDataClients +import org.wfanet.measurement.duchy.mill.Certificate +import org.wfanet.measurement.duchy.mill.shareshuffle.HonestMajorityShareShuffleMill +import org.wfanet.measurement.duchy.mill.shareshuffle.crypto.JniHonestMajorityShareShuffleCryptor import org.wfanet.measurement.duchy.storage.TinkKeyStore import org.wfanet.measurement.internal.duchy.ComputationStatsGrpcKt.ComputationStatsCoroutineStub import org.wfanet.measurement.internal.duchy.ComputationsGrpcKt.ComputationsCoroutineStub @@ -136,8 +136,7 @@ abstract class HonestMajorityShareShuffleMillDaemon : Runnable { // This will be the name of the pod when deployed to Kubernetes. Note that the millId is // included in mill logs to help debugging. - // TODO(@renjiez): add the parameter to override the millId. - val millId = System.getenv("HOSTNAME") + val millId = flags.millId val privateKeyStore = flags.keyEncryptionKeyTinkFile?.let { file -> @@ -186,9 +185,7 @@ abstract class HonestMajorityShareShuffleMillDaemon : Runnable { runBlocking { withContext(CoroutineName("Mill $millId")) { val throttler = MinimumIntervalThrottler(Clock.systemUTC(), flags.pollingInterval) - throttler.loopOnReady { - logAndSuppressExceptionSuspend { mill.pollAndProcessNextComputation() } - } + throttler.loopOnReady { logAndSuppressExceptionSuspend { mill.claimAndProcessWork() } } } } } diff --git a/src/main/kotlin/org/wfanet/measurement/duchy/deploy/common/daemon/mill/shareshuffle/HonestMajorityShareShuffleMillFlags.kt b/src/main/kotlin/org/wfanet/measurement/duchy/deploy/common/daemon/mill/shareshuffle/HonestMajorityShareShuffleMillFlags.kt index 13f1d016c92..ccc4dcad00b 100644 --- a/src/main/kotlin/org/wfanet/measurement/duchy/deploy/common/daemon/mill/shareshuffle/HonestMajorityShareShuffleMillFlags.kt +++ b/src/main/kotlin/org/wfanet/measurement/duchy/deploy/common/daemon/mill/shareshuffle/HonestMajorityShareShuffleMillFlags.kt @@ -16,43 +16,28 @@ package org.wfanet.measurement.duchy.deploy.common.daemon.mill.shareshuffle import java.io.File import java.time.Duration -import kotlin.properties.Delegates -import org.wfanet.measurement.common.grpc.TlsFlags -import org.wfanet.measurement.common.identity.DuchyInfoFlags import org.wfanet.measurement.duchy.deploy.common.CommonDuchyFlags import org.wfanet.measurement.duchy.deploy.common.ComputationsServiceFlags import org.wfanet.measurement.duchy.deploy.common.KingdomPublicApiFlags import org.wfanet.measurement.duchy.deploy.common.SystemApiFlags +import org.wfanet.measurement.duchy.mill.MillFlags import picocli.CommandLine -class HonestMajorityShareShuffleMillFlags { - // TODO(@renjie): Extract common mill flags. +class HonestMajorityShareShuffleMillFlags : MillFlags() { @CommandLine.Mixin lateinit var duchy: CommonDuchyFlags private set @CommandLine.Mixin - lateinit var tlsFlags: TlsFlags + lateinit var systemApiFlags: SystemApiFlags private set @CommandLine.Mixin - lateinit var duchyInfoFlags: DuchyInfoFlags - private set - - @CommandLine.Option( - names = ["--duchy-computation-control-target"], - description = ["Key-value pair of Duchy ID to ComputationControl service target."], - required = true, - ) - lateinit var computationControlServiceTargets: Map + lateinit var computationsServiceFlags: ComputationsServiceFlags private set - @CommandLine.Option( - names = ["--channel-shutdown-timeout"], - defaultValue = "3s", - description = ["How long to allow for the gRPC channel to shutdown."], - ) - lateinit var channelShutdownTimeout: Duration + @CommandLine.Mixin + lateinit var publicApiFlags: KingdomPublicApiFlags private set @CommandLine.Option( @@ -63,58 +48,6 @@ class HonestMajorityShareShuffleMillFlags { lateinit var pollingInterval: Duration private set - @CommandLine.Option( - names = ["--work-lock-duration"], - defaultValue = "5m", - description = ["How long to hold work locks."], - ) - lateinit var workLockDuration: Duration - private set - - @CommandLine.Mixin - lateinit var systemApiFlags: SystemApiFlags - private set - - @CommandLine.Mixin - lateinit var publicApiFlags: KingdomPublicApiFlags - private set - - @CommandLine.Mixin - lateinit var computationsServiceFlags: ComputationsServiceFlags - private set - - @set:CommandLine.Option( - names = ["--bytes-per-chunk"], - description = ["The number of bytes in a chunk when sending rpc result to other duchy."], - defaultValue = "32768", // 32 KiB. See https://github.com/grpc/grpc.github.io/issues/371. - ) - var requestChunkSizeBytes by Delegates.notNull() - private set - - @CommandLine.Option( - names = ["--consent-signaling-certificate-resource-name"], - description = ["The resource name of the duchy's consent signaling certificate."], - required = true, - ) - lateinit var csCertificateName: String - private set - - @CommandLine.Option( - names = ["--consent-signaling-private-key-der-file"], - description = ["The duchy's consent signaling private key (DER format) file."], - required = true, - ) - lateinit var csPrivateKeyDerFile: File - private set - - @CommandLine.Option( - names = ["--consent-signaling-certificate-der-file"], - description = ["The duchy's consent signaling certificate (DER format) file."], - required = true, - ) - lateinit var csCertificateDerFile: File - private set - @CommandLine.Option( names = ["--protocols-setup-config"], description = ["ProtocolsSetupConfig proto message in text format."], diff --git a/src/main/kotlin/org/wfanet/measurement/duchy/deploy/common/daemon/mill/liquidlegionsv2/BUILD.bazel b/src/main/kotlin/org/wfanet/measurement/duchy/deploy/common/job/mill/liquidlegionsv2/BUILD.bazel similarity index 65% rename from src/main/kotlin/org/wfanet/measurement/duchy/deploy/common/daemon/mill/liquidlegionsv2/BUILD.bazel rename to src/main/kotlin/org/wfanet/measurement/duchy/deploy/common/job/mill/liquidlegionsv2/BUILD.bazel index ddc92a24593..20062c38804 100644 --- a/src/main/kotlin/org/wfanet/measurement/duchy/deploy/common/daemon/mill/liquidlegionsv2/BUILD.bazel +++ b/src/main/kotlin/org/wfanet/measurement/duchy/deploy/common/job/mill/liquidlegionsv2/BUILD.bazel @@ -6,18 +6,19 @@ load("//src/main/docker:macros.bzl", "java_image") package(default_testonly = True) #TODO: delete when InMemoryKeyStore and FakeHybridCipher are not used. kt_jvm_library( - name = "liquid_legions_v2_mill_daemon", + name = "liquid_legions_v2_mill_job", srcs = [ - "LiquidLegionsV2MillDaemon.kt", "LiquidLegionsV2MillFlags.kt", + "LiquidLegionsV2MillJob.kt", ], visibility = ["//src/main/kotlin/org/wfanet/measurement/duchy/deploy:__subpackages__"], deps = [ "//src/main/kotlin/org/wfanet/measurement/common/identity", - "//src/main/kotlin/org/wfanet/measurement/duchy/daemon/mill/liquidlegionsv2:reach_frequency_liquid_legions_v2_mill", - "//src/main/kotlin/org/wfanet/measurement/duchy/daemon/mill/liquidlegionsv2:reach_only_liquid_legions_v2_mill", "//src/main/kotlin/org/wfanet/measurement/duchy/db/computation", "//src/main/kotlin/org/wfanet/measurement/duchy/deploy/common:flags", + "//src/main/kotlin/org/wfanet/measurement/duchy/mill:mill_flags", + "//src/main/kotlin/org/wfanet/measurement/duchy/mill/liquidlegionsv2:reach_frequency_liquid_legions_v2_mill", + "//src/main/kotlin/org/wfanet/measurement/duchy/mill/liquidlegionsv2:reach_only_liquid_legions_v2_mill", "//src/main/proto/wfa/measurement/internal/duchy:computations_service_kt_jvm_grpc_proto", "//src/main/proto/wfa/measurement/internal/duchy/protocol:liquid_legions_v2_kt_jvm_proto", "//src/main/proto/wfa/measurement/system/v1alpha:computation_control_service_kt_jvm_grpc_proto", @@ -34,10 +35,10 @@ kt_jvm_library( ) kt_jvm_library( - name = "forwarded_storage_liquid_legions_v2_mill_daemon", - srcs = ["ForwardedStorageLiquidLegionsV2MillDaemon.kt"], + name = "forwarded_storage_liquid_legions_v2_mill_job", + srcs = ["ForwardedStorageLiquidLegionsV2MillJob.kt"], deps = [ - ":liquid_legions_v2_mill_daemon", + ":liquid_legions_v2_mill_job", "@wfa_common_jvm//imports/java/picocli", "@wfa_common_jvm//src/main/kotlin/org/wfanet/measurement/common", "@wfa_common_jvm//src/main/kotlin/org/wfanet/measurement/storage/forwarded", @@ -45,15 +46,15 @@ kt_jvm_library( ) java_binary( - name = "ForwardedStorageLiquidLegionsV2MillDaemon", - main_class = "org.wfanet.measurement.duchy.deploy.common.daemon.mill.liquidlegionsv2.ForwardedStorageLiquidLegionsV2MillDaemonKt", - runtime_deps = [":forwarded_storage_liquid_legions_v2_mill_daemon"], + name = "ForwardedStorageLiquidLegionsV2MillJob", + main_class = "org.wfanet.measurement.duchy.deploy.common.job.mill.liquidlegionsv2.ForwardedStorageLiquidLegionsV2MillJobKt", + runtime_deps = [":forwarded_storage_liquid_legions_v2_mill_job"], ) java_image( - name = "forwarded_storage_liquid_legions_v2_mill_daemon_image", - binary = ":ForwardedStorageLiquidLegionsV2MillDaemon", - main_class = "org.wfanet.measurement.duchy.deploy.common.daemon.mill.liquidlegionsv2.ForwardedStorageLiquidLegionsV2MillDaemonKt", + name = "forwarded_storage_liquid_legions_v2_mill_job_image", + binary = ":ForwardedStorageLiquidLegionsV2MillJob", + main_class = "org.wfanet.measurement.duchy.deploy.common.job.mill.liquidlegionsv2.ForwardedStorageLiquidLegionsV2MillJobKt", target_compatible_with = DISTROLESS_JAVA, visibility = ["//src:docker_image_deployment"], ) diff --git a/src/main/kotlin/org/wfanet/measurement/duchy/deploy/common/daemon/mill/liquidlegionsv2/ForwardedStorageLiquidLegionsV2MillDaemon.kt b/src/main/kotlin/org/wfanet/measurement/duchy/deploy/common/job/mill/liquidlegionsv2/ForwardedStorageLiquidLegionsV2MillJob.kt similarity index 79% rename from src/main/kotlin/org/wfanet/measurement/duchy/deploy/common/daemon/mill/liquidlegionsv2/ForwardedStorageLiquidLegionsV2MillDaemon.kt rename to src/main/kotlin/org/wfanet/measurement/duchy/deploy/common/job/mill/liquidlegionsv2/ForwardedStorageLiquidLegionsV2MillJob.kt index 7db6d57a8ac..2d277a6bbe4 100644 --- a/src/main/kotlin/org/wfanet/measurement/duchy/deploy/common/daemon/mill/liquidlegionsv2/ForwardedStorageLiquidLegionsV2MillDaemon.kt +++ b/src/main/kotlin/org/wfanet/measurement/duchy/deploy/common/job/mill/liquidlegionsv2/ForwardedStorageLiquidLegionsV2MillJob.kt @@ -12,19 +12,19 @@ // See the License for the specific language governing permissions and // limitations under the License. -package org.wfanet.measurement.duchy.deploy.common.daemon.mill.liquidlegionsv2 +package org.wfanet.measurement.duchy.deploy.common.job.mill.liquidlegionsv2 import org.wfanet.measurement.common.commandLineMain import org.wfanet.measurement.storage.forwarded.ForwardedStorageFromFlags import picocli.CommandLine @CommandLine.Command( - name = "ForwardedStorageLiquidLegionsV2MillDaemon", - description = ["LiquidLegionsV2 Mill daemon."], + name = "ForwardedStorageLiquidLegionsV2MillJob", + description = ["LiquidLegionsV2 Mill job."], mixinStandardHelpOptions = true, showDefaultValues = true, ) -class ForwardedStorageLiquidLegionsV2MillDaemon : LiquidLegionsV2MillDaemon() { +class ForwardedStorageLiquidLegionsV2MillJob : LiquidLegionsV2MillJob() { @CommandLine.Mixin private lateinit var forwardedStorageFlags: ForwardedStorageFromFlags.Flags override fun run() { @@ -32,4 +32,4 @@ class ForwardedStorageLiquidLegionsV2MillDaemon : LiquidLegionsV2MillDaemon() { } } -fun main(args: Array) = commandLineMain(ForwardedStorageLiquidLegionsV2MillDaemon(), args) +fun main(args: Array) = commandLineMain(ForwardedStorageLiquidLegionsV2MillJob(), args) diff --git a/src/main/kotlin/org/wfanet/measurement/duchy/deploy/common/job/mill/liquidlegionsv2/LiquidLegionsV2MillFlags.kt b/src/main/kotlin/org/wfanet/measurement/duchy/deploy/common/job/mill/liquidlegionsv2/LiquidLegionsV2MillFlags.kt new file mode 100644 index 00000000000..766772e38e3 --- /dev/null +++ b/src/main/kotlin/org/wfanet/measurement/duchy/deploy/common/job/mill/liquidlegionsv2/LiquidLegionsV2MillFlags.kt @@ -0,0 +1,61 @@ +// Copyright 2020 The Cross-Media Measurement Authors +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package org.wfanet.measurement.duchy.deploy.common.job.mill.liquidlegionsv2 + +import kotlin.properties.Delegates +import org.wfanet.measurement.duchy.deploy.common.CommonDuchyFlags +import org.wfanet.measurement.duchy.deploy.common.ComputationsServiceFlags +import org.wfanet.measurement.duchy.deploy.common.SystemApiFlags +import org.wfanet.measurement.duchy.mill.MillFlags +import org.wfanet.measurement.internal.duchy.ComputationTypeEnum.ComputationType +import picocli.CommandLine + +class LiquidLegionsV2MillFlags : MillFlags() { + @CommandLine.Mixin + lateinit var duchy: CommonDuchyFlags + private set + + @CommandLine.Mixin + lateinit var systemApiFlags: SystemApiFlags + private set + + @CommandLine.Mixin + lateinit var computationsServiceFlags: ComputationsServiceFlags + private set + + @set:CommandLine.Option( + names = ["--parallelism"], + description = ["Maximum number of threads used in crypto actions"], + defaultValue = "1", + ) + var parallelism by Delegates.notNull() + private set + + @CommandLine.Option( + names = ["--claimed-computation-id"], + description = ["Global Computation ID of the claimed work item"], + required = true, + ) + lateinit var claimedGlobalComputationId: String + private set + + @CommandLine.Option( + names = ["--computation-type"], + description = ["Type (protocol) of the computation"], + required = true, + ) + lateinit var computationType: ComputationType + private set +} diff --git a/src/main/kotlin/org/wfanet/measurement/duchy/deploy/common/daemon/mill/liquidlegionsv2/LiquidLegionsV2MillDaemon.kt b/src/main/kotlin/org/wfanet/measurement/duchy/deploy/common/job/mill/liquidlegionsv2/LiquidLegionsV2MillJob.kt similarity index 59% rename from src/main/kotlin/org/wfanet/measurement/duchy/deploy/common/daemon/mill/liquidlegionsv2/LiquidLegionsV2MillDaemon.kt rename to src/main/kotlin/org/wfanet/measurement/duchy/deploy/common/job/mill/liquidlegionsv2/LiquidLegionsV2MillJob.kt index 7e4136f8ea1..c91745e3f04 100644 --- a/src/main/kotlin/org/wfanet/measurement/duchy/deploy/common/daemon/mill/liquidlegionsv2/LiquidLegionsV2MillDaemon.kt +++ b/src/main/kotlin/org/wfanet/measurement/duchy/deploy/common/job/mill/liquidlegionsv2/LiquidLegionsV2MillJob.kt @@ -12,16 +12,15 @@ // See the License for the specific language governing permissions and // limitations under the License. -package org.wfanet.measurement.duchy.deploy.common.daemon.mill.liquidlegionsv2 +package org.wfanet.measurement.duchy.deploy.common.job.mill.liquidlegionsv2 import com.google.protobuf.ByteString import io.grpc.Channel import io.opentelemetry.api.GlobalOpenTelemetry import io.opentelemetry.api.OpenTelemetry -import java.time.Clock import kotlinx.coroutines.CoroutineName +import kotlinx.coroutines.isActive import kotlinx.coroutines.runBlocking -import kotlinx.coroutines.withContext import org.wfanet.measurement.common.crypto.SigningCerts import org.wfanet.measurement.common.crypto.SigningKeyHandle import org.wfanet.measurement.common.crypto.readCertificate @@ -31,15 +30,15 @@ import org.wfanet.measurement.common.grpc.withDefaultDeadline import org.wfanet.measurement.common.grpc.withShutdownTimeout import org.wfanet.measurement.common.identity.DuchyInfo import org.wfanet.measurement.common.identity.withDuchyId -import org.wfanet.measurement.common.logAndSuppressExceptionSuspend -import org.wfanet.measurement.common.throttler.MinimumIntervalThrottler -import org.wfanet.measurement.duchy.daemon.mill.Certificate -import org.wfanet.measurement.duchy.daemon.mill.liquidlegionsv2.ReachFrequencyLiquidLegionsV2Mill -import org.wfanet.measurement.duchy.daemon.mill.liquidlegionsv2.ReachOnlyLiquidLegionsV2Mill -import org.wfanet.measurement.duchy.daemon.mill.liquidlegionsv2.crypto.JniLiquidLegionsV2Encryption -import org.wfanet.measurement.duchy.daemon.mill.liquidlegionsv2.crypto.JniReachOnlyLiquidLegionsV2Encryption import org.wfanet.measurement.duchy.db.computation.ComputationDataClients +import org.wfanet.measurement.duchy.mill.Certificate +import org.wfanet.measurement.duchy.mill.liquidlegionsv2.LiquidLegionsV2Mill +import org.wfanet.measurement.duchy.mill.liquidlegionsv2.ReachFrequencyLiquidLegionsV2Mill +import org.wfanet.measurement.duchy.mill.liquidlegionsv2.ReachOnlyLiquidLegionsV2Mill +import org.wfanet.measurement.duchy.mill.liquidlegionsv2.crypto.JniLiquidLegionsV2Encryption +import org.wfanet.measurement.duchy.mill.liquidlegionsv2.crypto.JniReachOnlyLiquidLegionsV2Encryption import org.wfanet.measurement.internal.duchy.ComputationStatsGrpcKt.ComputationStatsCoroutineStub +import org.wfanet.measurement.internal.duchy.ComputationTypeEnum import org.wfanet.measurement.internal.duchy.ComputationsGrpcKt.ComputationsCoroutineStub import org.wfanet.measurement.storage.StorageClient import org.wfanet.measurement.system.v1alpha.ComputationControlGrpcKt.ComputationControlCoroutineStub @@ -48,7 +47,8 @@ import org.wfanet.measurement.system.v1alpha.ComputationParticipantsGrpcKt.Compu import org.wfanet.measurement.system.v1alpha.ComputationsGrpcKt.ComputationsCoroutineStub as SystemComputationsCoroutineStub import picocli.CommandLine -abstract class LiquidLegionsV2MillDaemon : Runnable { +/** Job which processes claimed work using a Liquid Legions v2 protocol. */ +abstract class LiquidLegionsV2MillJob : Runnable { @CommandLine.Mixin protected lateinit var flags: LiquidLegionsV2MillFlags private set @@ -120,64 +120,65 @@ abstract class LiquidLegionsV2MillDaemon : Runnable { // This will be the name of the pod when deployed to Kubernetes. Note that the millId is // included in mill logs to help debugging. - val millId = System.getenv("HOSTNAME") + val millId = flags.millId // OpenTelemetry is usually enabled using the Java agent, in which case this will grab the // shared instance. val openTelemetry: OpenTelemetry = GlobalOpenTelemetry.get() - val reachFrequencyliquidLegionsV2Mill = - ReachFrequencyLiquidLegionsV2Mill( - millId = millId, - duchyId = flags.duchy.duchyName, - signingKey = csSigningKey, - consentSignalCert = csCertificate, - trustedCertificates = flags.tlsFlags.signingCerts.trustedCertificates, - dataClients = dataClients, - systemComputationParticipantsClient = systemComputationParticipantsClient, - systemComputationsClient = systemComputationsClient, - systemComputationLogEntriesClient = systemComputationLogEntriesClient, - computationStatsClient = computationStatsClient, - workerStubs = computationControlClientMap, - cryptoWorker = JniLiquidLegionsV2Encryption(), - workLockDuration = flags.workLockDuration, - openTelemetry = openTelemetry, - requestChunkSizeBytes = flags.requestChunkSizeBytes, - parallelism = flags.parallelism, - ) + val mill: LiquidLegionsV2Mill = + when (flags.computationType) { + ComputationTypeEnum.ComputationType.LIQUID_LEGIONS_SKETCH_AGGREGATION_V2 -> + ReachFrequencyLiquidLegionsV2Mill( + millId = millId, + duchyId = flags.duchy.duchyName, + signingKey = csSigningKey, + consentSignalCert = csCertificate, + trustedCertificates = flags.tlsFlags.signingCerts.trustedCertificates, + dataClients = dataClients, + systemComputationParticipantsClient = systemComputationParticipantsClient, + systemComputationsClient = systemComputationsClient, + systemComputationLogEntriesClient = systemComputationLogEntriesClient, + computationStatsClient = computationStatsClient, + workerStubs = computationControlClientMap, + cryptoWorker = JniLiquidLegionsV2Encryption(), + workLockDuration = flags.workLockDuration, + openTelemetry = openTelemetry, + requestChunkSizeBytes = flags.requestChunkSizeBytes, + parallelism = flags.parallelism, + ) + ComputationTypeEnum.ComputationType.REACH_ONLY_LIQUID_LEGIONS_SKETCH_AGGREGATION_V2 -> + ReachOnlyLiquidLegionsV2Mill( + millId = millId, + duchyId = flags.duchy.duchyName, + signingKey = csSigningKey, + consentSignalCert = csCertificate, + trustedCertificates = flags.tlsFlags.signingCerts.trustedCertificates, + dataClients = dataClients, + systemComputationParticipantsClient = systemComputationParticipantsClient, + systemComputationsClient = systemComputationsClient, + systemComputationLogEntriesClient = systemComputationLogEntriesClient, + computationStatsClient = computationStatsClient, + workerStubs = computationControlClientMap, + cryptoWorker = JniReachOnlyLiquidLegionsV2Encryption(), + workLockDuration = flags.workLockDuration, + openTelemetry = openTelemetry, + requestChunkSizeBytes = flags.requestChunkSizeBytes, + parallelism = flags.parallelism, + ) + ComputationTypeEnum.ComputationType.HONEST_MAJORITY_SHARE_SHUFFLE, + ComputationTypeEnum.ComputationType.UNSPECIFIED, + ComputationTypeEnum.ComputationType.UNRECOGNIZED -> + error("Unsupported ComputationType ${flags.computationType}") + } - val reachOnlyLiquidLegionsV2Mill = - ReachOnlyLiquidLegionsV2Mill( - millId = millId, - duchyId = flags.duchy.duchyName, - signingKey = csSigningKey, - consentSignalCert = csCertificate, - trustedCertificates = flags.tlsFlags.signingCerts.trustedCertificates, - dataClients = dataClients, - systemComputationParticipantsClient = systemComputationParticipantsClient, - systemComputationsClient = systemComputationsClient, - systemComputationLogEntriesClient = systemComputationLogEntriesClient, - computationStatsClient = computationStatsClient, - workerStubs = computationControlClientMap, - cryptoWorker = JniReachOnlyLiquidLegionsV2Encryption(), - workLockDuration = flags.workLockDuration, - openTelemetry = openTelemetry, - requestChunkSizeBytes = flags.requestChunkSizeBytes, - parallelism = flags.parallelism, - ) + runBlocking(CoroutineName("Mill $millId")) { + mill.processClaimedWork(flags.claimedGlobalComputationId) - runBlocking { - withContext(CoroutineName("Mill $millId")) { - val throttler = MinimumIntervalThrottler(Clock.systemUTC(), flags.pollingInterval) - throttler.loopOnReady { - logAndSuppressExceptionSuspend { - reachFrequencyliquidLegionsV2Mill.pollAndProcessNextComputation() - } - logAndSuppressExceptionSuspend { - reachOnlyLiquidLegionsV2Mill.pollAndProcessNextComputation() - } - } - } + // Continue processing until work is exhausted. + do { + val workExhausted = !mill.claimAndProcessWork() + } while (isActive && !workExhausted) } } } diff --git a/src/main/kotlin/org/wfanet/measurement/duchy/deploy/gcloud/daemon/mill/liquidlegionsv2/BUILD.bazel b/src/main/kotlin/org/wfanet/measurement/duchy/deploy/gcloud/job/mill/liquidlegionsv2/BUILD.bazel similarity index 67% rename from src/main/kotlin/org/wfanet/measurement/duchy/deploy/gcloud/daemon/mill/liquidlegionsv2/BUILD.bazel rename to src/main/kotlin/org/wfanet/measurement/duchy/deploy/gcloud/job/mill/liquidlegionsv2/BUILD.bazel index d1fc1e4009f..03b4652dac4 100644 --- a/src/main/kotlin/org/wfanet/measurement/duchy/deploy/gcloud/daemon/mill/liquidlegionsv2/BUILD.bazel +++ b/src/main/kotlin/org/wfanet/measurement/duchy/deploy/gcloud/job/mill/liquidlegionsv2/BUILD.bazel @@ -14,10 +14,10 @@ package( ) kt_jvm_library( - name = "gcs_liquid_legions_v2_mill_daemon", - srcs = ["GcsLiquidLegionsV2MillDaemon.kt"], + name = "gcs_liquid_legions_v2_mill_job", + srcs = ["GcsLiquidLegionsV2MillJob.kt"], deps = [ - "//src/main/kotlin/org/wfanet/measurement/duchy/deploy/common/daemon/mill/liquidlegionsv2:liquid_legions_v2_mill_daemon", + "//src/main/kotlin/org/wfanet/measurement/duchy/deploy/common/job/mill/liquidlegionsv2:liquid_legions_v2_mill_job", "//src/main/kotlin/org/wfanet/measurement/duchy/deploy/gcloud/spanner/computation", "@wfa_common_jvm//imports/java/picocli", "@wfa_common_jvm//src/main/kotlin/org/wfanet/measurement/common", @@ -26,18 +26,18 @@ kt_jvm_library( ) java_binary( - name = "GcsLiquidLegionsV2MillDaemon", - main_class = "org.wfanet.measurement.duchy.deploy.gcloud.daemon.mill.liquidlegionsv2.GcsLiquidLegionsV2MillDaemonKt", + name = "GcsLiquidLegionsV2MillJob", + main_class = "org.wfanet.measurement.duchy.deploy.gcloud.job.mill.liquidlegionsv2.GcsLiquidLegionsV2MillJobKt", runtime_deps = [ - ":gcs_liquid_legions_v2_mill_daemon", + ":gcs_liquid_legions_v2_mill_job", "@wfa_common_jvm//src/main/kotlin/org/wfanet/measurement/gcloud/logging", ], ) java_image( - name = "gcs_liquid_legions_v2_mill_daemon_image", - binary = ":GcsLiquidLegionsV2MillDaemon", - main_class = "org.wfanet.measurement.duchy.deploy.gcloud.daemon.mill.liquidlegionsv2.GcsLiquidLegionsV2MillDaemonKt", + name = "gcs_liquid_legions_v2_mill_job_image", + binary = ":GcsLiquidLegionsV2MillJob", + main_class = "org.wfanet.measurement.duchy.deploy.gcloud.job.mill.liquidlegionsv2.GcsLiquidLegionsV2MillJobKt", target_compatible_with = DISTROLESS_JAVA, visibility = ["//src:docker_image_deployment"], ) diff --git a/src/main/kotlin/org/wfanet/measurement/duchy/deploy/gcloud/daemon/mill/liquidlegionsv2/GcsLiquidLegionsV2MillDaemon.kt b/src/main/kotlin/org/wfanet/measurement/duchy/deploy/gcloud/job/mill/liquidlegionsv2/GcsLiquidLegionsV2MillJob.kt similarity index 75% rename from src/main/kotlin/org/wfanet/measurement/duchy/deploy/gcloud/daemon/mill/liquidlegionsv2/GcsLiquidLegionsV2MillDaemon.kt rename to src/main/kotlin/org/wfanet/measurement/duchy/deploy/gcloud/job/mill/liquidlegionsv2/GcsLiquidLegionsV2MillJob.kt index be4a06b0efd..119d5e54c2a 100644 --- a/src/main/kotlin/org/wfanet/measurement/duchy/deploy/gcloud/daemon/mill/liquidlegionsv2/GcsLiquidLegionsV2MillDaemon.kt +++ b/src/main/kotlin/org/wfanet/measurement/duchy/deploy/gcloud/job/mill/liquidlegionsv2/GcsLiquidLegionsV2MillJob.kt @@ -12,21 +12,21 @@ // See the License for the specific language governing permissions and // limitations under the License. -package org.wfanet.measurement.duchy.deploy.gcloud.daemon.mill.liquidlegionsv2 +package org.wfanet.measurement.duchy.deploy.gcloud.job.mill.liquidlegionsv2 import org.wfanet.measurement.common.commandLineMain -import org.wfanet.measurement.duchy.deploy.common.daemon.mill.liquidlegionsv2.LiquidLegionsV2MillDaemon +import org.wfanet.measurement.duchy.deploy.common.job.mill.liquidlegionsv2.LiquidLegionsV2MillJob import org.wfanet.measurement.gcloud.gcs.GcsFromFlags import org.wfanet.measurement.gcloud.gcs.GcsStorageClient import picocli.CommandLine @CommandLine.Command( - name = "GcsLiquidLegionsV2MillDaemon", - description = ["Liquid Legions V2 Mill daemon."], + name = "GcsLiquidLegionsV2MillJob", + description = ["Liquid Legions V2 Mill job."], mixinStandardHelpOptions = true, showDefaultValues = true, ) -class GcsLiquidLegionsV2MillDaemon : LiquidLegionsV2MillDaemon() { +class GcsLiquidLegionsV2MillJob : LiquidLegionsV2MillJob() { @CommandLine.Mixin private lateinit var gcsFlags: GcsFromFlags.Flags override fun run() { @@ -35,4 +35,4 @@ class GcsLiquidLegionsV2MillDaemon : LiquidLegionsV2MillDaemon() { } } -fun main(args: Array) = commandLineMain(GcsLiquidLegionsV2MillDaemon(), args) +fun main(args: Array) = commandLineMain(GcsLiquidLegionsV2MillJob(), args) diff --git a/src/main/kotlin/org/wfanet/measurement/duchy/daemon/herald/BUILD.bazel b/src/main/kotlin/org/wfanet/measurement/duchy/herald/BUILD.bazel similarity index 92% rename from src/main/kotlin/org/wfanet/measurement/duchy/daemon/herald/BUILD.bazel rename to src/main/kotlin/org/wfanet/measurement/duchy/herald/BUILD.bazel index 49ba8ba95b6..5a672b9dc53 100644 --- a/src/main/kotlin/org/wfanet/measurement/duchy/daemon/herald/BUILD.bazel +++ b/src/main/kotlin/org/wfanet/measurement/duchy/herald/BUILD.bazel @@ -32,9 +32,9 @@ kt_jvm_library( "continuation_token_manager", "//src/main/kotlin/org/wfanet/measurement/api:public_api_version", "//src/main/kotlin/org/wfanet/measurement/duchy:computation_stage", - "//src/main/kotlin/org/wfanet/measurement/duchy/daemon/utils:computation_conversions", - "//src/main/kotlin/org/wfanet/measurement/duchy/daemon/utils:duchy_order", "//src/main/kotlin/org/wfanet/measurement/duchy/db/computation", + "//src/main/kotlin/org/wfanet/measurement/duchy/utils:computation_conversions", + "//src/main/kotlin/org/wfanet/measurement/duchy/utils:duchy_order", "//src/main/kotlin/org/wfanet/measurement/system/v1alpha:resource_key", "//src/main/proto/wfa/measurement/api/v2alpha:crypto_kt_jvm_proto", "//src/main/proto/wfa/measurement/api/v2alpha:measurement_kt_jvm_proto", diff --git a/src/main/kotlin/org/wfanet/measurement/duchy/daemon/herald/ContinuationTokenManager.kt b/src/main/kotlin/org/wfanet/measurement/duchy/herald/ContinuationTokenManager.kt similarity index 98% rename from src/main/kotlin/org/wfanet/measurement/duchy/daemon/herald/ContinuationTokenManager.kt rename to src/main/kotlin/org/wfanet/measurement/duchy/herald/ContinuationTokenManager.kt index e2d14c866ce..ef0dbc99ee8 100644 --- a/src/main/kotlin/org/wfanet/measurement/duchy/daemon/herald/ContinuationTokenManager.kt +++ b/src/main/kotlin/org/wfanet/measurement/duchy/herald/ContinuationTokenManager.kt @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -package org.wfanet.measurement.duchy.daemon.herald +package org.wfanet.measurement.duchy.herald import io.grpc.Status import io.grpc.StatusException diff --git a/src/main/kotlin/org/wfanet/measurement/duchy/daemon/herald/Herald.kt b/src/main/kotlin/org/wfanet/measurement/duchy/herald/Herald.kt similarity index 99% rename from src/main/kotlin/org/wfanet/measurement/duchy/daemon/herald/Herald.kt rename to src/main/kotlin/org/wfanet/measurement/duchy/herald/Herald.kt index 865690d753d..a14a649d113 100644 --- a/src/main/kotlin/org/wfanet/measurement/duchy/daemon/herald/Herald.kt +++ b/src/main/kotlin/org/wfanet/measurement/duchy/herald/Herald.kt @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -package org.wfanet.measurement.duchy.daemon.herald +package org.wfanet.measurement.duchy.herald import io.grpc.Status import io.grpc.StatusException @@ -33,8 +33,8 @@ import org.wfanet.measurement.common.crypto.tink.TinkKeyId import org.wfanet.measurement.common.crypto.tink.TinkPrivateKeyHandle import org.wfanet.measurement.common.grpc.grpcStatusCode import org.wfanet.measurement.common.protoTimestamp -import org.wfanet.measurement.duchy.daemon.utils.key import org.wfanet.measurement.duchy.service.internal.computations.toGetTokenRequest +import org.wfanet.measurement.duchy.utils.key import org.wfanet.measurement.internal.duchy.ComputationDetails import org.wfanet.measurement.internal.duchy.ComputationsGrpcKt.ComputationsCoroutineStub import org.wfanet.measurement.internal.duchy.config.ProtocolsSetupConfig diff --git a/src/main/kotlin/org/wfanet/measurement/duchy/daemon/herald/HonestMajorityShareShuffleStarter.kt b/src/main/kotlin/org/wfanet/measurement/duchy/herald/HonestMajorityShareShuffleStarter.kt similarity index 96% rename from src/main/kotlin/org/wfanet/measurement/duchy/daemon/herald/HonestMajorityShareShuffleStarter.kt rename to src/main/kotlin/org/wfanet/measurement/duchy/herald/HonestMajorityShareShuffleStarter.kt index 45159398e35..a1c4105eb70 100644 --- a/src/main/kotlin/org/wfanet/measurement/duchy/daemon/herald/HonestMajorityShareShuffleStarter.kt +++ b/src/main/kotlin/org/wfanet/measurement/duchy/herald/HonestMajorityShareShuffleStarter.kt @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -package org.wfanet.measurement.duchy.daemon.herald +package org.wfanet.measurement.duchy.herald import com.google.protobuf.ByteString import com.google.protobuf.kotlin.toByteString @@ -24,12 +24,12 @@ import org.wfanet.measurement.api.v2alpha.MeasurementSpec import org.wfanet.measurement.common.crypto.PrivateKeyStore import org.wfanet.measurement.common.crypto.tink.TinkKeyId import org.wfanet.measurement.common.crypto.tink.TinkPrivateKeyHandle -import org.wfanet.measurement.duchy.daemon.utils.key -import org.wfanet.measurement.duchy.daemon.utils.toDuchyDifferentialPrivacyParams -import org.wfanet.measurement.duchy.daemon.utils.toKingdomComputationDetails -import org.wfanet.measurement.duchy.daemon.utils.toRequisitionEntries import org.wfanet.measurement.duchy.db.computation.advanceComputationStage import org.wfanet.measurement.duchy.toProtocolStage +import org.wfanet.measurement.duchy.utils.key +import org.wfanet.measurement.duchy.utils.toDuchyDifferentialPrivacyParams +import org.wfanet.measurement.duchy.utils.toKingdomComputationDetails +import org.wfanet.measurement.duchy.utils.toRequisitionEntries import org.wfanet.measurement.internal.duchy.ComputationToken import org.wfanet.measurement.internal.duchy.ComputationTypeEnum import org.wfanet.measurement.internal.duchy.ComputationsGrpcKt diff --git a/src/main/kotlin/org/wfanet/measurement/duchy/daemon/herald/LiquidLegionsV2Starter.kt b/src/main/kotlin/org/wfanet/measurement/duchy/herald/LiquidLegionsV2Starter.kt similarity index 97% rename from src/main/kotlin/org/wfanet/measurement/duchy/daemon/herald/LiquidLegionsV2Starter.kt rename to src/main/kotlin/org/wfanet/measurement/duchy/herald/LiquidLegionsV2Starter.kt index 293cfe9049e..32b11f0b68c 100644 --- a/src/main/kotlin/org/wfanet/measurement/duchy/daemon/herald/LiquidLegionsV2Starter.kt +++ b/src/main/kotlin/org/wfanet/measurement/duchy/herald/LiquidLegionsV2Starter.kt @@ -12,21 +12,21 @@ // See the License for the specific language governing permissions and // limitations under the License. -package org.wfanet.measurement.duchy.daemon.herald +package org.wfanet.measurement.duchy.herald import java.util.logging.Level import java.util.logging.Logger import org.wfanet.measurement.api.Version import org.wfanet.measurement.api.v2alpha.DifferentialPrivacyParams import org.wfanet.measurement.api.v2alpha.MeasurementSpec -import org.wfanet.measurement.duchy.daemon.utils.key -import org.wfanet.measurement.duchy.daemon.utils.sha1Hash -import org.wfanet.measurement.duchy.daemon.utils.toDuchyDifferentialPrivacyParams -import org.wfanet.measurement.duchy.daemon.utils.toDuchyElGamalPublicKey -import org.wfanet.measurement.duchy.daemon.utils.toKingdomComputationDetails -import org.wfanet.measurement.duchy.daemon.utils.toRequisitionEntries import org.wfanet.measurement.duchy.db.computation.advanceComputationStage import org.wfanet.measurement.duchy.toProtocolStage +import org.wfanet.measurement.duchy.utils.key +import org.wfanet.measurement.duchy.utils.sha1Hash +import org.wfanet.measurement.duchy.utils.toDuchyDifferentialPrivacyParams +import org.wfanet.measurement.duchy.utils.toDuchyElGamalPublicKey +import org.wfanet.measurement.duchy.utils.toKingdomComputationDetails +import org.wfanet.measurement.duchy.utils.toRequisitionEntries import org.wfanet.measurement.internal.duchy.ComputationDetails import org.wfanet.measurement.internal.duchy.ComputationToken import org.wfanet.measurement.internal.duchy.ComputationTypeEnum diff --git a/src/main/kotlin/org/wfanet/measurement/duchy/daemon/herald/ReachOnlyLiquidLegionsV2Starter.kt b/src/main/kotlin/org/wfanet/measurement/duchy/herald/ReachOnlyLiquidLegionsV2Starter.kt similarity index 97% rename from src/main/kotlin/org/wfanet/measurement/duchy/daemon/herald/ReachOnlyLiquidLegionsV2Starter.kt rename to src/main/kotlin/org/wfanet/measurement/duchy/herald/ReachOnlyLiquidLegionsV2Starter.kt index 815fb181c1b..2a37c223ff2 100644 --- a/src/main/kotlin/org/wfanet/measurement/duchy/daemon/herald/ReachOnlyLiquidLegionsV2Starter.kt +++ b/src/main/kotlin/org/wfanet/measurement/duchy/herald/ReachOnlyLiquidLegionsV2Starter.kt @@ -12,20 +12,20 @@ // See the License for the specific language governing permissions and // limitations under the License. -package org.wfanet.measurement.duchy.daemon.herald +package org.wfanet.measurement.duchy.herald import java.util.logging.Logger import org.wfanet.measurement.api.Version import org.wfanet.measurement.api.v2alpha.MeasurementSpec -import org.wfanet.measurement.duchy.daemon.utils.key -import org.wfanet.measurement.duchy.daemon.utils.sha1Hash -import org.wfanet.measurement.duchy.daemon.utils.toDuchyDifferentialPrivacyParams -import org.wfanet.measurement.duchy.daemon.utils.toDuchyElGamalPublicKey -import org.wfanet.measurement.duchy.daemon.utils.toKingdomComputationDetails -import org.wfanet.measurement.duchy.daemon.utils.toRequisitionEntries import org.wfanet.measurement.duchy.db.computation.advanceComputationStage import org.wfanet.measurement.duchy.service.internal.computations.outputPathList import org.wfanet.measurement.duchy.toProtocolStage +import org.wfanet.measurement.duchy.utils.key +import org.wfanet.measurement.duchy.utils.sha1Hash +import org.wfanet.measurement.duchy.utils.toDuchyDifferentialPrivacyParams +import org.wfanet.measurement.duchy.utils.toDuchyElGamalPublicKey +import org.wfanet.measurement.duchy.utils.toKingdomComputationDetails +import org.wfanet.measurement.duchy.utils.toRequisitionEntries import org.wfanet.measurement.internal.duchy.ComputationToken import org.wfanet.measurement.internal.duchy.ComputationTypeEnum import org.wfanet.measurement.internal.duchy.ComputationsGrpcKt diff --git a/src/main/kotlin/org/wfanet/measurement/duchy/daemon/mill/BUILD.bazel b/src/main/kotlin/org/wfanet/measurement/duchy/mill/BUILD.bazel similarity index 78% rename from src/main/kotlin/org/wfanet/measurement/duchy/daemon/mill/BUILD.bazel rename to src/main/kotlin/org/wfanet/measurement/duchy/mill/BUILD.bazel index b102af78edd..73bd57e54df 100644 --- a/src/main/kotlin/org/wfanet/measurement/duchy/daemon/mill/BUILD.bazel +++ b/src/main/kotlin/org/wfanet/measurement/duchy/mill/BUILD.bazel @@ -13,8 +13,8 @@ kt_jvm_library( deps = [ "//src/main/kotlin/org/wfanet/measurement/api:public_api_version", "//src/main/kotlin/org/wfanet/measurement/duchy:computation_stage", - "//src/main/kotlin/org/wfanet/measurement/duchy/daemon/utils:computation_conversions", "//src/main/kotlin/org/wfanet/measurement/duchy/db/computation", + "//src/main/kotlin/org/wfanet/measurement/duchy/utils:computation_conversions", "//src/main/kotlin/org/wfanet/measurement/system/v1alpha:resource_key", "//src/main/proto/wfa/measurement/internal/duchy:computation_stats_service_kt_jvm_grpc_proto", "//src/main/proto/wfa/measurement/internal/duchy:computation_token_kt_jvm_proto", @@ -28,3 +28,14 @@ kt_jvm_library( "@wfa_consent_signaling_client//src/main/kotlin/org/wfanet/measurement/consent/client/duchy", ], ) + +kt_jvm_library( + name = "mill_flags", + srcs = ["MillFlags.kt"], + deps = [ + "//src/main/kotlin/org/wfanet/measurement/common/identity", + "@wfa_common_jvm//imports/java/picocli", + "@wfa_common_jvm//src/main/kotlin/org/wfanet/measurement/common", + "@wfa_common_jvm//src/main/kotlin/org/wfanet/measurement/common/grpc", + ], +) diff --git a/src/main/kotlin/org/wfanet/measurement/duchy/daemon/mill/MillBase.kt b/src/main/kotlin/org/wfanet/measurement/duchy/mill/MillBase.kt similarity index 95% rename from src/main/kotlin/org/wfanet/measurement/duchy/daemon/mill/MillBase.kt rename to src/main/kotlin/org/wfanet/measurement/duchy/mill/MillBase.kt index 35799d3cd90..71b789b5a34 100644 --- a/src/main/kotlin/org/wfanet/measurement/duchy/daemon/mill/MillBase.kt +++ b/src/main/kotlin/org/wfanet/measurement/duchy/mill/MillBase.kt @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -package org.wfanet.measurement.duchy.daemon.mill +package org.wfanet.measurement.duchy.mill import com.google.protobuf.ByteString import io.grpc.Status @@ -50,13 +50,13 @@ import org.wfanet.measurement.common.protoTimestamp import org.wfanet.measurement.common.toProtoDuration import org.wfanet.measurement.consent.client.duchy.encryptResult import org.wfanet.measurement.consent.client.duchy.signResult -import org.wfanet.measurement.duchy.daemon.utils.ComputationResult -import org.wfanet.measurement.duchy.daemon.utils.toV2AlphaEncryptionPublicKey import org.wfanet.measurement.duchy.db.computation.BlobRef import org.wfanet.measurement.duchy.db.computation.ComputationDataClients import org.wfanet.measurement.duchy.db.computation.singleOutputBlobMetadata import org.wfanet.measurement.duchy.name import org.wfanet.measurement.duchy.number +import org.wfanet.measurement.duchy.utils.ComputationResult +import org.wfanet.measurement.duchy.utils.toV2AlphaEncryptionPublicKey import org.wfanet.measurement.internal.duchy.ComputationDetails.CompletedReason import org.wfanet.measurement.internal.duchy.ComputationStage import org.wfanet.measurement.internal.duchy.ComputationStatsGrpcKt.ComputationStatsCoroutineStub @@ -80,7 +80,7 @@ import org.wfanet.measurement.system.v1alpha.ComputationParticipant import org.wfanet.measurement.system.v1alpha.ComputationParticipantKey import org.wfanet.measurement.system.v1alpha.ComputationParticipantKt import org.wfanet.measurement.system.v1alpha.ComputationParticipantsGrpcKt.ComputationParticipantsCoroutineStub -import org.wfanet.measurement.system.v1alpha.ComputationsGrpcKt.ComputationsCoroutineStub as SystemComputationsCoroutineStub +import org.wfanet.measurement.system.v1alpha.ComputationsGrpcKt import org.wfanet.measurement.system.v1alpha.CreateComputationLogEntryRequest import org.wfanet.measurement.system.v1alpha.failComputationParticipantRequest import org.wfanet.measurement.system.v1alpha.getComputationParticipantRequest @@ -117,10 +117,10 @@ abstract class MillBase( protected val consentSignalCert: Certificate, protected val dataClients: ComputationDataClients, protected val systemComputationParticipantsClient: ComputationParticipantsCoroutineStub, - private val systemComputationsClient: SystemComputationsCoroutineStub, + private val systemComputationsClient: ComputationsGrpcKt.ComputationsCoroutineStub, private val systemComputationLogEntriesClient: ComputationLogEntriesCoroutineStub, private val computationStatsClient: ComputationStatsCoroutineStub, - private val computationType: ComputationType, + protected val computationType: ComputationType, private val workLockDuration: Duration, private val requestChunkSizeBytes: Int, private val maximumAttempts: Int, @@ -164,9 +164,13 @@ abstract class MillBase( .build() private var computationsServerReady = false - /** Poll and work on the next available computations. */ - suspend fun pollAndProcessNextComputation() { - logger.fine("@Mill $millId: Polling available computations...") + /** + * Claim the next available work item and process it. + * + * @return whether a work item was claimed + */ + suspend fun claimAndProcessWork(): Boolean { + logger.fine("@Mill $millId: Polling available work...") val claimWorkRequest = claimWorkRequest { computationType = this@MillBase.computationType @@ -179,38 +183,30 @@ abstract class MillBase( dataClients.computationsClient.claimWork(claimWorkRequest) } catch (e: StatusException) { if (!computationsServerReady && e.status.code == Status.Code.UNAVAILABLE) { - logger.info("ComputationServer not ready") - return + logger.info("Computations server not ready") + return false } throw Exception("Error claiming work", e) } computationsServerReady = true - if (claimWorkResponse.hasToken()) { - val token = claimWorkResponse.token - if (token.attempt > maximumAttempts) { - failComputation( - token, - "Failing computation due to too many failed ComputationStageAttempts.", - ) - return - } - - val wallDurationLogger = wallDurationLogger() - - processComputation(token) - wallDurationLogger.logStageDurationMetric( - token, - STAGE_WALL_CLOCK_DURATION, - stageWallClockDurationHistogram, - ) - } else { - logger.fine("@Mill $millId: No computation available, waiting for the next poll...") + if (!claimWorkResponse.hasToken()) { + return false } + + processComputation(claimWorkResponse.token) + return true } /** Process the computation according to its protocol and status. */ - private suspend fun processComputation(token: ComputationToken) { + protected suspend fun processComputation(token: ComputationToken) { + if (token.attempt > maximumAttempts) { + failComputation(token, "Failing computation due to too many failed ComputationStageAttempts.") + return + } + + val wallDurationLogger = wallDurationLogger() + // Log the current mill memory usage before processing. logStageMetric(token, CURRENT_RUNTIME_MEMORY_MAXIMUM, Runtime.getRuntime().maxMemory()) logStageMetric(token, CURRENT_RUNTIME_MEMORY_TOTAL, Runtime.getRuntime().totalMemory()) @@ -232,6 +228,12 @@ abstract class MillBase( } } logger.info("$globalId@$millId: Processed computation ") + + wallDurationLogger.logStageDurationMetric( + token, + STAGE_WALL_CLOCK_DURATION, + stageWallClockDurationHistogram, + ) } private suspend fun handleExceptions(token: ComputationToken, e: Exception) { @@ -680,7 +682,7 @@ abstract class MillBase( } /** Gets the latest [ComputationToken] for computation with [globalId]. */ - private suspend fun getLatestComputationToken(globalId: String): ComputationToken { + protected suspend fun getLatestComputationToken(globalId: String): ComputationToken { val response: GetComputationTokenResponse = try { dataClients.computationsClient.getComputationToken( diff --git a/src/main/kotlin/org/wfanet/measurement/duchy/deploy/common/daemon/mill/liquidlegionsv2/LiquidLegionsV2MillFlags.kt b/src/main/kotlin/org/wfanet/measurement/duchy/mill/MillFlags.kt similarity index 58% rename from src/main/kotlin/org/wfanet/measurement/duchy/deploy/common/daemon/mill/liquidlegionsv2/LiquidLegionsV2MillFlags.kt rename to src/main/kotlin/org/wfanet/measurement/duchy/mill/MillFlags.kt index 3f8d00b08ec..6ddf7c0e099 100644 --- a/src/main/kotlin/org/wfanet/measurement/duchy/deploy/common/daemon/mill/liquidlegionsv2/LiquidLegionsV2MillFlags.kt +++ b/src/main/kotlin/org/wfanet/measurement/duchy/mill/MillFlags.kt @@ -1,34 +1,29 @@ -// Copyright 2020 The Cross-Media Measurement Authors -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package org.wfanet.measurement.duchy.deploy.common.daemon.mill.liquidlegionsv2 +/* + * Copyright 2024 The Cross-Media Measurement Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.wfanet.measurement.duchy.mill import java.io.File import java.time.Duration import kotlin.properties.Delegates import org.wfanet.measurement.common.grpc.TlsFlags import org.wfanet.measurement.common.identity.DuchyInfoFlags -import org.wfanet.measurement.duchy.deploy.common.CommonDuchyFlags -import org.wfanet.measurement.duchy.deploy.common.ComputationsServiceFlags -import org.wfanet.measurement.duchy.deploy.common.SystemApiFlags import picocli.CommandLine -class LiquidLegionsV2MillFlags { - @CommandLine.Mixin - lateinit var duchy: CommonDuchyFlags - private set - +abstract class MillFlags { @CommandLine.Mixin lateinit var tlsFlags: TlsFlags private set @@ -37,6 +32,13 @@ class LiquidLegionsV2MillFlags { lateinit var duchyInfoFlags: DuchyInfoFlags private set + @CommandLine.Option( + names = ["--mill-id"], + description = ["ID of this Mill instance. Defaults to HOSTNAME."], + ) + var millId: String = System.getenv("HOSTNAME") + private set + @CommandLine.Option( names = ["--duchy-computation-control-target"], description = ["Key-value pair of Duchy ID to ComputationControl service target."], @@ -53,14 +55,6 @@ class LiquidLegionsV2MillFlags { lateinit var channelShutdownTimeout: Duration private set - @CommandLine.Option( - names = ["--polling-interval"], - defaultValue = "2s", - description = ["How long to sleep before polling the computation queue again if it is empty."], - ) - lateinit var pollingInterval: Duration - private set - @CommandLine.Option( names = ["--work-lock-duration"], defaultValue = "5m", @@ -69,14 +63,6 @@ class LiquidLegionsV2MillFlags { lateinit var workLockDuration: Duration private set - @CommandLine.Mixin - lateinit var systemApiFlags: SystemApiFlags - private set - - @CommandLine.Mixin - lateinit var computationsServiceFlags: ComputationsServiceFlags - private set - @set:CommandLine.Option( names = ["--bytes-per-chunk"], description = ["The number of bytes in a chunk when sending rpc result to other duchy."], @@ -108,12 +94,4 @@ class LiquidLegionsV2MillFlags { ) lateinit var csCertificateDerFile: File private set - - @set:CommandLine.Option( - names = ["--parallelism"], - description = ["Maximum number of threads used in crypto actions"], - defaultValue = "1", - ) - var parallelism by Delegates.notNull() - private set } diff --git a/src/main/kotlin/org/wfanet/measurement/duchy/daemon/mill/liquidlegionsv2/BUILD.bazel b/src/main/kotlin/org/wfanet/measurement/duchy/mill/liquidlegionsv2/BUILD.bazel similarity index 89% rename from src/main/kotlin/org/wfanet/measurement/duchy/daemon/mill/liquidlegionsv2/BUILD.bazel rename to src/main/kotlin/org/wfanet/measurement/duchy/mill/liquidlegionsv2/BUILD.bazel index 94873a960a5..9707c81b97d 100644 --- a/src/main/kotlin/org/wfanet/measurement/duchy/daemon/mill/liquidlegionsv2/BUILD.bazel +++ b/src/main/kotlin/org/wfanet/measurement/duchy/mill/liquidlegionsv2/BUILD.bazel @@ -22,12 +22,12 @@ kt_jvm_library( deps = [ "//src/main/kotlin/org/wfanet/measurement/api:public_api_version", "//src/main/kotlin/org/wfanet/measurement/common/identity", - "//src/main/kotlin/org/wfanet/measurement/duchy/daemon/mill:mill_base", - "//src/main/kotlin/org/wfanet/measurement/duchy/daemon/mill/liquidlegionsv2/crypto:liquidlegionsv2encryption", - "//src/main/kotlin/org/wfanet/measurement/duchy/daemon/utils:computation_conversions", - "//src/main/kotlin/org/wfanet/measurement/duchy/daemon/utils:duchy_order", "//src/main/kotlin/org/wfanet/measurement/duchy/db/computation", + "//src/main/kotlin/org/wfanet/measurement/duchy/mill:mill_base", + "//src/main/kotlin/org/wfanet/measurement/duchy/mill/liquidlegionsv2/crypto:liquidlegionsv2encryption", "//src/main/kotlin/org/wfanet/measurement/duchy/service/system/v1alpha:advance_computation_request_headers", + "//src/main/kotlin/org/wfanet/measurement/duchy/utils:computation_conversions", + "//src/main/kotlin/org/wfanet/measurement/duchy/utils:duchy_order", "//src/main/kotlin/org/wfanet/measurement/system/v1alpha:resource_key", "//src/main/proto/wfa/any_sketch/crypto:sketch_encryption_methods_kt_jvm_proto", "//src/main/proto/wfa/measurement/internal/duchy:crypto_kt_jvm_proto", @@ -60,10 +60,10 @@ kt_jvm_library( ":liquid_legions_v2_mill", "//src/main/kotlin/org/wfanet/measurement/api:public_api_version", "//src/main/kotlin/org/wfanet/measurement/common/identity", - "//src/main/kotlin/org/wfanet/measurement/duchy/daemon/mill/liquidlegionsv2/crypto:liquidlegionsv2encryption", - "//src/main/kotlin/org/wfanet/measurement/duchy/daemon/utils:duchy_order", "//src/main/kotlin/org/wfanet/measurement/duchy/db/computation", + "//src/main/kotlin/org/wfanet/measurement/duchy/mill/liquidlegionsv2/crypto:liquidlegionsv2encryption", "//src/main/kotlin/org/wfanet/measurement/duchy/service/system/v1alpha:advance_computation_request_headers", + "//src/main/kotlin/org/wfanet/measurement/duchy/utils:duchy_order", "//src/main/kotlin/org/wfanet/measurement/system/v1alpha:resource_key", "//src/main/proto/wfa/any_sketch/crypto:sketch_encryption_methods_kt_jvm_proto", "//src/main/proto/wfa/measurement/internal/duchy:crypto_kt_jvm_proto", @@ -98,11 +98,11 @@ kt_jvm_library( ":liquid_legions_v2_mill", "//src/main/kotlin/org/wfanet/measurement/api:public_api_version", "//src/main/kotlin/org/wfanet/measurement/common/identity", - "//src/main/kotlin/org/wfanet/measurement/duchy/daemon/mill/liquidlegionsv2/crypto:reachonlyliquidlegionsv2encryption", - "//src/main/kotlin/org/wfanet/measurement/duchy/daemon/utils:computation_conversions", - "//src/main/kotlin/org/wfanet/measurement/duchy/daemon/utils:duchy_order", "//src/main/kotlin/org/wfanet/measurement/duchy/db/computation", + "//src/main/kotlin/org/wfanet/measurement/duchy/mill/liquidlegionsv2/crypto:reachonlyliquidlegionsv2encryption", "//src/main/kotlin/org/wfanet/measurement/duchy/service/system/v1alpha:advance_computation_request_headers", + "//src/main/kotlin/org/wfanet/measurement/duchy/utils:computation_conversions", + "//src/main/kotlin/org/wfanet/measurement/duchy/utils:duchy_order", "//src/main/kotlin/org/wfanet/measurement/system/v1alpha:resource_key", "//src/main/proto/wfa/any_sketch/crypto:sketch_encryption_methods_kt_jvm_proto", "//src/main/proto/wfa/measurement/internal/duchy:crypto_kt_jvm_proto", diff --git a/src/main/kotlin/org/wfanet/measurement/duchy/daemon/mill/liquidlegionsv2/LiquidLegionsV2Mill.kt b/src/main/kotlin/org/wfanet/measurement/duchy/mill/liquidlegionsv2/LiquidLegionsV2Mill.kt similarity index 96% rename from src/main/kotlin/org/wfanet/measurement/duchy/daemon/mill/liquidlegionsv2/LiquidLegionsV2Mill.kt rename to src/main/kotlin/org/wfanet/measurement/duchy/mill/liquidlegionsv2/LiquidLegionsV2Mill.kt index f53770fad63..3c960f20e2f 100644 --- a/src/main/kotlin/org/wfanet/measurement/duchy/daemon/mill/liquidlegionsv2/LiquidLegionsV2Mill.kt +++ b/src/main/kotlin/org/wfanet/measurement/duchy/mill/liquidlegionsv2/LiquidLegionsV2Mill.kt @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -package org.wfanet.measurement.duchy.daemon.mill.liquidlegionsv2 +package org.wfanet.measurement.duchy.mill.liquidlegionsv2 import com.google.protobuf.ByteString import io.opentelemetry.api.OpenTelemetry @@ -30,9 +30,9 @@ import org.wfanet.measurement.common.crypto.readCertificate import org.wfanet.measurement.common.identity.DuchyInfo import org.wfanet.measurement.consent.client.duchy.verifyDataProviderParticipation import org.wfanet.measurement.consent.client.duchy.verifyElGamalPublicKey -import org.wfanet.measurement.duchy.daemon.mill.Certificate -import org.wfanet.measurement.duchy.daemon.mill.MillBase import org.wfanet.measurement.duchy.db.computation.ComputationDataClients +import org.wfanet.measurement.duchy.mill.Certificate +import org.wfanet.measurement.duchy.mill.MillBase import org.wfanet.measurement.internal.duchy.ComputationDetails.KingdomComputationDetails import org.wfanet.measurement.internal.duchy.ComputationStatsGrpcKt.ComputationStatsCoroutineStub import org.wfanet.measurement.internal.duchy.ComputationToken @@ -105,6 +105,12 @@ abstract class LiquidLegionsV2Mill( clock, openTelemetry, ) { + + suspend fun processClaimedWork(globalComputationId: String) { + val token: ComputationToken = getLatestComputationToken(globalComputationId) + processComputation(token) + } + /** * Verifies that all EDPs have participated. * diff --git a/src/main/kotlin/org/wfanet/measurement/duchy/daemon/mill/liquidlegionsv2/ReachFrequencyLiquidLegionsV2Mill.kt b/src/main/kotlin/org/wfanet/measurement/duchy/mill/liquidlegionsv2/ReachFrequencyLiquidLegionsV2Mill.kt similarity index 98% rename from src/main/kotlin/org/wfanet/measurement/duchy/daemon/mill/liquidlegionsv2/ReachFrequencyLiquidLegionsV2Mill.kt rename to src/main/kotlin/org/wfanet/measurement/duchy/mill/liquidlegionsv2/ReachFrequencyLiquidLegionsV2Mill.kt index f20ed3fa08d..4e6252af711 100644 --- a/src/main/kotlin/org/wfanet/measurement/duchy/daemon/mill/liquidlegionsv2/ReachFrequencyLiquidLegionsV2Mill.kt +++ b/src/main/kotlin/org/wfanet/measurement/duchy/mill/liquidlegionsv2/ReachFrequencyLiquidLegionsV2Mill.kt @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -package org.wfanet.measurement.duchy.daemon.mill.liquidlegionsv2 +package org.wfanet.measurement.duchy.mill.liquidlegionsv2 import com.google.protobuf.ByteString import io.opentelemetry.api.OpenTelemetry @@ -26,25 +26,24 @@ import org.wfanet.measurement.api.v2alpha.MeasurementSpec import org.wfanet.measurement.common.crypto.SigningKeyHandle import org.wfanet.measurement.common.flatten import org.wfanet.measurement.consent.client.duchy.signElgamalPublicKey -import org.wfanet.measurement.duchy.daemon.mill.CRYPTO_CPU_DURATION -import org.wfanet.measurement.duchy.daemon.mill.Certificate -import org.wfanet.measurement.duchy.daemon.mill.liquidlegionsv2.crypto.LiquidLegionsV2Encryption -import org.wfanet.measurement.duchy.daemon.utils.ReachAndFrequencyResult -import org.wfanet.measurement.duchy.daemon.utils.ReachResult -import org.wfanet.measurement.duchy.daemon.utils.toAnySketchElGamalPublicKey -import org.wfanet.measurement.duchy.daemon.utils.toCmmsElGamalPublicKey -import org.wfanet.measurement.duchy.daemon.utils.toV2AlphaElGamalPublicKey import org.wfanet.measurement.duchy.db.computation.ComputationDataClients +import org.wfanet.measurement.duchy.mill.CRYPTO_CPU_DURATION +import org.wfanet.measurement.duchy.mill.Certificate +import org.wfanet.measurement.duchy.mill.liquidlegionsv2.crypto.LiquidLegionsV2Encryption import org.wfanet.measurement.duchy.service.internal.computations.outputPathList import org.wfanet.measurement.duchy.service.system.v1alpha.advanceComputationHeader import org.wfanet.measurement.duchy.toProtocolStage +import org.wfanet.measurement.duchy.utils.ReachAndFrequencyResult +import org.wfanet.measurement.duchy.utils.ReachResult +import org.wfanet.measurement.duchy.utils.toAnySketchElGamalPublicKey +import org.wfanet.measurement.duchy.utils.toCmmsElGamalPublicKey +import org.wfanet.measurement.duchy.utils.toV2AlphaElGamalPublicKey import org.wfanet.measurement.internal.duchy.ComputationDetails.CompletedReason import org.wfanet.measurement.internal.duchy.ComputationStatsGrpcKt.ComputationStatsCoroutineStub import org.wfanet.measurement.internal.duchy.ComputationToken import org.wfanet.measurement.internal.duchy.ComputationTypeEnum.ComputationType import org.wfanet.measurement.internal.duchy.ElGamalPublicKey import org.wfanet.measurement.internal.duchy.UpdateComputationDetailsRequest -import org.wfanet.measurement.internal.duchy.computationStage import org.wfanet.measurement.internal.duchy.config.RoleInComputation import org.wfanet.measurement.internal.duchy.config.RoleInComputation.AGGREGATOR import org.wfanet.measurement.internal.duchy.config.RoleInComputation.NON_AGGREGATOR diff --git a/src/main/kotlin/org/wfanet/measurement/duchy/daemon/mill/liquidlegionsv2/ReachOnlyLiquidLegionsV2Mill.kt b/src/main/kotlin/org/wfanet/measurement/duchy/mill/liquidlegionsv2/ReachOnlyLiquidLegionsV2Mill.kt similarity index 97% rename from src/main/kotlin/org/wfanet/measurement/duchy/daemon/mill/liquidlegionsv2/ReachOnlyLiquidLegionsV2Mill.kt rename to src/main/kotlin/org/wfanet/measurement/duchy/mill/liquidlegionsv2/ReachOnlyLiquidLegionsV2Mill.kt index cca1e857a0d..10f40664199 100644 --- a/src/main/kotlin/org/wfanet/measurement/duchy/daemon/mill/liquidlegionsv2/ReachOnlyLiquidLegionsV2Mill.kt +++ b/src/main/kotlin/org/wfanet/measurement/duchy/mill/liquidlegionsv2/ReachOnlyLiquidLegionsV2Mill.kt @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -package org.wfanet.measurement.duchy.daemon.mill.liquidlegionsv2 +package org.wfanet.measurement.duchy.mill.liquidlegionsv2 import com.google.protobuf.ByteString import io.opentelemetry.api.OpenTelemetry @@ -25,25 +25,24 @@ import org.wfanet.measurement.api.Version import org.wfanet.measurement.api.v2alpha.MeasurementSpec import org.wfanet.measurement.common.crypto.SigningKeyHandle import org.wfanet.measurement.consent.client.duchy.signElgamalPublicKey -import org.wfanet.measurement.duchy.daemon.mill.CRYPTO_CPU_DURATION -import org.wfanet.measurement.duchy.daemon.mill.Certificate -import org.wfanet.measurement.duchy.daemon.mill.liquidlegionsv2.crypto.ReachOnlyLiquidLegionsV2Encryption -import org.wfanet.measurement.duchy.daemon.utils.ReachResult -import org.wfanet.measurement.duchy.daemon.utils.toAnySketchElGamalPublicKey -import org.wfanet.measurement.duchy.daemon.utils.toCmmsElGamalPublicKey -import org.wfanet.measurement.duchy.daemon.utils.toV2AlphaElGamalPublicKey import org.wfanet.measurement.duchy.db.computation.BlobRef import org.wfanet.measurement.duchy.db.computation.ComputationDataClients +import org.wfanet.measurement.duchy.mill.CRYPTO_CPU_DURATION +import org.wfanet.measurement.duchy.mill.Certificate +import org.wfanet.measurement.duchy.mill.liquidlegionsv2.crypto.ReachOnlyLiquidLegionsV2Encryption import org.wfanet.measurement.duchy.service.internal.computations.outputPathList import org.wfanet.measurement.duchy.service.system.v1alpha.advanceComputationHeader import org.wfanet.measurement.duchy.toProtocolStage +import org.wfanet.measurement.duchy.utils.ReachResult +import org.wfanet.measurement.duchy.utils.toAnySketchElGamalPublicKey +import org.wfanet.measurement.duchy.utils.toCmmsElGamalPublicKey +import org.wfanet.measurement.duchy.utils.toV2AlphaElGamalPublicKey import org.wfanet.measurement.internal.duchy.ComputationDetails.CompletedReason import org.wfanet.measurement.internal.duchy.ComputationStatsGrpcKt.ComputationStatsCoroutineStub import org.wfanet.measurement.internal.duchy.ComputationToken import org.wfanet.measurement.internal.duchy.ComputationTypeEnum.ComputationType import org.wfanet.measurement.internal.duchy.ElGamalPublicKey import org.wfanet.measurement.internal.duchy.computationDetails -import org.wfanet.measurement.internal.duchy.computationStage import org.wfanet.measurement.internal.duchy.config.RoleInComputation import org.wfanet.measurement.internal.duchy.config.RoleInComputation.AGGREGATOR import org.wfanet.measurement.internal.duchy.config.RoleInComputation.NON_AGGREGATOR diff --git a/src/main/kotlin/org/wfanet/measurement/duchy/daemon/mill/liquidlegionsv2/crypto/BUILD.bazel b/src/main/kotlin/org/wfanet/measurement/duchy/mill/liquidlegionsv2/crypto/BUILD.bazel similarity index 100% rename from src/main/kotlin/org/wfanet/measurement/duchy/daemon/mill/liquidlegionsv2/crypto/BUILD.bazel rename to src/main/kotlin/org/wfanet/measurement/duchy/mill/liquidlegionsv2/crypto/BUILD.bazel diff --git a/src/main/kotlin/org/wfanet/measurement/duchy/daemon/mill/liquidlegionsv2/crypto/JniLiquidLegionsV2Encryption.kt b/src/main/kotlin/org/wfanet/measurement/duchy/mill/liquidlegionsv2/crypto/JniLiquidLegionsV2Encryption.kt similarity index 98% rename from src/main/kotlin/org/wfanet/measurement/duchy/daemon/mill/liquidlegionsv2/crypto/JniLiquidLegionsV2Encryption.kt rename to src/main/kotlin/org/wfanet/measurement/duchy/mill/liquidlegionsv2/crypto/JniLiquidLegionsV2Encryption.kt index e092ff269d1..9f3a45406a2 100644 --- a/src/main/kotlin/org/wfanet/measurement/duchy/daemon/mill/liquidlegionsv2/crypto/JniLiquidLegionsV2Encryption.kt +++ b/src/main/kotlin/org/wfanet/measurement/duchy/mill/liquidlegionsv2/crypto/JniLiquidLegionsV2Encryption.kt @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -package org.wfanet.measurement.duchy.daemon.mill.liquidlegionsv2.crypto +package org.wfanet.measurement.duchy.mill.liquidlegionsv2.crypto import org.wfanet.anysketch.crypto.CombineElGamalPublicKeysRequest import org.wfanet.anysketch.crypto.CombineElGamalPublicKeysResponse diff --git a/src/main/kotlin/org/wfanet/measurement/duchy/daemon/mill/liquidlegionsv2/crypto/JniReachOnlyLiquidLegionsV2Encryption.kt b/src/main/kotlin/org/wfanet/measurement/duchy/mill/liquidlegionsv2/crypto/JniReachOnlyLiquidLegionsV2Encryption.kt similarity index 98% rename from src/main/kotlin/org/wfanet/measurement/duchy/daemon/mill/liquidlegionsv2/crypto/JniReachOnlyLiquidLegionsV2Encryption.kt rename to src/main/kotlin/org/wfanet/measurement/duchy/mill/liquidlegionsv2/crypto/JniReachOnlyLiquidLegionsV2Encryption.kt index 376dc8ca1a6..a907a292db9 100644 --- a/src/main/kotlin/org/wfanet/measurement/duchy/daemon/mill/liquidlegionsv2/crypto/JniReachOnlyLiquidLegionsV2Encryption.kt +++ b/src/main/kotlin/org/wfanet/measurement/duchy/mill/liquidlegionsv2/crypto/JniReachOnlyLiquidLegionsV2Encryption.kt @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -package org.wfanet.measurement.duchy.daemon.mill.liquidlegionsv2.crypto +package org.wfanet.measurement.duchy.mill.liquidlegionsv2.crypto import org.wfanet.anysketch.crypto.CombineElGamalPublicKeysRequest import org.wfanet.anysketch.crypto.CombineElGamalPublicKeysResponse diff --git a/src/main/kotlin/org/wfanet/measurement/duchy/daemon/mill/liquidlegionsv2/crypto/LiquidLegionsV2Encryption.kt b/src/main/kotlin/org/wfanet/measurement/duchy/mill/liquidlegionsv2/crypto/LiquidLegionsV2Encryption.kt similarity index 98% rename from src/main/kotlin/org/wfanet/measurement/duchy/daemon/mill/liquidlegionsv2/crypto/LiquidLegionsV2Encryption.kt rename to src/main/kotlin/org/wfanet/measurement/duchy/mill/liquidlegionsv2/crypto/LiquidLegionsV2Encryption.kt index 547bb83edeb..26c787da2d5 100644 --- a/src/main/kotlin/org/wfanet/measurement/duchy/daemon/mill/liquidlegionsv2/crypto/LiquidLegionsV2Encryption.kt +++ b/src/main/kotlin/org/wfanet/measurement/duchy/mill/liquidlegionsv2/crypto/LiquidLegionsV2Encryption.kt @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -package org.wfanet.measurement.duchy.daemon.mill.liquidlegionsv2.crypto +package org.wfanet.measurement.duchy.mill.liquidlegionsv2.crypto import org.wfanet.anysketch.crypto.CombineElGamalPublicKeysRequest import org.wfanet.anysketch.crypto.CombineElGamalPublicKeysResponse diff --git a/src/main/kotlin/org/wfanet/measurement/duchy/daemon/mill/liquidlegionsv2/crypto/ReachOnlyLiquidLegionsV2Encryption.kt b/src/main/kotlin/org/wfanet/measurement/duchy/mill/liquidlegionsv2/crypto/ReachOnlyLiquidLegionsV2Encryption.kt similarity index 97% rename from src/main/kotlin/org/wfanet/measurement/duchy/daemon/mill/liquidlegionsv2/crypto/ReachOnlyLiquidLegionsV2Encryption.kt rename to src/main/kotlin/org/wfanet/measurement/duchy/mill/liquidlegionsv2/crypto/ReachOnlyLiquidLegionsV2Encryption.kt index 733cb3cba81..f82fb503060 100644 --- a/src/main/kotlin/org/wfanet/measurement/duchy/daemon/mill/liquidlegionsv2/crypto/ReachOnlyLiquidLegionsV2Encryption.kt +++ b/src/main/kotlin/org/wfanet/measurement/duchy/mill/liquidlegionsv2/crypto/ReachOnlyLiquidLegionsV2Encryption.kt @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -package org.wfanet.measurement.duchy.daemon.mill.liquidlegionsv2.crypto +package org.wfanet.measurement.duchy.mill.liquidlegionsv2.crypto import org.wfanet.anysketch.crypto.CombineElGamalPublicKeysRequest import org.wfanet.anysketch.crypto.CombineElGamalPublicKeysResponse diff --git a/src/main/kotlin/org/wfanet/measurement/duchy/daemon/mill/shareshuffle/BUILD.bazel b/src/main/kotlin/org/wfanet/measurement/duchy/mill/shareshuffle/BUILD.bazel similarity index 87% rename from src/main/kotlin/org/wfanet/measurement/duchy/daemon/mill/shareshuffle/BUILD.bazel rename to src/main/kotlin/org/wfanet/measurement/duchy/mill/shareshuffle/BUILD.bazel index 58e3fe26bfc..4a4988ec71e 100644 --- a/src/main/kotlin/org/wfanet/measurement/duchy/daemon/mill/shareshuffle/BUILD.bazel +++ b/src/main/kotlin/org/wfanet/measurement/duchy/mill/shareshuffle/BUILD.bazel @@ -23,12 +23,12 @@ kt_jvm_library( "//src/main/kotlin/org/wfanet/measurement/api/v2alpha:packed_messages", "//src/main/kotlin/org/wfanet/measurement/common:bitwise_operations", "//src/main/kotlin/org/wfanet/measurement/common/identity", - "//src/main/kotlin/org/wfanet/measurement/duchy/daemon/mill:mill_base", - "//src/main/kotlin/org/wfanet/measurement/duchy/daemon/mill/shareshuffle/crypto:honest_majority_share_shuffle_cryptor", - "//src/main/kotlin/org/wfanet/measurement/duchy/daemon/utils:computation_conversions", - "//src/main/kotlin/org/wfanet/measurement/duchy/daemon/utils:duchy_order", "//src/main/kotlin/org/wfanet/measurement/duchy/db/computation", + "//src/main/kotlin/org/wfanet/measurement/duchy/mill:mill_base", + "//src/main/kotlin/org/wfanet/measurement/duchy/mill/shareshuffle/crypto:honest_majority_share_shuffle_cryptor", "//src/main/kotlin/org/wfanet/measurement/duchy/service/system/v1alpha:advance_computation_request_headers", + "//src/main/kotlin/org/wfanet/measurement/duchy/utils:computation_conversions", + "//src/main/kotlin/org/wfanet/measurement/duchy/utils:duchy_order", "//src/main/kotlin/org/wfanet/measurement/system/v1alpha:resource_key", "//src/main/proto/wfa/any_sketch:frequency_vector_kt_jvm_proto", "//src/main/proto/wfa/measurement/api/v2alpha:certificates_service_kt_jvm_grpc", diff --git a/src/main/kotlin/org/wfanet/measurement/duchy/daemon/mill/shareshuffle/HonestMajorityShareShuffleMill.kt b/src/main/kotlin/org/wfanet/measurement/duchy/mill/shareshuffle/HonestMajorityShareShuffleMill.kt similarity index 97% rename from src/main/kotlin/org/wfanet/measurement/duchy/daemon/mill/shareshuffle/HonestMajorityShareShuffleMill.kt rename to src/main/kotlin/org/wfanet/measurement/duchy/mill/shareshuffle/HonestMajorityShareShuffleMill.kt index d3b29c60d60..266c5bdaf8e 100644 --- a/src/main/kotlin/org/wfanet/measurement/duchy/daemon/mill/shareshuffle/HonestMajorityShareShuffleMill.kt +++ b/src/main/kotlin/org/wfanet/measurement/duchy/mill/shareshuffle/HonestMajorityShareShuffleMill.kt @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -package org.wfanet.measurement.duchy.daemon.mill.shareshuffle +package org.wfanet.measurement.duchy.mill.shareshuffle import com.google.protobuf.ByteString import io.grpc.StatusException @@ -44,19 +44,19 @@ import org.wfanet.measurement.common.xor import org.wfanet.measurement.consent.client.duchy.decryptRandomSeed import org.wfanet.measurement.consent.client.duchy.signEncryptionPublicKey import org.wfanet.measurement.consent.client.duchy.verifyRandomSeed -import org.wfanet.measurement.duchy.daemon.mill.CRYPTO_CPU_DURATION -import org.wfanet.measurement.duchy.daemon.mill.CRYPTO_WALL_CLOCK_DURATION -import org.wfanet.measurement.duchy.daemon.mill.Certificate -import org.wfanet.measurement.duchy.daemon.mill.DATA_TRANSMISSION_RPC_WALL_CLOCK_DURATION -import org.wfanet.measurement.duchy.daemon.mill.MillBase -import org.wfanet.measurement.duchy.daemon.mill.shareshuffle.crypto.HonestMajorityShareShuffleCryptor -import org.wfanet.measurement.duchy.daemon.utils.ReachAndFrequencyResult -import org.wfanet.measurement.duchy.daemon.utils.toV2AlphaEncryptionPublicKey import org.wfanet.measurement.duchy.db.computation.ComputationDataClients import org.wfanet.measurement.duchy.db.computation.ComputationDataClients.PermanentErrorException +import org.wfanet.measurement.duchy.mill.CRYPTO_CPU_DURATION +import org.wfanet.measurement.duchy.mill.CRYPTO_WALL_CLOCK_DURATION +import org.wfanet.measurement.duchy.mill.Certificate +import org.wfanet.measurement.duchy.mill.DATA_TRANSMISSION_RPC_WALL_CLOCK_DURATION +import org.wfanet.measurement.duchy.mill.MillBase +import org.wfanet.measurement.duchy.mill.shareshuffle.crypto.HonestMajorityShareShuffleCryptor import org.wfanet.measurement.duchy.service.internal.computations.inputPathList import org.wfanet.measurement.duchy.service.system.v1alpha.advanceComputationHeader import org.wfanet.measurement.duchy.toProtocolStage +import org.wfanet.measurement.duchy.utils.ReachAndFrequencyResult +import org.wfanet.measurement.duchy.utils.toV2AlphaEncryptionPublicKey import org.wfanet.measurement.internal.duchy.ComputationDetails import org.wfanet.measurement.internal.duchy.ComputationStatsGrpcKt import org.wfanet.measurement.internal.duchy.ComputationToken diff --git a/src/main/kotlin/org/wfanet/measurement/duchy/daemon/mill/shareshuffle/crypto/BUILD.bazel b/src/main/kotlin/org/wfanet/measurement/duchy/mill/shareshuffle/crypto/BUILD.bazel similarity index 100% rename from src/main/kotlin/org/wfanet/measurement/duchy/daemon/mill/shareshuffle/crypto/BUILD.bazel rename to src/main/kotlin/org/wfanet/measurement/duchy/mill/shareshuffle/crypto/BUILD.bazel diff --git a/src/main/kotlin/org/wfanet/measurement/duchy/daemon/mill/shareshuffle/crypto/HonestMajorityShareShuffleCryptor.kt b/src/main/kotlin/org/wfanet/measurement/duchy/mill/shareshuffle/crypto/HonestMajorityShareShuffleCryptor.kt similarity index 96% rename from src/main/kotlin/org/wfanet/measurement/duchy/daemon/mill/shareshuffle/crypto/HonestMajorityShareShuffleCryptor.kt rename to src/main/kotlin/org/wfanet/measurement/duchy/mill/shareshuffle/crypto/HonestMajorityShareShuffleCryptor.kt index efd4c933561..440fb2e16f8 100644 --- a/src/main/kotlin/org/wfanet/measurement/duchy/daemon/mill/shareshuffle/crypto/HonestMajorityShareShuffleCryptor.kt +++ b/src/main/kotlin/org/wfanet/measurement/duchy/mill/shareshuffle/crypto/HonestMajorityShareShuffleCryptor.kt @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -package org.wfanet.measurement.duchy.daemon.mill.shareshuffle.crypto +package org.wfanet.measurement.duchy.mill.shareshuffle.crypto import org.wfanet.measurement.internal.duchy.protocol.CompleteAggregationPhaseRequest import org.wfanet.measurement.internal.duchy.protocol.CompleteAggregationPhaseResponse diff --git a/src/main/kotlin/org/wfanet/measurement/duchy/daemon/mill/shareshuffle/crypto/JniHonestMajorityShareShuffleCryptor.kt b/src/main/kotlin/org/wfanet/measurement/duchy/mill/shareshuffle/crypto/JniHonestMajorityShareShuffleCryptor.kt similarity index 97% rename from src/main/kotlin/org/wfanet/measurement/duchy/daemon/mill/shareshuffle/crypto/JniHonestMajorityShareShuffleCryptor.kt rename to src/main/kotlin/org/wfanet/measurement/duchy/mill/shareshuffle/crypto/JniHonestMajorityShareShuffleCryptor.kt index 9e30e0d7083..c22ac8ffc02 100644 --- a/src/main/kotlin/org/wfanet/measurement/duchy/daemon/mill/shareshuffle/crypto/JniHonestMajorityShareShuffleCryptor.kt +++ b/src/main/kotlin/org/wfanet/measurement/duchy/mill/shareshuffle/crypto/JniHonestMajorityShareShuffleCryptor.kt @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -package org.wfanet.measurement.duchy.daemon.mill.shareshuffle.crypto +package org.wfanet.measurement.duchy.mill.shareshuffle.crypto import org.wfanet.measurement.internal.duchy.protocol.CompleteAggregationPhaseRequest import org.wfanet.measurement.internal.duchy.protocol.CompleteAggregationPhaseResponse diff --git a/src/main/kotlin/org/wfanet/measurement/duchy/service/system/v1alpha/BUILD.bazel b/src/main/kotlin/org/wfanet/measurement/duchy/service/system/v1alpha/BUILD.bazel index 82c4e333b2d..e0780eaf8b4 100644 --- a/src/main/kotlin/org/wfanet/measurement/duchy/service/system/v1alpha/BUILD.bazel +++ b/src/main/kotlin/org/wfanet/measurement/duchy/service/system/v1alpha/BUILD.bazel @@ -14,7 +14,7 @@ kt_jvm_library( srcs = ["AdvanceComputationRequestHeaders.kt"], visibility = [ ":__subpackages__", - "//src/main/kotlin/org/wfanet/measurement/duchy/daemon/mill:__subpackages__", + "//src/main/kotlin/org/wfanet/measurement/duchy/mill:__subpackages__", ], deps = [ "//src/main/kotlin/org/wfanet/measurement/duchy:computation_stage", diff --git a/src/main/kotlin/org/wfanet/measurement/duchy/daemon/testing/BUILD.bazel b/src/main/kotlin/org/wfanet/measurement/duchy/testing/BUILD.bazel similarity index 92% rename from src/main/kotlin/org/wfanet/measurement/duchy/daemon/testing/BUILD.bazel rename to src/main/kotlin/org/wfanet/measurement/duchy/testing/BUILD.bazel index f7778e7f88b..28a0803195a 100644 --- a/src/main/kotlin/org/wfanet/measurement/duchy/daemon/testing/BUILD.bazel +++ b/src/main/kotlin/org/wfanet/measurement/duchy/testing/BUILD.bazel @@ -3,7 +3,7 @@ load("@wfa_rules_kotlin_jvm//kotlin:defs.bzl", "kt_jvm_library") package( default_testonly = True, default_visibility = [ - "//src/test/kotlin/org/wfanet/measurement/duchy/daemon:__subpackages__", + "//src/test/kotlin/org/wfanet/measurement/duchy:__subpackages__", ], ) diff --git a/src/main/kotlin/org/wfanet/measurement/duchy/daemon/testing/TestRequisition.kt b/src/main/kotlin/org/wfanet/measurement/duchy/testing/TestRequisition.kt similarity index 98% rename from src/main/kotlin/org/wfanet/measurement/duchy/daemon/testing/TestRequisition.kt rename to src/main/kotlin/org/wfanet/measurement/duchy/testing/TestRequisition.kt index 8800c832342..c1d3245bf35 100644 --- a/src/main/kotlin/org/wfanet/measurement/duchy/daemon/testing/TestRequisition.kt +++ b/src/main/kotlin/org/wfanet/measurement/duchy/testing/TestRequisition.kt @@ -14,7 +14,7 @@ * limitations under the License. */ -package org.wfanet.measurement.duchy.daemon.testing +package org.wfanet.measurement.duchy.testing import com.google.protobuf.ByteString import com.google.protobuf.kotlin.toByteString diff --git a/src/main/kotlin/org/wfanet/measurement/duchy/daemon/utils/BUILD.bazel b/src/main/kotlin/org/wfanet/measurement/duchy/utils/BUILD.bazel similarity index 92% rename from src/main/kotlin/org/wfanet/measurement/duchy/daemon/utils/BUILD.bazel rename to src/main/kotlin/org/wfanet/measurement/duchy/utils/BUILD.bazel index c0223c49e97..a1787d3ec05 100644 --- a/src/main/kotlin/org/wfanet/measurement/duchy/daemon/utils/BUILD.bazel +++ b/src/main/kotlin/org/wfanet/measurement/duchy/utils/BUILD.bazel @@ -1,8 +1,8 @@ load("@wfa_rules_kotlin_jvm//kotlin:defs.bzl", "kt_jvm_library") package(default_visibility = [ - "//src/main/kotlin/org/wfanet/measurement/duchy/daemon:__subpackages__", - "//src/test/kotlin/org/wfanet/measurement/duchy/daemon:__subpackages__", + "//src/main/kotlin/org/wfanet/measurement/duchy:__subpackages__", + "//src/test/kotlin/org/wfanet/measurement/duchy:__subpackages__", ]) kt_jvm_library( diff --git a/src/main/kotlin/org/wfanet/measurement/duchy/daemon/utils/ComputationConversions.kt b/src/main/kotlin/org/wfanet/measurement/duchy/utils/ComputationConversions.kt similarity index 99% rename from src/main/kotlin/org/wfanet/measurement/duchy/daemon/utils/ComputationConversions.kt rename to src/main/kotlin/org/wfanet/measurement/duchy/utils/ComputationConversions.kt index 860b4a900bb..a9d8915f72a 100644 --- a/src/main/kotlin/org/wfanet/measurement/duchy/daemon/utils/ComputationConversions.kt +++ b/src/main/kotlin/org/wfanet/measurement/duchy/utils/ComputationConversions.kt @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -package org.wfanet.measurement.duchy.daemon.utils +package org.wfanet.measurement.duchy.utils import com.google.protobuf.Any as ProtoAny import com.google.protobuf.ByteString diff --git a/src/main/kotlin/org/wfanet/measurement/duchy/daemon/utils/DuchyOrder.kt b/src/main/kotlin/org/wfanet/measurement/duchy/utils/DuchyOrder.kt similarity index 98% rename from src/main/kotlin/org/wfanet/measurement/duchy/daemon/utils/DuchyOrder.kt rename to src/main/kotlin/org/wfanet/measurement/duchy/utils/DuchyOrder.kt index 1ab807daa92..b76ad01af61 100644 --- a/src/main/kotlin/org/wfanet/measurement/duchy/daemon/utils/DuchyOrder.kt +++ b/src/main/kotlin/org/wfanet/measurement/duchy/utils/DuchyOrder.kt @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -package org.wfanet.measurement.duchy.daemon.utils +package org.wfanet.measurement.duchy.utils import java.math.BigInteger import java.security.MessageDigest diff --git a/src/main/kotlin/org/wfanet/measurement/integration/common/BUILD.bazel b/src/main/kotlin/org/wfanet/measurement/integration/common/BUILD.bazel index efbb302bd24..8a559e4e0ff 100644 --- a/src/main/kotlin/org/wfanet/measurement/integration/common/BUILD.bazel +++ b/src/main/kotlin/org/wfanet/measurement/integration/common/BUILD.bazel @@ -148,10 +148,10 @@ kt_jvm_library( "//src/main/kotlin/org/wfanet/measurement/common/identity", "//src/main/kotlin/org/wfanet/measurement/common/identity:principal_identity", "//src/main/kotlin/org/wfanet/measurement/common/identity/testing", - "//src/main/kotlin/org/wfanet/measurement/duchy/daemon/herald", - "//src/main/kotlin/org/wfanet/measurement/duchy/daemon/mill/liquidlegionsv2:reach_frequency_liquid_legions_v2_mill", - "//src/main/kotlin/org/wfanet/measurement/duchy/daemon/mill/liquidlegionsv2:reach_only_liquid_legions_v2_mill", "//src/main/kotlin/org/wfanet/measurement/duchy/deploy/common/server:duchy_data_server", + "//src/main/kotlin/org/wfanet/measurement/duchy/herald", + "//src/main/kotlin/org/wfanet/measurement/duchy/mill/liquidlegionsv2:reach_frequency_liquid_legions_v2_mill", + "//src/main/kotlin/org/wfanet/measurement/duchy/mill/liquidlegionsv2:reach_only_liquid_legions_v2_mill", "//src/main/kotlin/org/wfanet/measurement/duchy/storage:tink_key_store", "@wfa_common_jvm//imports/java/io/grpc:api", "@wfa_common_jvm//imports/java/io/grpc:core", diff --git a/src/main/kotlin/org/wfanet/measurement/integration/common/InProcessDuchy.kt b/src/main/kotlin/org/wfanet/measurement/integration/common/InProcessDuchy.kt index adef4caf53c..2da4b8dcec7 100644 --- a/src/main/kotlin/org/wfanet/measurement/integration/common/InProcessDuchy.kt +++ b/src/main/kotlin/org/wfanet/measurement/integration/common/InProcessDuchy.kt @@ -51,15 +51,15 @@ import org.wfanet.measurement.common.identity.withDuchyId import org.wfanet.measurement.common.testing.ProviderRule import org.wfanet.measurement.common.testing.chainRulesSequentially import org.wfanet.measurement.common.throttler.MinimumIntervalThrottler -import org.wfanet.measurement.duchy.daemon.herald.ContinuationTokenManager -import org.wfanet.measurement.duchy.daemon.herald.Herald -import org.wfanet.measurement.duchy.daemon.mill.Certificate -import org.wfanet.measurement.duchy.daemon.mill.liquidlegionsv2.ReachFrequencyLiquidLegionsV2Mill -import org.wfanet.measurement.duchy.daemon.mill.liquidlegionsv2.ReachOnlyLiquidLegionsV2Mill -import org.wfanet.measurement.duchy.daemon.mill.liquidlegionsv2.crypto.JniLiquidLegionsV2Encryption -import org.wfanet.measurement.duchy.daemon.mill.liquidlegionsv2.crypto.JniReachOnlyLiquidLegionsV2Encryption import org.wfanet.measurement.duchy.db.computation.ComputationDataClients import org.wfanet.measurement.duchy.deploy.common.service.DuchyDataServices +import org.wfanet.measurement.duchy.herald.ContinuationTokenManager +import org.wfanet.measurement.duchy.herald.Herald +import org.wfanet.measurement.duchy.mill.Certificate +import org.wfanet.measurement.duchy.mill.liquidlegionsv2.ReachFrequencyLiquidLegionsV2Mill +import org.wfanet.measurement.duchy.mill.liquidlegionsv2.ReachOnlyLiquidLegionsV2Mill +import org.wfanet.measurement.duchy.mill.liquidlegionsv2.crypto.JniLiquidLegionsV2Encryption +import org.wfanet.measurement.duchy.mill.liquidlegionsv2.crypto.JniReachOnlyLiquidLegionsV2Encryption import org.wfanet.measurement.duchy.service.api.v2alpha.RequisitionFulfillmentService import org.wfanet.measurement.duchy.service.internal.computationcontrol.AsyncComputationControlService import org.wfanet.measurement.duchy.service.system.v1alpha.ComputationControlService @@ -273,8 +273,8 @@ class InProcessDuchy( val throttler = MinimumIntervalThrottler(Clock.systemUTC(), Duration.ofSeconds(1)) throttler.loopOnReady { - reachFrequencyLiquidLegionsV2Mill.pollAndProcessNextComputation() - reachOnlyLiquidLegionsV2Mill.pollAndProcessNextComputation() + reachFrequencyLiquidLegionsV2Mill.claimAndProcessWork() + reachOnlyLiquidLegionsV2Mill.claimAndProcessWork() } } } diff --git a/src/main/terraform/aws/cmms/duchies.tf b/src/main/terraform/aws/cmms/duchies.tf index be924122cd0..21dd4f96015 100644 --- a/src/main/terraform/aws/cmms/duchies.tf +++ b/src/main/terraform/aws/cmms/duchies.tf @@ -30,7 +30,7 @@ module "clusters" { default_instance_types = ["m5.large"] default_max_node_count = 2 high_perf_instance_types = ["c5.xlarge"] - high_perf_max_node_count = 3 + high_perf_max_node_count = 20 } # IAM role to be used by aws_load_balancer_controller addons across all clusters diff --git a/src/main/terraform/aws/examples/duchy/main.tf b/src/main/terraform/aws/examples/duchy/main.tf index 60443dce271..e082e3e9b11 100644 --- a/src/main/terraform/aws/examples/duchy/main.tf +++ b/src/main/terraform/aws/examples/duchy/main.tf @@ -87,7 +87,7 @@ module "cluster" { default_instance_types = ["m5.large"] default_max_node_count = 2 high_perf_instance_types = ["c5.xlarge"] - high_perf_max_node_count = 2 + high_perf_max_node_count = 20 } module "load_balancer_controller_irsa_role" { diff --git a/src/main/terraform/gcloud/cmms/.terraform.lock.hcl b/src/main/terraform/gcloud/cmms/.terraform.lock.hcl index b68810a5e8c..4357aa33bd1 100644 --- a/src/main/terraform/gcloud/cmms/.terraform.lock.hcl +++ b/src/main/terraform/gcloud/cmms/.terraform.lock.hcl @@ -24,21 +24,21 @@ provider "registry.terraform.io/cyrilgdn/postgresql" { } provider "registry.terraform.io/hashicorp/google" { - version = "5.4.0" - constraints = ">= 4.63.0, ~> 5.4.0" + version = "5.37.0" + constraints = ">= 4.63.0, >= 5.4.0, >= 5.37.0, ~> 5.37.0" hashes = [ - "h1:9YorEXI71ZzyY2LuR1+8m0kYyr9iQjcU0lZQRnxh5iA=", - "zh:045586ab10eab1ab6fc85b7ac238533c65a2517ad5592464b28ad83a9c1bb384", - "zh:126087ff0ffcd2b917cf9379b73197a6a62832ff054adc930391ec9162542aeb", - "zh:2280b05503e83a447b5d8a01394c7941cbad49809f1f7579401ce59707938127", - "zh:22be8d8c4dce29434295f83ee09b9bcdcb05ac5d8c19952884e09cfcf667495e", - "zh:245e32e49d457e35ee2a7dd14228cec768c9ab7fc023cdc820d018638ad84102", - "zh:4caaa9e7cdf820825841dc11a8006b5eeb2825e10c5c5e0dc8226599132efcb6", - "zh:4e7de78b5cc7b71849f9f20a305f145c1915abc260be23814bb27bcb3dc684e6", - "zh:713ceb0b26598712b4e8f8a6b23df5564798265a1afbc71d177329d0a39f45bd", - "zh:966a6091a7b2237e82b530a89afefeb639fdfc59f1e3780dae52e90491d1a27d", - "zh:dd4c936f9da7bb6b1c1b442276afe7933f41a3dbba763a58773006bae8416e80", + "h1:+0FfxQm8+OxXM6tDF8UCA2ir59Gel4oSnLoJoKkGCpc=", + "zh:00754c426ec8c2416d54904b25ebf3d831b5af4f02a20336d0a60d91897497d7", + "zh:0b4128affc12ace78d9cbd8aac308b992cc1829effba0ac761b4c3e9e37a36e2", + "zh:46afda4cdadc242b0a27acb181967b7a37768aa084c84593506490483fbdefd6", + "zh:4ad56e9d5ed5e05184bd0a23e17b6eb985c28e7ce23ce29fdf50daa5594ead69", + "zh:557b021f77bd97462b20519148f454158e6cf89eba5461ae6016568b317e1c87", + "zh:6552bb5e901e9fbcd95dc95b4e017dfd7c1eca465dd749f9c8afdc4ae2bf3213", + "zh:687347e31e69f2c4518abc057ad157ed4fd2b78e399b3a172bcab4277f7a235b", + "zh:a66bfc55856693fe82a81554abf7fd72b8ca2d56a08cb59c4769c15b1a1acea5", + "zh:a8b242c5aab000f2a27e934930a75656efb4a96fdb06a419b22ae0daffa6fba3", + "zh:da9e9b40d632f218a3e0bb88b8cf95b91485cee1eb2fd2a384d45c2619c36da4", "zh:f569b65999264a9416862bca5cd2a6177d94ccb0424f3a4ef424428912b9cb3c", - "zh:f9162d90318bdc237f2b4e3462143f60d752ab69ed08d50467f6a4025a504df0", + "zh:fda7bd1578e4b40e9e2d0298c77b0ad9f4486abee8ad38755299dd4b06be54c7", ] } diff --git a/src/main/terraform/gcloud/cmms/duchies.tf b/src/main/terraform/gcloud/cmms/duchies.tf index 1080cd40505..7e52433cf60 100644 --- a/src/main/terraform/gcloud/cmms/duchies.tf +++ b/src/main/terraform/gcloud/cmms/duchies.tf @@ -16,10 +16,11 @@ module "clusters" { source = "../modules/cluster" for_each = local.duchy_names - name = "${each.key}-duchy" - location = local.cluster_location - release_channel = var.cluster_release_channel - secret_key = module.common.cluster_secret_key + name = "${each.key}-duchy" + location = local.cluster_location + release_channel = var.cluster_release_channel + secret_key = module.common.cluster_secret_key + autoscaling_profile = "OPTIMIZE_UTILIZATION" } module "default_node_pools" { @@ -41,7 +42,7 @@ module "highmem_node_pools" { name = "highmem" service_account = module.common.cluster_service_account machine_type = "c2-standard-4" - max_node_count = 3 + max_node_count = 20 spot = true } diff --git a/src/main/terraform/gcloud/cmms/versions.tf b/src/main/terraform/gcloud/cmms/versions.tf index 653db3aba6a..cfd1dd1294a 100644 --- a/src/main/terraform/gcloud/cmms/versions.tf +++ b/src/main/terraform/gcloud/cmms/versions.tf @@ -16,7 +16,7 @@ terraform { required_providers { google = { source = "hashicorp/google" - version = "~> 5.4.0" + version = "~> 5.37.0" } postgresql = { source = "cyrilgdn/postgresql" diff --git a/src/main/terraform/gcloud/examples/duchy/main.tf b/src/main/terraform/gcloud/examples/duchy/main.tf index 3ca228f642c..eaa1cf70bb6 100644 --- a/src/main/terraform/gcloud/examples/duchy/main.tf +++ b/src/main/terraform/gcloud/examples/duchy/main.tf @@ -47,10 +47,11 @@ module "storage" { module "cluster" { source = "../../modules/cluster" - name = local.cluster_name - location = local.cluster_location - release_channel = var.cluster_release_channel - secret_key = module.common.cluster_secret_key + name = local.cluster_name + location = local.cluster_location + release_channel = var.cluster_release_channel + secret_key = module.common.cluster_secret_key + autoscaling_profile = "OPTIMIZE_UTILIZATION" } module "default_node_pool" { @@ -70,7 +71,7 @@ module "spot_node_pool" { name = "spot" service_account = module.common.cluster_service_account machine_type = "c2-standard-4" - max_node_count = 2 + max_node_count = 20 spot = true } diff --git a/src/main/terraform/gcloud/modules/cluster/main.tf b/src/main/terraform/gcloud/modules/cluster/main.tf index 9f2fafcdf6f..9faedd43a8e 100644 --- a/src/main/terraform/gcloud/modules/cluster/main.tf +++ b/src/main/terraform/gcloud/modules/cluster/main.tf @@ -48,6 +48,10 @@ resource "google_container_cluster" "cluster" { channel = var.release_channel } + cluster_autoscaling { + autoscaling_profile = var.autoscaling_profile + } + # We can't create a cluster with no node pool defined, but we want to only use # separately managed node pools. So we create the smallest possible default # node pool and immediately delete it. diff --git a/src/main/terraform/gcloud/modules/cluster/variables.tf b/src/main/terraform/gcloud/modules/cluster/variables.tf index 720c7c200f4..6a5b0416010 100644 --- a/src/main/terraform/gcloud/modules/cluster/variables.tf +++ b/src/main/terraform/gcloud/modules/cluster/variables.tf @@ -36,3 +36,9 @@ variable "release_channel" { type = string nullable = false } + +variable "autoscaling_profile" { + description = "Autoscaling profile" + type = string + default = null +} diff --git a/src/main/terraform/gcloud/modules/cluster/versions.tf b/src/main/terraform/gcloud/modules/cluster/versions.tf index f7b31289a0c..1d73575d325 100644 --- a/src/main/terraform/gcloud/modules/cluster/versions.tf +++ b/src/main/terraform/gcloud/modules/cluster/versions.tf @@ -16,7 +16,7 @@ terraform { required_providers { google = { source = "hashicorp/google" - version = ">= 4.63.0" + version = ">= 5.37.0" } } } diff --git a/src/test/kotlin/org/wfanet/measurement/duchy/deploy/common/daemon/mill/BUILD.bazel b/src/test/kotlin/org/wfanet/measurement/duchy/deploy/common/daemon/mill/BUILD.bazel new file mode 100644 index 00000000000..7282b4a406e --- /dev/null +++ b/src/test/kotlin/org/wfanet/measurement/duchy/deploy/common/daemon/mill/BUILD.bazel @@ -0,0 +1,16 @@ +load("@wfa_rules_kotlin_jvm//kotlin:defs.bzl", "kt_jvm_test") + +kt_jvm_test( + name = "MillJobSchedulerTest", + srcs = ["MillJobSchedulerTest.kt"], + test_class = "org.wfanet.measurement.duchy.deploy.common.daemon.mill.MillJobSchedulerTest", + deps = [ + "//src/main/kotlin/org/wfanet/measurement/duchy/deploy/common/daemon/mill:job_scheduler", + "//src/main/proto/wfa/measurement/internal/duchy:computations_service_kt_jvm_grpc_proto", + "@wfa_common_jvm//imports/java/com/google/common/truth", + "@wfa_common_jvm//imports/java/org/junit", + "@wfa_common_jvm//imports/kotlin/kotlinx/coroutines/test", + "@wfa_common_jvm//src/main/kotlin/org/wfanet/measurement/common/grpc/testing", + "@wfa_common_jvm//src/main/kotlin/org/wfanet/measurement/common/testing", + ], +) diff --git a/src/test/kotlin/org/wfanet/measurement/duchy/deploy/common/daemon/mill/MillJobSchedulerTest.kt b/src/test/kotlin/org/wfanet/measurement/duchy/deploy/common/daemon/mill/MillJobSchedulerTest.kt new file mode 100644 index 00000000000..5489d19ae9c --- /dev/null +++ b/src/test/kotlin/org/wfanet/measurement/duchy/deploy/common/daemon/mill/MillJobSchedulerTest.kt @@ -0,0 +1,281 @@ +/* + * Copyright 2024 The Cross-Media Measurement Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.wfanet.measurement.duchy.deploy.common.daemon.mill + +import com.google.common.truth.Truth.assertThat +import com.google.common.truth.extensions.proto.ProtoTruth.assertThat +import io.kubernetes.client.openapi.JSON +import io.kubernetes.client.openapi.models.V1Container +import io.kubernetes.client.openapi.models.V1Deployment +import io.kubernetes.client.openapi.models.V1Job +import io.kubernetes.client.openapi.models.V1JobCondition +import io.kubernetes.client.openapi.models.V1JobList +import io.kubernetes.client.openapi.models.V1JobStatus +import io.kubernetes.client.openapi.models.V1ObjectMeta +import io.kubernetes.client.openapi.models.V1OwnerReference +import io.kubernetes.client.openapi.models.V1PodSpec +import io.kubernetes.client.openapi.models.V1PodTemplate +import io.kubernetes.client.openapi.models.V1PodTemplateSpec +import io.kubernetes.client.util.Namespaces +import java.time.Duration +import java.time.Instant +import java.time.ZoneOffset +import kotlin.random.Random +import kotlinx.coroutines.CompletableDeferred +import kotlinx.coroutines.launch +import kotlinx.coroutines.test.runTest +import org.junit.Rule +import org.junit.Test +import org.junit.runner.RunWith +import org.junit.runners.JUnit4 +import org.mockito.ArgumentMatcher +import org.mockito.kotlin.any +import org.mockito.kotlin.argThat +import org.mockito.kotlin.argumentCaptor +import org.mockito.kotlin.doReturn +import org.mockito.kotlin.eq +import org.mockito.kotlin.mock +import org.mockito.kotlin.never +import org.mockito.kotlin.times +import org.mockito.kotlin.verify +import org.mockito.kotlin.whenever +import org.wfanet.measurement.common.CountDownLatch +import org.wfanet.measurement.common.grpc.testing.GrpcTestServerRule +import org.wfanet.measurement.common.grpc.testing.mockService +import org.wfanet.measurement.common.k8s.KubernetesClient +import org.wfanet.measurement.common.testing.verifyAndCapture +import org.wfanet.measurement.common.toProtoDuration +import org.wfanet.measurement.duchy.toProtocolStage +import org.wfanet.measurement.internal.duchy.ClaimWorkRequest +import org.wfanet.measurement.internal.duchy.ComputationTypeEnum.ComputationType +import org.wfanet.measurement.internal.duchy.ComputationsGrpcKt +import org.wfanet.measurement.internal.duchy.claimWorkRequest +import org.wfanet.measurement.internal.duchy.claimWorkResponse +import org.wfanet.measurement.internal.duchy.computationToken +import org.wfanet.measurement.internal.duchy.protocol.LiquidLegionsSketchAggregationV2 + +@RunWith(JUnit4::class) +class MillJobSchedulerTest { + private val computationsServiceMock: ComputationsGrpcKt.ComputationsCoroutineImplBase = + mockService() + private val k8sClientMock: KubernetesClient = mock { + onBlocking { getDeployment(any(), any()) } doReturn deployment + onBlocking { getPodTemplate(eq(LLV2_POD_TEMPLATE_NAME), any()) } doReturn + liquidLegionsV2PodTemplate + } + + @JvmField @Rule val serverRule = GrpcTestServerRule { addService(computationsServiceMock) } + + private fun createMillJobScheduler(): MillJobScheduler { + return MillJobScheduler( + DUCHY_ID, + ComputationsGrpcKt.ComputationsCoroutineStub(serverRule.channel), + DEPLOYMENT_NAME, + POLLING_DELAY, + SUCCESSFUL_JOB_HISTORY_LIMIT, + FAILED_JOB_HISTORY_LIMIT, + LLV2_POD_TEMPLATE_NAME, + LLV2_MAX_CONCURRENCY, + LLV2_WORK_LOCK_DURATION, + Random(123), + k8sClientMock, + ) + } + + @Test + fun `claims work and creates job when below maximum concurrency`() = runTest { + whenever(k8sClientMock.listJobs(any(), any())).thenReturn(V1JobList()) + whenever( + computationsServiceMock.claimWork( + hasComputationType(ComputationType.LIQUID_LEGIONS_SKETCH_AGGREGATION_V2) + ) + ) + .thenReturn(claimWorkResponse { token = computationToken { globalComputationId = "comp-1" } }) + val jobCreated = CompletableDeferred() + whenever(k8sClientMock.createJob(any())).thenAnswer { invocation -> + jobCreated.complete(invocation.arguments.first() as V1Job) + null + } + val jobScheduler = createMillJobScheduler() + + backgroundScope.launch { jobScheduler.run() } + // Wait until createJob has been called. + val createJobRequest: V1Job = jobCreated.await() + + val claimWorkRequest: ClaimWorkRequest = + verifyAndCapture( + computationsServiceMock, + ComputationsGrpcKt.ComputationsCoroutineImplBase::claimWork, + ) + assertThat(claimWorkRequest) + .ignoringFields(ClaimWorkRequest.OWNER_FIELD_NUMBER) + .isEqualTo( + claimWorkRequest { + computationType = ComputationType.LIQUID_LEGIONS_SKETCH_AGGREGATION_V2 + prioritizedStages += + LiquidLegionsSketchAggregationV2.Stage.INITIALIZATION_PHASE.toProtocolStage() + lockDuration = LLV2_WORK_LOCK_DURATION.toProtoDuration() + } + ) + assertThat(createJobRequest.metadata.labels).containsEntry("mill-type", "LIQUID_LEGIONS_V2") + assertThat(createJobRequest.metadata.ownerReferences).containsExactly(deploymentOwnerReference) + assertThat(createJobRequest.spec.template.metadata.labels) + .containsAtLeastEntriesIn(liquidLegionsV2PodTemplate.template.metadata.labels) + } + + @Test + fun `does not claim work when at maximum concurrency`() = runTest { + val listJobsLatch = CountDownLatch(2) + whenever(k8sClientMock.listJobs(any(), any())).thenAnswer { + listJobsLatch.countDown() + V1JobList().apply { + repeat(LLV2_MAX_CONCURRENCY) { addItemsItem(buildOwnedJob { status.active = 1 }) } + } + } + val jobScheduler = createMillJobScheduler() + + backgroundScope.launch { jobScheduler.run() } + // Wait until claimWork has had a chance to be called. + listJobsLatch.await() + + verify(computationsServiceMock, never()).claimWork(any()) + } + + @Test + fun `cleans up jobs above history limits`() = runTest { + val listJobsLatch = CountDownLatch(2) + val newestCompletionTime = Instant.now() + whenever(k8sClientMock.listJobs(any(), any())).thenAnswer { + listJobsLatch.countDown() + V1JobList().apply { + repeat(SUCCESSFUL_JOB_HISTORY_LIMIT + 1) { index -> + addItemsItem( + buildOwnedJob { + metadata.name = "successful-job-$index" + status.addConditionsItem( + V1JobCondition().apply { + type = "Complete" + status = "True" + } + ) + status.completionTime = + newestCompletionTime.minusSeconds(index.toLong()).atOffset(ZoneOffset.UTC) + } + ) + } + repeat(FAILED_JOB_HISTORY_LIMIT + 1) { index -> + addItemsItem( + buildOwnedJob { + metadata.name = "failed-job-$index" + status.addConditionsItem( + V1JobCondition().apply { + type = "Complete" + status = "True" + } + ) + status.addConditionsItem( + V1JobCondition().apply { + type = "Failed" + status = "True" + } + ) + status.completionTime = + newestCompletionTime.minusSeconds(index.toLong()).atOffset(ZoneOffset.UTC) + } + ) + } + } + } + val jobScheduler = createMillJobScheduler() + + backgroundScope.launch { jobScheduler.run() } + // Wait until job cleanup has had a chance to be called. + listJobsLatch.await() + + val jobNameCaptor = argumentCaptor() + verify(k8sClientMock, times(2)).deleteJob(jobNameCaptor.capture(), any(), any()) + assertThat(jobNameCaptor.allValues).containsExactly("successful-job-3", "failed-job-3") + } + + private class ComputationTypeMatcher(private val computationType: ComputationType) : + ArgumentMatcher { + override fun matches(argument: ClaimWorkRequest): Boolean { + return argument.computationType == computationType + } + } + + private fun hasComputationType(computationType: ComputationType) = + argThat(ComputationTypeMatcher(computationType)) + + companion object { + init { + // Ensure JSON for Kubernetes API is initialized. + JSON() + } + + private const val DUCHY_ID = "worker1" + private const val DEPLOYMENT_NAME = "$DUCHY_ID-mill-job-scheduler-deployment" + private const val LLV2_POD_TEMPLATE_NAME = "$DUCHY_ID-llv2-mill" + private val LLV2_WORK_LOCK_DURATION = Duration.ofMinutes(5) + private const val LLV2_MAX_CONCURRENCY = 5 + private val POLLING_DELAY = Duration.ofSeconds(2) + private const val SUCCESSFUL_JOB_HISTORY_LIMIT = 3 + private const val FAILED_JOB_HISTORY_LIMIT = 3 + + private val deployment = + V1Deployment().apply { + apiVersion = "apps/v1" + kind = "Deployment" + metadata = + V1ObjectMeta().apply { + name = DEPLOYMENT_NAME + uid = "$DEPLOYMENT_NAME-uid" + } + } + private val deploymentOwnerReference = + V1OwnerReference().apply { + apiVersion = deployment.apiVersion + kind = deployment.kind + name = deployment.metadata.name + uid = deployment.metadata.uid + } + private val liquidLegionsV2PodTemplate = + V1PodTemplate().apply { + metadata = V1ObjectMeta().apply { name = LLV2_POD_TEMPLATE_NAME } + template = + V1PodTemplateSpec().apply { + metadata = V1ObjectMeta().apply { putLabelsItem("app", "llv2-mill") } + spec = + V1PodSpec().apply { + addContainersItem(V1Container().apply { name = "llv2-mill-container" }) + } + } + } + + private fun buildOwnedJob(fill: V1Job.() -> Unit): V1Job { + return V1Job().apply { + metadata = + V1ObjectMeta().apply { + namespace = Namespaces.NAMESPACE_DEFAULT + addOwnerReferencesItem(V1OwnerReference().apply { uid = deployment.metadata.uid }) + } + status = V1JobStatus() + fill() + } + } + } +} diff --git a/src/test/kotlin/org/wfanet/measurement/duchy/daemon/herald/BUILD.bazel b/src/test/kotlin/org/wfanet/measurement/duchy/herald/BUILD.bazel similarity index 86% rename from src/test/kotlin/org/wfanet/measurement/duchy/daemon/herald/BUILD.bazel rename to src/test/kotlin/org/wfanet/measurement/duchy/herald/BUILD.bazel index b2db80f4c40..a3a8e700668 100644 --- a/src/test/kotlin/org/wfanet/measurement/duchy/daemon/herald/BUILD.bazel +++ b/src/test/kotlin/org/wfanet/measurement/duchy/herald/BUILD.bazel @@ -3,17 +3,17 @@ load("@wfa_rules_kotlin_jvm//kotlin:defs.bzl", "kt_jvm_test") kt_jvm_test( name = "HeraldTest", srcs = ["HeraldTest.kt"], - test_class = "org.wfanet.measurement.duchy.daemon.herald.HeraldTest", + test_class = "org.wfanet.measurement.duchy.herald.HeraldTest", deps = [ "//src/main/kotlin/org/wfanet/measurement/api/v2alpha:resource_key", "//src/main/kotlin/org/wfanet/measurement/duchy:computation_stage", - "//src/main/kotlin/org/wfanet/measurement/duchy/daemon/herald", - "//src/main/kotlin/org/wfanet/measurement/duchy/daemon/testing", "//src/main/kotlin/org/wfanet/measurement/duchy/db/computation", "//src/main/kotlin/org/wfanet/measurement/duchy/db/computation/testing", + "//src/main/kotlin/org/wfanet/measurement/duchy/herald", "//src/main/kotlin/org/wfanet/measurement/duchy/service/internal/computations", "//src/main/kotlin/org/wfanet/measurement/duchy/service/internal/testing", "//src/main/kotlin/org/wfanet/measurement/duchy/storage:tink_key_store", + "//src/main/kotlin/org/wfanet/measurement/duchy/testing", "//src/main/kotlin/org/wfanet/measurement/system/v1alpha:resource_key", "//src/main/proto/wfa/measurement/internal/duchy:noise_mechanism_kt_jvm_proto", "//src/main/proto/wfa/measurement/internal/duchy:requisition_details_kt_jvm_proto", @@ -38,9 +38,9 @@ kt_jvm_test( name = "ContinuationTokenManagerTest", testonly = True, srcs = ["ContinuationTokenManagerTest.kt"], - test_class = "org.wfanet.measurement.duchy.daemon.herald.ContinuationTokenManagerTest", + test_class = "org.wfanet.measurement.duchy.herald.ContinuationTokenManagerTest", deps = [ - "//src/main/kotlin/org/wfanet/measurement/duchy/daemon/herald:continuation_token_manager", + "//src/main/kotlin/org/wfanet/measurement/duchy/herald:continuation_token_manager", "//src/main/kotlin/org/wfanet/measurement/duchy/service/internal/testing", "@wfa_common_jvm//imports/java/com/google/common/truth", "@wfa_common_jvm//imports/java/org/junit", diff --git a/src/test/kotlin/org/wfanet/measurement/duchy/daemon/herald/ContinuationTokenManagerTest.kt b/src/test/kotlin/org/wfanet/measurement/duchy/herald/ContinuationTokenManagerTest.kt similarity index 98% rename from src/test/kotlin/org/wfanet/measurement/duchy/daemon/herald/ContinuationTokenManagerTest.kt rename to src/test/kotlin/org/wfanet/measurement/duchy/herald/ContinuationTokenManagerTest.kt index 0eb755ed7e7..0fd57958526 100644 --- a/src/test/kotlin/org/wfanet/measurement/duchy/daemon/herald/ContinuationTokenManagerTest.kt +++ b/src/test/kotlin/org/wfanet/measurement/duchy/herald/ContinuationTokenManagerTest.kt @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -package org.wfanet.measurement.duchy.daemon.herald +package org.wfanet.measurement.duchy.herald import com.google.common.truth.Truth.assertThat import io.grpc.inprocess.InProcessChannelBuilder diff --git a/src/test/kotlin/org/wfanet/measurement/duchy/daemon/herald/HeraldTest.kt b/src/test/kotlin/org/wfanet/measurement/duchy/herald/HeraldTest.kt similarity index 99% rename from src/test/kotlin/org/wfanet/measurement/duchy/daemon/herald/HeraldTest.kt rename to src/test/kotlin/org/wfanet/measurement/duchy/herald/HeraldTest.kt index b1b9732829f..24897b09242 100644 --- a/src/test/kotlin/org/wfanet/measurement/duchy/daemon/herald/HeraldTest.kt +++ b/src/test/kotlin/org/wfanet/measurement/duchy/herald/HeraldTest.kt @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -package org.wfanet.measurement.duchy.daemon.herald +package org.wfanet.measurement.duchy.herald import com.google.common.truth.Truth.assertThat import com.google.common.truth.extensions.proto.ProtoTruth.assertThat @@ -66,9 +66,6 @@ import org.wfanet.measurement.common.pack import org.wfanet.measurement.common.testing.captureFirst import org.wfanet.measurement.common.testing.chainRulesSequentially import org.wfanet.measurement.common.testing.verifyProtoArgument -import org.wfanet.measurement.duchy.daemon.testing.TestRequisition -import org.wfanet.measurement.duchy.daemon.utils.key -import org.wfanet.measurement.duchy.daemon.utils.toDuchyEncryptionPublicKey import org.wfanet.measurement.duchy.db.computation.testing.FakeComputationsDatabase import org.wfanet.measurement.duchy.service.internal.computations.ComputationsService import org.wfanet.measurement.duchy.service.internal.computations.newEmptyOutputBlobMetadata @@ -77,7 +74,10 @@ import org.wfanet.measurement.duchy.service.internal.computations.newPassThrough import org.wfanet.measurement.duchy.storage.ComputationStore import org.wfanet.measurement.duchy.storage.RequisitionStore import org.wfanet.measurement.duchy.storage.TinkKeyStore +import org.wfanet.measurement.duchy.testing.TestRequisition import org.wfanet.measurement.duchy.toProtocolStage +import org.wfanet.measurement.duchy.utils.key +import org.wfanet.measurement.duchy.utils.toDuchyEncryptionPublicKey import org.wfanet.measurement.internal.duchy.ComputationDetails import org.wfanet.measurement.internal.duchy.ComputationDetailsKt.kingdomComputationDetails import org.wfanet.measurement.internal.duchy.ComputationsGrpcKt.ComputationsCoroutineImplBase as InternalComputationsCoroutineImplBase diff --git a/src/test/kotlin/org/wfanet/measurement/duchy/daemon/mill/liquidlegionsv2/BUILD.bazel b/src/test/kotlin/org/wfanet/measurement/duchy/mill/liquidlegionsv2/BUILD.bazel similarity index 87% rename from src/test/kotlin/org/wfanet/measurement/duchy/daemon/mill/liquidlegionsv2/BUILD.bazel rename to src/test/kotlin/org/wfanet/measurement/duchy/mill/liquidlegionsv2/BUILD.bazel index fff15e80c73..01fae2e6100 100644 --- a/src/test/kotlin/org/wfanet/measurement/duchy/daemon/mill/liquidlegionsv2/BUILD.bazel +++ b/src/test/kotlin/org/wfanet/measurement/duchy/mill/liquidlegionsv2/BUILD.bazel @@ -7,13 +7,13 @@ kt_jvm_library( name = "reach_frequency_liquid_legions_v2_mill_test", srcs = ["ReachFrequencyLiquidLegionsV2MillTest.kt"], deps = [ - "//src/main/kotlin/org/wfanet/measurement/duchy/daemon/mill/liquidlegionsv2:reach_frequency_liquid_legions_v2_mill", - "//src/main/kotlin/org/wfanet/measurement/duchy/daemon/testing", "//src/main/kotlin/org/wfanet/measurement/duchy/db/computation", "//src/main/kotlin/org/wfanet/measurement/duchy/db/computation/testing", + "//src/main/kotlin/org/wfanet/measurement/duchy/mill/liquidlegionsv2:reach_frequency_liquid_legions_v2_mill", "//src/main/kotlin/org/wfanet/measurement/duchy/service/internal/computations", "//src/main/kotlin/org/wfanet/measurement/duchy/storage:computation_store", "//src/main/kotlin/org/wfanet/measurement/duchy/storage:requisition_store", + "//src/main/kotlin/org/wfanet/measurement/duchy/testing", "//src/main/kotlin/org/wfanet/measurement/system/v1alpha:resource_key", "//src/main/proto/wfa/measurement/api/v2alpha:crypto_kt_jvm_proto", "//src/main/proto/wfa/measurement/internal/duchy/protocol:liquid_legions_v2_encryption_methods_kt_jvm_proto", @@ -36,7 +36,7 @@ kt_jvm_library( # TODO(bazelbuild/rules_kotlin#1088): Use kt_jvm_test when fixed. java_test( name = "ReachFrequencyLiquidLegionsV2MillTest", - test_class = "org.wfanet.measurement.duchy.daemon.mill.liquidlegionsv2.crypto.ReachFrequencyLiquidLegionsV2MillTest", + test_class = "org.wfanet.measurement.duchy.mill.liquidlegionsv2.ReachFrequencyLiquidLegionsV2MillTest", runtime_deps = [":reach_frequency_liquid_legions_v2_mill_test"], ) @@ -44,13 +44,13 @@ kt_jvm_library( name = "reach_only_liquid_legions_v2_mill_test", srcs = ["ReachOnlyLiquidLegionsV2MillTest.kt"], deps = [ - "//src/main/kotlin/org/wfanet/measurement/duchy/daemon/mill/liquidlegionsv2:reach_only_liquid_legions_v2_mill", - "//src/main/kotlin/org/wfanet/measurement/duchy/daemon/testing", "//src/main/kotlin/org/wfanet/measurement/duchy/db/computation", "//src/main/kotlin/org/wfanet/measurement/duchy/db/computation/testing", + "//src/main/kotlin/org/wfanet/measurement/duchy/mill/liquidlegionsv2:reach_only_liquid_legions_v2_mill", "//src/main/kotlin/org/wfanet/measurement/duchy/service/internal/computations", "//src/main/kotlin/org/wfanet/measurement/duchy/storage:computation_store", "//src/main/kotlin/org/wfanet/measurement/duchy/storage:requisition_store", + "//src/main/kotlin/org/wfanet/measurement/duchy/testing", "//src/main/kotlin/org/wfanet/measurement/system/v1alpha:resource_key", "//src/main/proto/wfa/measurement/api/v2alpha:crypto_kt_jvm_proto", "//src/main/proto/wfa/measurement/internal/duchy/protocol:reach_only_liquid_legions_v2_encryption_methods_kt_jvm_proto", @@ -72,6 +72,6 @@ kt_jvm_library( # TODO(bazelbuild/rules_kotlin#1088): Use kt_jvm_test when fixed. java_test( name = "ReachOnlyLiquidLegionsV2MillTest", - test_class = "org.wfanet.measurement.duchy.daemon.mill.liquidlegionsv2.crypto.ReachOnlyLiquidLegionsV2MillTest", + test_class = "org.wfanet.measurement.duchy.mill.liquidlegionsv2.ReachOnlyLiquidLegionsV2MillTest", runtime_deps = [":reach_only_liquid_legions_v2_mill_test"], ) diff --git a/src/test/kotlin/org/wfanet/measurement/duchy/daemon/mill/liquidlegionsv2/ReachFrequencyLiquidLegionsV2MillTest.kt b/src/test/kotlin/org/wfanet/measurement/duchy/mill/liquidlegionsv2/ReachFrequencyLiquidLegionsV2MillTest.kt similarity index 98% rename from src/test/kotlin/org/wfanet/measurement/duchy/daemon/mill/liquidlegionsv2/ReachFrequencyLiquidLegionsV2MillTest.kt rename to src/test/kotlin/org/wfanet/measurement/duchy/mill/liquidlegionsv2/ReachFrequencyLiquidLegionsV2MillTest.kt index 10e2074491f..1305541bd5b 100644 --- a/src/test/kotlin/org/wfanet/measurement/duchy/daemon/mill/liquidlegionsv2/ReachFrequencyLiquidLegionsV2MillTest.kt +++ b/src/test/kotlin/org/wfanet/measurement/duchy/mill/liquidlegionsv2/ReachFrequencyLiquidLegionsV2MillTest.kt @@ -23,7 +23,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -package org.wfanet.measurement.duchy.daemon.mill.liquidlegionsv2.crypto +package org.wfanet.measurement.duchy.mill.liquidlegionsv2 import com.google.common.truth.Truth.assertThat import com.google.common.truth.extensions.proto.ProtoTruth.assertThat @@ -67,12 +67,10 @@ import org.wfanet.measurement.common.identity.DuchyInfo import org.wfanet.measurement.common.testing.chainRulesSequentially import org.wfanet.measurement.common.testing.verifyProtoArgument import org.wfanet.measurement.consent.client.common.toEncryptionPublicKey -import org.wfanet.measurement.duchy.daemon.mill.Certificate -import org.wfanet.measurement.duchy.daemon.mill.liquidlegionsv2.ReachFrequencyLiquidLegionsV2Mill -import org.wfanet.measurement.duchy.daemon.testing.TestRequisition -import org.wfanet.measurement.duchy.daemon.utils.toDuchyEncryptionPublicKey import org.wfanet.measurement.duchy.db.computation.ComputationDataClients import org.wfanet.measurement.duchy.db.computation.testing.FakeComputationsDatabase +import org.wfanet.measurement.duchy.mill.Certificate +import org.wfanet.measurement.duchy.mill.liquidlegionsv2.crypto.LiquidLegionsV2Encryption import org.wfanet.measurement.duchy.service.internal.computations.ComputationsService import org.wfanet.measurement.duchy.service.internal.computations.newEmptyOutputBlobMetadata import org.wfanet.measurement.duchy.service.internal.computations.newInputBlobMetadata @@ -80,7 +78,9 @@ import org.wfanet.measurement.duchy.storage.ComputationBlobContext import org.wfanet.measurement.duchy.storage.ComputationStore import org.wfanet.measurement.duchy.storage.RequisitionBlobContext import org.wfanet.measurement.duchy.storage.RequisitionStore +import org.wfanet.measurement.duchy.testing.TestRequisition import org.wfanet.measurement.duchy.toProtocolStage +import org.wfanet.measurement.duchy.utils.toDuchyEncryptionPublicKey import org.wfanet.measurement.internal.duchy.ComputationBlobDependency import org.wfanet.measurement.internal.duchy.ComputationDetails.CompletedReason import org.wfanet.measurement.internal.duchy.ComputationDetailsKt.kingdomComputationDetails @@ -641,7 +641,7 @@ class ReachFrequencyLiquidLegionsV2MillTest { fakeComputationDb.claimedComputationIds.clear() } - nonAggregatorMill.pollAndProcessNextComputation() + nonAggregatorMill.claimAndProcessWork() assertThat(fakeComputationDb[LOCAL_ID]) .isEqualTo( @@ -691,7 +691,7 @@ class ReachFrequencyLiquidLegionsV2MillTest { } // Mill should claim computation1 of INITIALIZATION_PHASE. - nonAggregatorMill.pollAndProcessNextComputation() + nonAggregatorMill.claimAndProcessWork() assertThat(fakeComputationDb[2]!!.computationStage) .isEqualTo(WAIT_REQUISITIONS_AND_KEY_SET.toProtocolStage()) @@ -753,7 +753,7 @@ class ReachFrequencyLiquidLegionsV2MillTest { } // Stage 1. Process the above computation - nonAggregatorMill.pollAndProcessNextComputation() + nonAggregatorMill.claimAndProcessWork() // Stage 2. Check the status of the computation assertThat(fakeComputationDb[LOCAL_ID]) @@ -872,7 +872,7 @@ class ReachFrequencyLiquidLegionsV2MillTest { } // Stage 1. Process the above computation - nonAggregatorMill.pollAndProcessNextComputation() + nonAggregatorMill.claimAndProcessWork() testScheduler.advanceUntilIdle() // Stage 3. Verify stage was advanced successfully. @@ -910,7 +910,7 @@ class ReachFrequencyLiquidLegionsV2MillTest { .thenReturn(ComputationLogEntry.getDefaultInstance()) // Stage 1. Process the above computation - aggregatorMill.pollAndProcessNextComputation() + aggregatorMill.claimAndProcessWork() // Stage 2. Check the status of the computation assertThat(fakeComputationDb[LOCAL_ID]!!) @@ -989,7 +989,7 @@ class ReachFrequencyLiquidLegionsV2MillTest { .thenReturn(computationParticipant) // Stage 1. Process the above computation - aggregatorMill.pollAndProcessNextComputation() + aggregatorMill.claimAndProcessWork() // Stage 2. Check the status of the computation assertThat(fakeComputationDb[LOCAL_ID]!!) @@ -1054,7 +1054,7 @@ class ReachFrequencyLiquidLegionsV2MillTest { .thenReturn(ComputationLogEntry.getDefaultInstance()) // Stage 1. Process the above computation - aggregatorMill.pollAndProcessNextComputation() + aggregatorMill.claimAndProcessWork() // Stage 2. Check the status of the computation assertThat(fakeComputationDb[LOCAL_ID]!!) @@ -1133,7 +1133,7 @@ class ReachFrequencyLiquidLegionsV2MillTest { .thenReturn(ComputationLogEntry.getDefaultInstance()) // Stage 1. Process the above computation - aggregatorMill.pollAndProcessNextComputation() + aggregatorMill.claimAndProcessWork() // Stage 2. Check the status of the computation assertThat(fakeComputationDb[LOCAL_ID]!!) @@ -1221,7 +1221,7 @@ class ReachFrequencyLiquidLegionsV2MillTest { } // Stage 1. Process the above computation - nonAggregatorMill.pollAndProcessNextComputation() + nonAggregatorMill.claimAndProcessWork() // Stage 2. Check the status of the computation val blobKey = calculatedBlobContext.blobKey @@ -1330,7 +1330,7 @@ class ReachFrequencyLiquidLegionsV2MillTest { } // Stage 1. Process the above computation - aggregatorMill.pollAndProcessNextComputation() + aggregatorMill.claimAndProcessWork() // Stage 2. Check the status of the computation val blobKey = ComputationBlobContext(GLOBAL_ID, SETUP_PHASE.toProtocolStage(), 3L).blobKey @@ -1442,7 +1442,7 @@ class ReachFrequencyLiquidLegionsV2MillTest { } // Stage 1. Process the above computation - aggregatorMill.pollAndProcessNextComputation() + aggregatorMill.claimAndProcessWork() // Stage 2. Check the status of the computation val blobKey = ComputationBlobContext(GLOBAL_ID, SETUP_PHASE.toProtocolStage(), 3L).blobKey @@ -1535,7 +1535,7 @@ class ReachFrequencyLiquidLegionsV2MillTest { ) // Stage 1. Process the above computation - nonAggregatorMill.pollAndProcessNextComputation() + nonAggregatorMill.claimAndProcessWork() // Stage 2. Check the status of the computation assertThat(fakeComputationDb[LOCAL_ID]) @@ -1603,7 +1603,7 @@ class ReachFrequencyLiquidLegionsV2MillTest { } // Stage 1. Process the above computation - nonAggregatorMill.pollAndProcessNextComputation() + nonAggregatorMill.claimAndProcessWork() // Stage 2. Check the status of the computation val blobKey = calculatedBlobContext.blobKey @@ -1692,7 +1692,7 @@ class ReachFrequencyLiquidLegionsV2MillTest { } // Stage 1. Process the above computation - aggregatorMill.pollAndProcessNextComputation() + aggregatorMill.claimAndProcessWork() // Stage 2. Check the status of the computation val blobKey = calculatedBlobContext.blobKey @@ -1789,7 +1789,7 @@ class ReachFrequencyLiquidLegionsV2MillTest { } // Process the above computation - aggregatorMill.pollAndProcessNextComputation() + aggregatorMill.claimAndProcessWork() // Check that the request sent to the crypto worker was correct. assertThat(cryptoRequest) @@ -1842,7 +1842,7 @@ class ReachFrequencyLiquidLegionsV2MillTest { } // Process the above computation - aggregatorMill.pollAndProcessNextComputation() + aggregatorMill.claimAndProcessWork() // Check that the request sent to the crypto worker was correct. assertThat(cryptoRequest) @@ -1894,7 +1894,7 @@ class ReachFrequencyLiquidLegionsV2MillTest { } // Stage 1. Process the above computation - nonAggregatorMill.pollAndProcessNextComputation() + nonAggregatorMill.claimAndProcessWork() // Stage 2. Check the status of the computation val blobKey = calculatedBlobContext.blobKey @@ -1989,7 +1989,7 @@ class ReachFrequencyLiquidLegionsV2MillTest { } // Process the above computation - nonAggregatorMill.pollAndProcessNextComputation() + nonAggregatorMill.claimAndProcessWork() // Check that the request sent to the crypto worker was correct. assertThat(cryptoRequest) @@ -2041,7 +2041,7 @@ class ReachFrequencyLiquidLegionsV2MillTest { } // Process the above computation - nonAggregatorMill.pollAndProcessNextComputation() + nonAggregatorMill.claimAndProcessWork() // Check that the request sent to the crypto worker was correct. assertThat(cryptoRequest) @@ -2103,7 +2103,7 @@ class ReachFrequencyLiquidLegionsV2MillTest { } // Stage 1. Process the above computation - aggregatorMill.pollAndProcessNextComputation() + aggregatorMill.claimAndProcessWork() // Stage 2. Check the status of the computation val blobKey = calculatedBlobContext.blobKey @@ -2228,7 +2228,7 @@ class ReachFrequencyLiquidLegionsV2MillTest { } // Stage 1. Process the above computation - aggregatorMill.pollAndProcessNextComputation() + aggregatorMill.claimAndProcessWork() // Stage 2. Check the status of the computation assertThat(fakeComputationDb[LOCAL_ID]) @@ -2358,7 +2358,7 @@ class ReachFrequencyLiquidLegionsV2MillTest { } // Process the above computation - aggregatorMill.pollAndProcessNextComputation() + aggregatorMill.claimAndProcessWork() // Check the status of the computation val blobKey = calculatedBlobContext.blobKey @@ -2470,7 +2470,7 @@ class ReachFrequencyLiquidLegionsV2MillTest { } // Process the above computation - aggregatorMill.pollAndProcessNextComputation() + aggregatorMill.claimAndProcessWork() // Check the status of the computation val blobKey = calculatedBlobContext.blobKey @@ -2567,7 +2567,7 @@ class ReachFrequencyLiquidLegionsV2MillTest { } // Stage 1. Process the above computation - nonAggregatorMill.pollAndProcessNextComputation() + nonAggregatorMill.claimAndProcessWork() // Stage 2. Check the status of the computation val blobKey = calculatedBlobContext.blobKey @@ -2658,7 +2658,7 @@ class ReachFrequencyLiquidLegionsV2MillTest { } // Stage 1. Process the above computation - aggregatorMill.pollAndProcessNextComputation() + aggregatorMill.claimAndProcessWork() // Stage 2. Check the status of the computation val blobKey = calculatedBlobContext.blobKey diff --git a/src/test/kotlin/org/wfanet/measurement/duchy/daemon/mill/liquidlegionsv2/ReachOnlyLiquidLegionsV2MillTest.kt b/src/test/kotlin/org/wfanet/measurement/duchy/mill/liquidlegionsv2/ReachOnlyLiquidLegionsV2MillTest.kt similarity index 98% rename from src/test/kotlin/org/wfanet/measurement/duchy/daemon/mill/liquidlegionsv2/ReachOnlyLiquidLegionsV2MillTest.kt rename to src/test/kotlin/org/wfanet/measurement/duchy/mill/liquidlegionsv2/ReachOnlyLiquidLegionsV2MillTest.kt index c697255ac0b..c99f9b85380 100644 --- a/src/test/kotlin/org/wfanet/measurement/duchy/daemon/mill/liquidlegionsv2/ReachOnlyLiquidLegionsV2MillTest.kt +++ b/src/test/kotlin/org/wfanet/measurement/duchy/mill/liquidlegionsv2/ReachOnlyLiquidLegionsV2MillTest.kt @@ -23,7 +23,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -package org.wfanet.measurement.duchy.daemon.mill.liquidlegionsv2.crypto +package org.wfanet.measurement.duchy.mill.liquidlegionsv2 import com.google.common.truth.Truth.assertThat import com.google.common.truth.extensions.proto.ProtoTruth.assertThat @@ -66,12 +66,10 @@ import org.wfanet.measurement.common.identity.DuchyInfo import org.wfanet.measurement.common.testing.chainRulesSequentially import org.wfanet.measurement.common.testing.verifyProtoArgument import org.wfanet.measurement.consent.client.common.toEncryptionPublicKey -import org.wfanet.measurement.duchy.daemon.mill.Certificate -import org.wfanet.measurement.duchy.daemon.mill.liquidlegionsv2.ReachOnlyLiquidLegionsV2Mill -import org.wfanet.measurement.duchy.daemon.testing.TestRequisition -import org.wfanet.measurement.duchy.daemon.utils.toDuchyEncryptionPublicKey import org.wfanet.measurement.duchy.db.computation.ComputationDataClients import org.wfanet.measurement.duchy.db.computation.testing.FakeComputationsDatabase +import org.wfanet.measurement.duchy.mill.Certificate +import org.wfanet.measurement.duchy.mill.liquidlegionsv2.crypto.ReachOnlyLiquidLegionsV2Encryption import org.wfanet.measurement.duchy.service.internal.computations.ComputationsService import org.wfanet.measurement.duchy.service.internal.computations.newEmptyOutputBlobMetadata import org.wfanet.measurement.duchy.service.internal.computations.newInputBlobMetadata @@ -79,7 +77,9 @@ import org.wfanet.measurement.duchy.storage.ComputationBlobContext import org.wfanet.measurement.duchy.storage.ComputationStore import org.wfanet.measurement.duchy.storage.RequisitionBlobContext import org.wfanet.measurement.duchy.storage.RequisitionStore +import org.wfanet.measurement.duchy.testing.TestRequisition import org.wfanet.measurement.duchy.toProtocolStage +import org.wfanet.measurement.duchy.utils.toDuchyEncryptionPublicKey import org.wfanet.measurement.internal.duchy.ComputationBlobDependency import org.wfanet.measurement.internal.duchy.ComputationDetails.CompletedReason import org.wfanet.measurement.internal.duchy.ComputationDetailsKt.kingdomComputationDetails @@ -517,7 +517,6 @@ class ReachOnlyLiquidLegionsV2MillTest { computationStatsClient = computationStatsStub, workerStubs = workerStubs, cryptoWorker = mockCryptoWorker, - workLockDuration = Duration.ofMinutes(5), openTelemetry = GlobalOpenTelemetry.get(), requestChunkSizeBytes = 20, maximumAttempts = 2, @@ -537,7 +536,6 @@ class ReachOnlyLiquidLegionsV2MillTest { computationStatsClient = computationStatsStub, workerStubs = workerStubs, cryptoWorker = mockCryptoWorker, - workLockDuration = Duration.ofMinutes(5), openTelemetry = GlobalOpenTelemetry.get(), requestChunkSizeBytes = 20, maximumAttempts = 2, @@ -579,7 +577,7 @@ class ReachOnlyLiquidLegionsV2MillTest { fakeComputationDb.claimedComputationIds.clear() } - nonAggregatorMill.pollAndProcessNextComputation() + nonAggregatorMill.claimAndProcessWork() assertThat(fakeComputationDb[LOCAL_ID]) .isEqualTo( @@ -640,7 +638,7 @@ class ReachOnlyLiquidLegionsV2MillTest { } // Mill should claim computation1 of INITIALIZATION_PHASE. - nonAggregatorMill.pollAndProcessNextComputation() + nonAggregatorMill.claimAndProcessWork() assertThat(fakeComputationDb[2]!!.computationStage) .isEqualTo(WAIT_REQUISITIONS_AND_KEY_SET.toProtocolStage()) @@ -697,7 +695,7 @@ class ReachOnlyLiquidLegionsV2MillTest { } // Stage 1. Process the above computation - nonAggregatorMill.pollAndProcessNextComputation() + nonAggregatorMill.claimAndProcessWork() // Stage 2. Check the status of the computation assertThat(fakeComputationDb[LOCAL_ID]) @@ -796,7 +794,7 @@ class ReachOnlyLiquidLegionsV2MillTest { .thenReturn(computationParticipant) // Stage 1. Process the above computation - aggregatorMill.pollAndProcessNextComputation() + aggregatorMill.claimAndProcessWork() // Stage 2. Check the status of the computation assertThat(fakeComputationDb[LOCAL_ID]!!) @@ -880,7 +878,7 @@ class ReachOnlyLiquidLegionsV2MillTest { .thenReturn(computationParticipant) // Stage 1. Process the above computation - aggregatorMill.pollAndProcessNextComputation() + aggregatorMill.claimAndProcessWork() // Stage 2. Check the status of the computation assertThat(fakeComputationDb[LOCAL_ID]!!) @@ -945,7 +943,7 @@ class ReachOnlyLiquidLegionsV2MillTest { .thenReturn(computationParticipant) // Stage 1. Process the above computation - aggregatorMill.pollAndProcessNextComputation() + aggregatorMill.claimAndProcessWork() // Stage 2. Check the status of the computation assertThat(fakeComputationDb[LOCAL_ID]!!) @@ -1036,7 +1034,7 @@ class ReachOnlyLiquidLegionsV2MillTest { .thenReturn(computationParticipant) // Stage 1. Process the above computation - aggregatorMill.pollAndProcessNextComputation() + aggregatorMill.claimAndProcessWork() // Stage 2. Check the status of the computation assertThat(fakeComputationDb[LOCAL_ID]!!) @@ -1109,7 +1107,7 @@ class ReachOnlyLiquidLegionsV2MillTest { ) // Stage 1. Process the above computation - nonAggregatorMill.pollAndProcessNextComputation() + nonAggregatorMill.claimAndProcessWork() // Stage 2. Check the status of the computation assertThat(fakeComputationDb[LOCAL_ID]) @@ -1177,7 +1175,7 @@ class ReachOnlyLiquidLegionsV2MillTest { } // Stage 1. Process the above computation - nonAggregatorMill.pollAndProcessNextComputation() + nonAggregatorMill.claimAndProcessWork() // Stage 2. Check the status of the computation val blobKey = calculatedBlobContext.blobKey @@ -1267,7 +1265,7 @@ class ReachOnlyLiquidLegionsV2MillTest { ) // Stage 1. Process the above computation - aggregatorMill.pollAndProcessNextComputation() + aggregatorMill.claimAndProcessWork() // Stage 2. Check the status of the computation assertThat(fakeComputationDb[LOCAL_ID]) @@ -1343,7 +1341,7 @@ class ReachOnlyLiquidLegionsV2MillTest { } // Stage 1. Process the above computation - aggregatorMill.pollAndProcessNextComputation() + aggregatorMill.claimAndProcessWork() // Stage 2. Check the status of the computation val blobKey = ComputationBlobContext(GLOBAL_ID, SETUP_PHASE.toProtocolStage(), 3L).blobKey @@ -1459,7 +1457,7 @@ class ReachOnlyLiquidLegionsV2MillTest { } // Stage 1. Process the above computation - aggregatorMill.pollAndProcessNextComputation() + aggregatorMill.claimAndProcessWork() // Stage 2. Check the status of the computation val blobKey = ComputationBlobContext(GLOBAL_ID, SETUP_PHASE.toProtocolStage(), 3L).blobKey @@ -1566,7 +1564,7 @@ class ReachOnlyLiquidLegionsV2MillTest { .thenReturn(computationParticipant) // Stage 1. Process the above computation - aggregatorMill.pollAndProcessNextComputation() + aggregatorMill.claimAndProcessWork() // Stage 2. Check the status of the computation assertThat(fakeComputationDb[LOCAL_ID]!!) @@ -1637,7 +1635,7 @@ class ReachOnlyLiquidLegionsV2MillTest { ) // Stage 1. Process the above computation - nonAggregatorMill.pollAndProcessNextComputation() + nonAggregatorMill.claimAndProcessWork() // Stage 2. Check the status of the computation assertThat(fakeComputationDb[LOCAL_ID]) @@ -1700,7 +1698,7 @@ class ReachOnlyLiquidLegionsV2MillTest { } // Stage 1. Process the above computation - nonAggregatorMill.pollAndProcessNextComputation() + nonAggregatorMill.claimAndProcessWork() // Stage 2. Check the status of the computation val blobKey = calculatedBlobContext.blobKey @@ -1769,7 +1767,7 @@ class ReachOnlyLiquidLegionsV2MillTest { .thenReturn(computationParticipant) // Stage 1. Process the above computation - nonAggregatorMill.pollAndProcessNextComputation() + nonAggregatorMill.claimAndProcessWork() // Stage 2. Check the status of the computation assertThat(fakeComputationDb[LOCAL_ID]!!) @@ -1840,7 +1838,7 @@ class ReachOnlyLiquidLegionsV2MillTest { ) // Stage 1. Process the above computation - aggregatorMill.pollAndProcessNextComputation() + aggregatorMill.claimAndProcessWork() // Stage 2. Check the status of the computation assertThat(fakeComputationDb[LOCAL_ID]) @@ -1905,7 +1903,7 @@ class ReachOnlyLiquidLegionsV2MillTest { } // Stage 1. Process the above computation - aggregatorMill.pollAndProcessNextComputation() + aggregatorMill.claimAndProcessWork() // Stage 2. Check the status of the computation val blobKey = calculatedBlobContext.blobKey @@ -2007,7 +2005,7 @@ class ReachOnlyLiquidLegionsV2MillTest { .thenReturn(computationParticipant) // Stage 1. Process the above computation - aggregatorMill.pollAndProcessNextComputation() + aggregatorMill.claimAndProcessWork() // Stage 2. Check the status of the computation assertThat(fakeComputationDb[LOCAL_ID]!!) diff --git a/src/test/kotlin/org/wfanet/measurement/duchy/daemon/mill/liquidlegionsv2/crypto/BUILD.bazel b/src/test/kotlin/org/wfanet/measurement/duchy/mill/liquidlegionsv2/crypto/BUILD.bazel similarity index 79% rename from src/test/kotlin/org/wfanet/measurement/duchy/daemon/mill/liquidlegionsv2/crypto/BUILD.bazel rename to src/test/kotlin/org/wfanet/measurement/duchy/mill/liquidlegionsv2/crypto/BUILD.bazel index edb28d8eeeb..aafd43564d9 100644 --- a/src/test/kotlin/org/wfanet/measurement/duchy/daemon/mill/liquidlegionsv2/crypto/BUILD.bazel +++ b/src/test/kotlin/org/wfanet/measurement/duchy/mill/liquidlegionsv2/crypto/BUILD.bazel @@ -7,7 +7,7 @@ kt_jvm_library( name = "liquid_legions_v2_encryption_utility_test", srcs = ["LiquidLegionsV2EncryptionUtilityTest.kt"], deps = [ - "//src/main/kotlin/org/wfanet/measurement/duchy/daemon/utils:computation_conversions", + "//src/main/kotlin/org/wfanet/measurement/duchy/utils:computation_conversions", "//src/main/proto/wfa/any_sketch:sketch_kt_jvm_proto", "//src/main/proto/wfa/any_sketch/crypto:sketch_encryption_methods_kt_jvm_proto", "//src/main/proto/wfa/measurement/internal/duchy/protocol:liquid_legions_v2_encryption_methods_kt_jvm_proto", @@ -23,7 +23,7 @@ kt_jvm_library( # TODO(bazelbuild/rules_kotlin#1088): Use kt_jvm_test when fixed. java_test( name = "LiquidLegionsV2EncryptionUtilityTest", - test_class = "org.wfanet.measurement.duchy.daemon.mill.liquidlegionsv2.crypto.LiquidLegionsV2EncryptionUtilityTest", + test_class = "org.wfanet.measurement.duchy.mill.liquidlegionsv2.crypto.LiquidLegionsV2EncryptionUtilityTest", runtime_deps = [":liquid_legions_v2_encryption_utility_test"], ) @@ -31,7 +31,7 @@ kt_jvm_library( name = "jni_liquid_legions_v2_encryption_test", srcs = ["JniLiquidLegionsV2EncryptionTest.kt"], deps = [ - "//src/main/kotlin/org/wfanet/measurement/duchy/daemon/mill/liquidlegionsv2/crypto:liquidlegionsv2encryption", + "//src/main/kotlin/org/wfanet/measurement/duchy/mill/liquidlegionsv2/crypto:liquidlegionsv2encryption", "//src/main/proto/wfa/measurement/internal/duchy/protocol:liquid_legions_v2_encryption_methods_kt_jvm_proto", "@wfa_common_jvm//imports/java/com/google/common/truth", "@wfa_common_jvm//imports/kotlin/kotlin/test", @@ -41,7 +41,7 @@ kt_jvm_library( # TODO(bazelbuild/rules_kotlin#1088): Use kt_jvm_test when fixed. java_test( name = "JniLiquidLegionsV2EncryptionTest", - test_class = "org.wfanet.measurement.duchy.daemon.mill.liquidlegionsv2.crypto.JniLiquidLegionsV2EncryptionTest", + test_class = "org.wfanet.measurement.duchy.mill.liquidlegionsv2.crypto.JniLiquidLegionsV2EncryptionTest", runtime_deps = [":jni_liquid_legions_v2_encryption_test"], ) @@ -49,7 +49,7 @@ kt_jvm_library( name = "reach_only_liquid_legions_v2_encryption_utility_test", srcs = ["ReachOnlyLiquidLegionsV2EncryptionUtilityTest.kt"], deps = [ - "//src/main/kotlin/org/wfanet/measurement/duchy/daemon/utils:computation_conversions", + "//src/main/kotlin/org/wfanet/measurement/duchy/utils:computation_conversions", "//src/main/proto/wfa/any_sketch:sketch_kt_jvm_proto", "//src/main/proto/wfa/any_sketch/crypto:sketch_encryption_methods_kt_jvm_proto", "//src/main/proto/wfa/measurement/internal/duchy/protocol:reach_only_liquid_legions_v2_encryption_methods_kt_jvm_proto", @@ -67,7 +67,7 @@ kt_jvm_library( # TODO(bazelbuild/rules_kotlin#1088): Use kt_jvm_test when fixed. java_test( name = "ReachOnlyLiquidLegionsV2EncryptionUtilityTest", - test_class = "org.wfanet.measurement.duchy.daemon.mill.liquidlegionsv2.crypto.ReachOnlyLiquidLegionsV2EncryptionUtilityTest", + test_class = "org.wfanet.measurement.duchy.mill.liquidlegionsv2.crypto.ReachOnlyLiquidLegionsV2EncryptionUtilityTest", runtime_deps = [":reach_only_liquid_legions_v2_encryption_utility_test"], ) @@ -75,7 +75,7 @@ kt_jvm_library( name = "jni_reach_only_liquid_legions_v2_encryption_test", srcs = ["JniReachOnlyLiquidLegionsV2EncryptionTest.kt"], deps = [ - "//src/main/kotlin/org/wfanet/measurement/duchy/daemon/mill/liquidlegionsv2/crypto:reachonlyliquidlegionsv2encryption", + "//src/main/kotlin/org/wfanet/measurement/duchy/mill/liquidlegionsv2/crypto:reachonlyliquidlegionsv2encryption", "//src/main/proto/wfa/measurement/internal/duchy/protocol:reach_only_liquid_legions_v2_encryption_methods_kt_jvm_proto", "@wfa_common_jvm//imports/java/com/google/common/truth", "@wfa_common_jvm//imports/kotlin/kotlin/test", @@ -85,6 +85,6 @@ kt_jvm_library( # TODO(bazelbuild/rules_kotlin#1088): Use kt_jvm_test when fixed. java_test( name = "JniReachOnlyLiquidLegionsV2EncryptionTest", - test_class = "org.wfanet.measurement.duchy.daemon.mill.liquidlegionsv2.crypto.JniReachOnlyLiquidLegionsV2EncryptionTest", + test_class = "org.wfanet.measurement.duchy.mill.liquidlegionsv2.crypto.JniReachOnlyLiquidLegionsV2EncryptionTest", runtime_deps = [":jni_reach_only_liquid_legions_v2_encryption_test"], ) diff --git a/src/test/kotlin/org/wfanet/measurement/duchy/daemon/mill/liquidlegionsv2/crypto/JniLiquidLegionsV2EncryptionTest.kt b/src/test/kotlin/org/wfanet/measurement/duchy/mill/liquidlegionsv2/crypto/JniLiquidLegionsV2EncryptionTest.kt similarity index 95% rename from src/test/kotlin/org/wfanet/measurement/duchy/daemon/mill/liquidlegionsv2/crypto/JniLiquidLegionsV2EncryptionTest.kt rename to src/test/kotlin/org/wfanet/measurement/duchy/mill/liquidlegionsv2/crypto/JniLiquidLegionsV2EncryptionTest.kt index 7600c6cfe21..d7655dbb655 100644 --- a/src/test/kotlin/org/wfanet/measurement/duchy/daemon/mill/liquidlegionsv2/crypto/JniLiquidLegionsV2EncryptionTest.kt +++ b/src/test/kotlin/org/wfanet/measurement/duchy/mill/liquidlegionsv2/crypto/JniLiquidLegionsV2EncryptionTest.kt @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -package org.wfanet.measurement.duchy.daemon.mill.liquidlegionsv2.crypto +package org.wfanet.measurement.duchy.mill.liquidlegionsv2.crypto import com.google.common.truth.Truth.assertThat import kotlin.test.assertFailsWith diff --git a/src/test/kotlin/org/wfanet/measurement/duchy/daemon/mill/liquidlegionsv2/crypto/JniReachOnlyLiquidLegionsV2EncryptionTest.kt b/src/test/kotlin/org/wfanet/measurement/duchy/mill/liquidlegionsv2/crypto/JniReachOnlyLiquidLegionsV2EncryptionTest.kt similarity index 96% rename from src/test/kotlin/org/wfanet/measurement/duchy/daemon/mill/liquidlegionsv2/crypto/JniReachOnlyLiquidLegionsV2EncryptionTest.kt rename to src/test/kotlin/org/wfanet/measurement/duchy/mill/liquidlegionsv2/crypto/JniReachOnlyLiquidLegionsV2EncryptionTest.kt index 20676ac6984..7aec61efff1 100644 --- a/src/test/kotlin/org/wfanet/measurement/duchy/daemon/mill/liquidlegionsv2/crypto/JniReachOnlyLiquidLegionsV2EncryptionTest.kt +++ b/src/test/kotlin/org/wfanet/measurement/duchy/mill/liquidlegionsv2/crypto/JniReachOnlyLiquidLegionsV2EncryptionTest.kt @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -package org.wfanet.measurement.duchy.daemon.mill.liquidlegionsv2.crypto +package org.wfanet.measurement.duchy.mill.liquidlegionsv2.crypto import com.google.common.truth.Truth.assertThat import kotlin.test.assertFailsWith diff --git a/src/test/kotlin/org/wfanet/measurement/duchy/daemon/mill/liquidlegionsv2/crypto/LiquidLegionsV2EncryptionUtilityTest.kt b/src/test/kotlin/org/wfanet/measurement/duchy/mill/liquidlegionsv2/crypto/LiquidLegionsV2EncryptionUtilityTest.kt similarity index 98% rename from src/test/kotlin/org/wfanet/measurement/duchy/daemon/mill/liquidlegionsv2/crypto/LiquidLegionsV2EncryptionUtilityTest.kt rename to src/test/kotlin/org/wfanet/measurement/duchy/mill/liquidlegionsv2/crypto/LiquidLegionsV2EncryptionUtilityTest.kt index d869e3639d4..cf83ded3bfe 100644 --- a/src/test/kotlin/org/wfanet/measurement/duchy/daemon/mill/liquidlegionsv2/crypto/LiquidLegionsV2EncryptionUtilityTest.kt +++ b/src/test/kotlin/org/wfanet/measurement/duchy/mill/liquidlegionsv2/crypto/LiquidLegionsV2EncryptionUtilityTest.kt @@ -11,7 +11,7 @@ // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. -package org.wfanet.measurement.duchy.daemon.mill.liquidlegionsv2.crypto +package org.wfanet.measurement.duchy.mill.liquidlegionsv2.crypto import com.google.common.truth.Truth.assertThat import com.google.protobuf.ByteString @@ -27,8 +27,8 @@ import org.wfanet.anysketch.crypto.EncryptSketchRequest import org.wfanet.anysketch.crypto.EncryptSketchRequest.DestroyedRegisterStrategy.FLAGGED_KEY import org.wfanet.anysketch.crypto.EncryptSketchResponse import org.wfanet.anysketch.crypto.SketchEncrypterAdapter -import org.wfanet.measurement.duchy.daemon.utils.toAnySketchElGamalPublicKey -import org.wfanet.measurement.duchy.daemon.utils.toCmmsElGamalPublicKey +import org.wfanet.measurement.duchy.utils.toAnySketchElGamalPublicKey +import org.wfanet.measurement.duchy.utils.toCmmsElGamalPublicKey import org.wfanet.measurement.internal.duchy.protocol.CompleteExecutionPhaseOneAtAggregatorResponse import org.wfanet.measurement.internal.duchy.protocol.CompleteExecutionPhaseOneResponse import org.wfanet.measurement.internal.duchy.protocol.CompleteExecutionPhaseThreeAtAggregatorResponse diff --git a/src/test/kotlin/org/wfanet/measurement/duchy/daemon/mill/liquidlegionsv2/crypto/ReachOnlyLiquidLegionsV2EncryptionUtilityTest.kt b/src/test/kotlin/org/wfanet/measurement/duchy/mill/liquidlegionsv2/crypto/ReachOnlyLiquidLegionsV2EncryptionUtilityTest.kt similarity index 98% rename from src/test/kotlin/org/wfanet/measurement/duchy/daemon/mill/liquidlegionsv2/crypto/ReachOnlyLiquidLegionsV2EncryptionUtilityTest.kt rename to src/test/kotlin/org/wfanet/measurement/duchy/mill/liquidlegionsv2/crypto/ReachOnlyLiquidLegionsV2EncryptionUtilityTest.kt index 5208eda8385..2fb1cd54fe3 100644 --- a/src/test/kotlin/org/wfanet/measurement/duchy/daemon/mill/liquidlegionsv2/crypto/ReachOnlyLiquidLegionsV2EncryptionUtilityTest.kt +++ b/src/test/kotlin/org/wfanet/measurement/duchy/mill/liquidlegionsv2/crypto/ReachOnlyLiquidLegionsV2EncryptionUtilityTest.kt @@ -11,7 +11,7 @@ // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. -package org.wfanet.measurement.duchy.daemon.mill.liquidlegionsv2.crypto +package org.wfanet.measurement.duchy.mill.liquidlegionsv2.crypto import com.google.common.truth.Truth.assertThat import com.google.protobuf.ByteString @@ -29,8 +29,8 @@ import org.wfanet.anysketch.crypto.combineElGamalPublicKeysRequest import org.wfanet.anysketch.crypto.encryptSketchRequest import org.wfanet.anysketch.sketch import org.wfanet.estimation.Estimators -import org.wfanet.measurement.duchy.daemon.utils.toAnySketchElGamalPublicKey -import org.wfanet.measurement.duchy.daemon.utils.toCmmsElGamalPublicKey +import org.wfanet.measurement.duchy.utils.toAnySketchElGamalPublicKey +import org.wfanet.measurement.duchy.utils.toCmmsElGamalPublicKey import org.wfanet.measurement.internal.duchy.protocol.CompleteReachOnlyExecutionPhaseAtAggregatorResponse import org.wfanet.measurement.internal.duchy.protocol.CompleteReachOnlyExecutionPhaseResponse import org.wfanet.measurement.internal.duchy.protocol.CompleteReachOnlyInitializationPhaseResponse diff --git a/src/test/kotlin/org/wfanet/measurement/duchy/daemon/mill/shareshuffle/BUILD.bazel b/src/test/kotlin/org/wfanet/measurement/duchy/mill/shareshuffle/BUILD.bazel similarity index 89% rename from src/test/kotlin/org/wfanet/measurement/duchy/daemon/mill/shareshuffle/BUILD.bazel rename to src/test/kotlin/org/wfanet/measurement/duchy/mill/shareshuffle/BUILD.bazel index f9b48d088fd..0be73a3d421 100644 --- a/src/test/kotlin/org/wfanet/measurement/duchy/daemon/mill/shareshuffle/BUILD.bazel +++ b/src/test/kotlin/org/wfanet/measurement/duchy/mill/shareshuffle/BUILD.bazel @@ -7,14 +7,14 @@ kt_jvm_library( name = "honest_majority_share_shuffle_mill_test", srcs = ["HonestMajorityShareShuffleMillTest.kt"], deps = [ - "//src/main/kotlin/org/wfanet/measurement/duchy/daemon/mill/shareshuffle:honest_majority_share_shuffle_mill", - "//src/main/kotlin/org/wfanet/measurement/duchy/daemon/testing", "//src/main/kotlin/org/wfanet/measurement/duchy/db/computation", "//src/main/kotlin/org/wfanet/measurement/duchy/db/computation/testing", + "//src/main/kotlin/org/wfanet/measurement/duchy/mill/shareshuffle:honest_majority_share_shuffle_mill", "//src/main/kotlin/org/wfanet/measurement/duchy/service/internal/computations", "//src/main/kotlin/org/wfanet/measurement/duchy/storage:computation_store", "//src/main/kotlin/org/wfanet/measurement/duchy/storage:requisition_store", "//src/main/kotlin/org/wfanet/measurement/duchy/storage:tink_key_store", + "//src/main/kotlin/org/wfanet/measurement/duchy/testing", "//src/main/kotlin/org/wfanet/measurement/system/v1alpha:resource_key", "//src/main/proto/wfa/measurement/api/v2alpha:crypto_kt_jvm_proto", "//src/main/proto/wfa/measurement/internal/duchy/protocol:honest_majority_share_shuffle_methods_kt_jvm_proto", @@ -39,6 +39,6 @@ kt_jvm_library( # TODO(bazelbuild/rules_kotlin#1088): Use kt_jvm_test when fixed. java_test( name = "HonestMajorityShareShuffleMillTest", - test_class = "org.wfanet.measurement.duchy.daemon.mill.shareshuffle.HonestMajorityShareShuffleMillTest", + test_class = "org.wfanet.measurement.duchy.mill.shareshuffle.HonestMajorityShareShuffleMillTest", runtime_deps = [":honest_majority_share_shuffle_mill_test"], ) diff --git a/src/test/kotlin/org/wfanet/measurement/duchy/daemon/mill/shareshuffle/HonestMajorityShareShuffleMillTest.kt b/src/test/kotlin/org/wfanet/measurement/duchy/mill/shareshuffle/HonestMajorityShareShuffleMillTest.kt similarity index 98% rename from src/test/kotlin/org/wfanet/measurement/duchy/daemon/mill/shareshuffle/HonestMajorityShareShuffleMillTest.kt rename to src/test/kotlin/org/wfanet/measurement/duchy/mill/shareshuffle/HonestMajorityShareShuffleMillTest.kt index 9862d8b8d75..a2869045a7f 100644 --- a/src/test/kotlin/org/wfanet/measurement/duchy/daemon/mill/shareshuffle/HonestMajorityShareShuffleMillTest.kt +++ b/src/test/kotlin/org/wfanet/measurement/duchy/mill/shareshuffle/HonestMajorityShareShuffleMillTest.kt @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -package org.wfanet.measurement.duchy.daemon.mill.shareshuffle +package org.wfanet.measurement.duchy.mill.shareshuffle import com.google.common.truth.Truth.assertThat import com.google.common.truth.extensions.proto.FieldScopes @@ -78,13 +78,10 @@ import org.wfanet.measurement.consent.client.common.toEncryptionPublicKey import org.wfanet.measurement.consent.client.dataprovider.encryptRandomSeed import org.wfanet.measurement.consent.client.dataprovider.signRandomSeed import org.wfanet.measurement.consent.client.dataprovider.verifyEncryptionPublicKey -import org.wfanet.measurement.duchy.daemon.mill.Certificate -import org.wfanet.measurement.duchy.daemon.mill.shareshuffle.crypto.HonestMajorityShareShuffleCryptor -import org.wfanet.measurement.duchy.daemon.testing.TestRequisition -import org.wfanet.measurement.duchy.daemon.utils.toDuchyEncryptionPublicKey -import org.wfanet.measurement.duchy.daemon.utils.toV2AlphaEncryptionPublicKey import org.wfanet.measurement.duchy.db.computation.ComputationDataClients import org.wfanet.measurement.duchy.db.computation.testing.FakeComputationsDatabase +import org.wfanet.measurement.duchy.mill.Certificate +import org.wfanet.measurement.duchy.mill.shareshuffle.crypto.HonestMajorityShareShuffleCryptor import org.wfanet.measurement.duchy.service.internal.computations.ComputationsService import org.wfanet.measurement.duchy.service.internal.computations.newOutputBlobMetadata import org.wfanet.measurement.duchy.storage.ComputationBlobContext @@ -92,7 +89,10 @@ import org.wfanet.measurement.duchy.storage.ComputationStore import org.wfanet.measurement.duchy.storage.RequisitionBlobContext import org.wfanet.measurement.duchy.storage.RequisitionStore import org.wfanet.measurement.duchy.storage.TinkKeyStore +import org.wfanet.measurement.duchy.testing.TestRequisition import org.wfanet.measurement.duchy.toProtocolStage +import org.wfanet.measurement.duchy.utils.toDuchyEncryptionPublicKey +import org.wfanet.measurement.duchy.utils.toV2AlphaEncryptionPublicKey import org.wfanet.measurement.internal.duchy.ComputationBlobDependency import org.wfanet.measurement.internal.duchy.ComputationDetails import org.wfanet.measurement.internal.duchy.ComputationDetailsKt @@ -578,7 +578,7 @@ class HonestMajorityShareShuffleMillTest { requisitions = REACH_AND_FREQUENCY_REQUISITIONS, ) val mill = createHmssMill(DUCHY_ONE_ID) - mill.pollAndProcessNextComputation() + mill.claimAndProcessWork() assertThat(fakeComputationDb[2]!!.computationStage) .isEqualTo(Stage.WAIT_TO_START.toProtocolStage()) @@ -605,7 +605,7 @@ class HonestMajorityShareShuffleMillTest { ) val mill = createHmssMill(DUCHY_ONE_ID) - mill.pollAndProcessNextComputation() + mill.claimAndProcessWork() assertThat(fakeComputationDb[LOCAL_ID]) .isEqualTo( @@ -677,7 +677,7 @@ class HonestMajorityShareShuffleMillTest { ) val mill = createHmssMill(DUCHY_ONE_ID) - mill.pollAndProcessNextComputation() + mill.claimAndProcessWork() verify(mockComputationParticipants, never()).setParticipantRequisitionParams(any()) } @@ -712,7 +712,7 @@ class HonestMajorityShareShuffleMillTest { ) val mill = createHmssMill(DUCHY_ONE_ID) - mill.pollAndProcessNextComputation() + mill.claimAndProcessWork() testScheduler.advanceUntilIdle() // Verify stage was advanced successfully. @@ -733,7 +733,7 @@ class HonestMajorityShareShuffleMillTest { ) val mill = createHmssMill(DUCHY_ONE_ID) - mill.pollAndProcessNextComputation() + mill.claimAndProcessWork() val updatedToken = fakeComputationDb[LOCAL_ID] @@ -839,7 +839,7 @@ class HonestMajorityShareShuffleMillTest { ) val mill = createHmssMill(DUCHY_TWO_ID) - mill.pollAndProcessNextComputation() + mill.claimAndProcessWork() assertThat(fakeComputationDb[LOCAL_ID]) .isEqualTo( @@ -961,7 +961,7 @@ class HonestMajorityShareShuffleMillTest { } val mill = createHmssMill(DUCHY_ONE_ID) - mill.pollAndProcessNextComputation() + mill.claimAndProcessWork() assertThat(fakeComputationDb[LOCAL_ID]) .isEqualTo( @@ -1073,7 +1073,7 @@ class HonestMajorityShareShuffleMillTest { .thenThrow(Status.NOT_FOUND.asRuntimeException()) val mill = createHmssMill(DUCHY_ONE_ID) - mill.pollAndProcessNextComputation() + mill.claimAndProcessWork() assertThat(fakeComputationDb[LOCAL_ID]) .isEqualTo( @@ -1148,7 +1148,7 @@ class HonestMajorityShareShuffleMillTest { } val mill = createHmssMill(DUCHY_THREE_ID) - mill.pollAndProcessNextComputation() + mill.claimAndProcessWork() assertThat(fakeComputationDb[LOCAL_ID]) .isEqualTo( @@ -1270,7 +1270,7 @@ class HonestMajorityShareShuffleMillTest { } val mill = createHmssMill(DUCHY_ONE_ID) - mill.pollAndProcessNextComputation() + mill.claimAndProcessWork() assertThat(fakeComputationDb[LOCAL_ID]) .isEqualTo( @@ -1374,7 +1374,7 @@ class HonestMajorityShareShuffleMillTest { } val mill = createHmssMill(DUCHY_THREE_ID) - mill.pollAndProcessNextComputation() + mill.claimAndProcessWork() assertThat(fakeComputationDb[LOCAL_ID]) .isEqualTo( diff --git a/src/test/kotlin/org/wfanet/measurement/duchy/daemon/mill/shareshuffle/crypto/BUILD.bazel b/src/test/kotlin/org/wfanet/measurement/duchy/mill/shareshuffle/crypto/BUILD.bazel similarity index 75% rename from src/test/kotlin/org/wfanet/measurement/duchy/daemon/mill/shareshuffle/crypto/BUILD.bazel rename to src/test/kotlin/org/wfanet/measurement/duchy/mill/shareshuffle/crypto/BUILD.bazel index ef918f5c0b8..a8f59ba1a63 100644 --- a/src/test/kotlin/org/wfanet/measurement/duchy/daemon/mill/shareshuffle/crypto/BUILD.bazel +++ b/src/test/kotlin/org/wfanet/measurement/duchy/mill/shareshuffle/crypto/BUILD.bazel @@ -7,7 +7,7 @@ kt_jvm_library( name = "jni_honest_majority_share_shuffle_cryptor_test", srcs = ["JniHonestMajorityShareShuffleCryptorTest.kt"], deps = [ - "//src/main/kotlin/org/wfanet/measurement/duchy/daemon/mill/shareshuffle/crypto:honest_majority_share_shuffle_cryptor", + "//src/main/kotlin/org/wfanet/measurement/duchy/mill/shareshuffle/crypto:honest_majority_share_shuffle_cryptor", "//src/main/proto/wfa/measurement/internal/duchy/protocol:honest_majority_share_shuffle_methods_kt_jvm_proto", "@wfa_common_jvm//imports/java/com/google/common/truth", "@wfa_common_jvm//imports/kotlin/kotlin/test", @@ -17,6 +17,6 @@ kt_jvm_library( # TODO(bazelbuild/rules_kotlin#1088): Use kt_jvm_test when fixed. java_test( name = "JniHonestMajorityShareShuffleCryptorTest", - test_class = "org.wfanet.measurement.duchy.daemon.mill.shareshuffle.crypto.JniHonestMajorityShareShuffleCryptorTest", + test_class = "org.wfanet.measurement.duchy.mill.shareshuffle.crypto.JniHonestMajorityShareShuffleCryptorTest", runtime_deps = [":jni_honest_majority_share_shuffle_cryptor_test"], ) diff --git a/src/test/kotlin/org/wfanet/measurement/duchy/daemon/mill/shareshuffle/crypto/JniHonestMajorityShareShuffleCryptorTest.kt b/src/test/kotlin/org/wfanet/measurement/duchy/mill/shareshuffle/crypto/JniHonestMajorityShareShuffleCryptorTest.kt similarity index 94% rename from src/test/kotlin/org/wfanet/measurement/duchy/daemon/mill/shareshuffle/crypto/JniHonestMajorityShareShuffleCryptorTest.kt rename to src/test/kotlin/org/wfanet/measurement/duchy/mill/shareshuffle/crypto/JniHonestMajorityShareShuffleCryptorTest.kt index 14807ac770f..05d0fca3ca2 100644 --- a/src/test/kotlin/org/wfanet/measurement/duchy/daemon/mill/shareshuffle/crypto/JniHonestMajorityShareShuffleCryptorTest.kt +++ b/src/test/kotlin/org/wfanet/measurement/duchy/mill/shareshuffle/crypto/JniHonestMajorityShareShuffleCryptorTest.kt @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -package org.wfanet.measurement.duchy.daemon.mill.shareshuffle.crypto +package org.wfanet.measurement.duchy.mill.shareshuffle.crypto import com.google.common.truth.Truth.assertThat import kotlin.test.assertFailsWith diff --git a/src/test/kotlin/org/wfanet/measurement/duchy/daemon/utils/BUILD.bazel b/src/test/kotlin/org/wfanet/measurement/duchy/utils/BUILD.bazel similarity index 67% rename from src/test/kotlin/org/wfanet/measurement/duchy/daemon/utils/BUILD.bazel rename to src/test/kotlin/org/wfanet/measurement/duchy/utils/BUILD.bazel index f6a83455a29..4063d30aaf7 100644 --- a/src/test/kotlin/org/wfanet/measurement/duchy/daemon/utils/BUILD.bazel +++ b/src/test/kotlin/org/wfanet/measurement/duchy/utils/BUILD.bazel @@ -3,9 +3,9 @@ load("@wfa_rules_kotlin_jvm//kotlin:defs.bzl", "kt_jvm_test") kt_jvm_test( name = "DuchyOrderTest", srcs = ["DuchyOrderTest.kt"], - test_class = "org.wfanet.measurement.duchy.daemon.utils.DuchyOrderTest", + test_class = "org.wfanet.measurement.duchy.utils.DuchyOrderTest", deps = [ - "//src/main/kotlin/org/wfanet/measurement/duchy/daemon/utils:duchy_order", + "//src/main/kotlin/org/wfanet/measurement/duchy/utils:duchy_order", "@wfa_common_jvm//imports/java/com/google/common/truth", "@wfa_common_jvm//imports/java/org/junit", "@wfa_common_jvm//imports/kotlin/kotlin/test", diff --git a/src/test/kotlin/org/wfanet/measurement/duchy/daemon/utils/DuchyOrderTest.kt b/src/test/kotlin/org/wfanet/measurement/duchy/utils/DuchyOrderTest.kt similarity index 98% rename from src/test/kotlin/org/wfanet/measurement/duchy/daemon/utils/DuchyOrderTest.kt rename to src/test/kotlin/org/wfanet/measurement/duchy/utils/DuchyOrderTest.kt index d9ec2f836f1..ca8c1910052 100644 --- a/src/test/kotlin/org/wfanet/measurement/duchy/daemon/utils/DuchyOrderTest.kt +++ b/src/test/kotlin/org/wfanet/measurement/duchy/utils/DuchyOrderTest.kt @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -package org.wfanet.measurement.duchy.daemon.utils +package org.wfanet.measurement.duchy.utils import com.google.common.truth.Truth.assertThat import org.junit.Assert.assertEquals diff --git a/src/test/kotlin/org/wfanet/measurement/integration/k8s/BUILD.bazel b/src/test/kotlin/org/wfanet/measurement/integration/k8s/BUILD.bazel index 27c06392f6e..a032483be03 100644 --- a/src/test/kotlin/org/wfanet/measurement/integration/k8s/BUILD.bazel +++ b/src/test/kotlin/org/wfanet/measurement/integration/k8s/BUILD.bazel @@ -34,7 +34,7 @@ kt_jvm_library( ":abstract_correctness_test", "//imports/java/io/kubernetes/client", "//imports/java/io/kubernetes/client/extended", - "//src/main/kotlin/org/wfanet/measurement/common/k8s", + "//src/main/kotlin/org/wfanet/measurement/common/k8s:client", "//src/main/kotlin/org/wfanet/measurement/common/k8s/testing:port_forwarder", "//src/main/kotlin/org/wfanet/measurement/common/k8s/testing:processes", "//src/main/kotlin/org/wfanet/measurement/integration/common:configs", diff --git a/src/test/kotlin/org/wfanet/measurement/integration/k8s/EmptyClusterCorrectnessTest.kt b/src/test/kotlin/org/wfanet/measurement/integration/k8s/EmptyClusterCorrectnessTest.kt index 78e5991bdd6..0a39445fe6c 100644 --- a/src/test/kotlin/org/wfanet/measurement/integration/k8s/EmptyClusterCorrectnessTest.kt +++ b/src/test/kotlin/org/wfanet/measurement/integration/k8s/EmptyClusterCorrectnessTest.kt @@ -62,6 +62,7 @@ import org.wfanet.measurement.common.crypto.tink.TinkPrivateKeyHandle import org.wfanet.measurement.common.grpc.buildMutualTlsChannel import org.wfanet.measurement.common.grpc.withDefaultDeadline import org.wfanet.measurement.common.k8s.KubernetesClient +import org.wfanet.measurement.common.k8s.KubernetesClientImpl import org.wfanet.measurement.common.k8s.testing.PortForwarder import org.wfanet.measurement.common.k8s.testing.Processes import org.wfanet.measurement.common.testing.chainRulesSequentially @@ -484,7 +485,7 @@ class EmptyClusterCorrectnessTest : AbstractCorrectnessTest(measurementSystem) { private val measurementSystem = LocalMeasurementSystem( - lazy { KubernetesClient(ClientBuilder.defaultClient()) }, + lazy { KubernetesClientImpl(ClientBuilder.defaultClient()) }, lazy { tempDir }, lazy { UUID.randomUUID().toString() }, ) diff --git a/src/test/kotlin/org/wfanet/panelmatch/integration/k8s/AbstractPanelMatchCorrectnessTest.kt b/src/test/kotlin/org/wfanet/panelmatch/integration/k8s/AbstractPanelMatchCorrectnessTest.kt index c184ca45986..87c7084ad59 100644 --- a/src/test/kotlin/org/wfanet/panelmatch/integration/k8s/AbstractPanelMatchCorrectnessTest.kt +++ b/src/test/kotlin/org/wfanet/panelmatch/integration/k8s/AbstractPanelMatchCorrectnessTest.kt @@ -32,7 +32,7 @@ import org.junit.Test import org.wfanet.measurement.api.v2alpha.ExchangeWorkflow import org.wfanet.measurement.common.crypto.SigningCerts import org.wfanet.measurement.common.crypto.jceProvider -import org.wfanet.measurement.common.k8s.KubernetesClient +import org.wfanet.measurement.common.k8s.KubernetesClientImpl import org.wfanet.measurement.common.parseTextProto import org.wfanet.measurement.loadtest.panelmatch.PanelMatchSimulator @@ -91,7 +91,7 @@ abstract class AbstractPanelMatchCorrectnessTest(private val localSystem: PanelM Security.removeProvider(jceProvider.name) } - val k8sClient = KubernetesClient(ClientBuilder.defaultClient()) + val k8sClient = KubernetesClientImpl(ClientBuilder.defaultClient()) @JvmStatic protected val READY_TIMEOUT = Duration.ofMinutes(2L) private val TEST_DATA_PATH = Paths.get("wfa_measurement_system", "src", "main", "k8s", "panelmatch", "testing", "data") diff --git a/src/test/kotlin/org/wfanet/panelmatch/integration/k8s/BUILD.bazel b/src/test/kotlin/org/wfanet/panelmatch/integration/k8s/BUILD.bazel index 13bddab9a2f..6f2e717a111 100644 --- a/src/test/kotlin/org/wfanet/panelmatch/integration/k8s/BUILD.bazel +++ b/src/test/kotlin/org/wfanet/panelmatch/integration/k8s/BUILD.bazel @@ -16,7 +16,7 @@ kt_jvm_library( deps = [ "//imports/java/io/kubernetes/client", "//imports/java/io/kubernetes/client/extended", - "//src/main/kotlin/org/wfanet/measurement/common/k8s", + "//src/main/kotlin/org/wfanet/measurement/common/k8s:client", "//src/main/kotlin/org/wfanet/measurement/common/k8s/testing:port_forwarder", "//src/main/kotlin/org/wfanet/measurement/common/k8s/testing:processes", "//src/main/kotlin/org/wfanet/measurement/integration/common:configs", diff --git a/src/test/kotlin/org/wfanet/panelmatch/integration/k8s/EmptyClusterPanelMatchCorrectnessTest.kt b/src/test/kotlin/org/wfanet/panelmatch/integration/k8s/EmptyClusterPanelMatchCorrectnessTest.kt index c24ad217717..8c126916597 100644 --- a/src/test/kotlin/org/wfanet/panelmatch/integration/k8s/EmptyClusterPanelMatchCorrectnessTest.kt +++ b/src/test/kotlin/org/wfanet/panelmatch/integration/k8s/EmptyClusterPanelMatchCorrectnessTest.kt @@ -50,6 +50,7 @@ import org.wfanet.measurement.common.crypto.jceProvider import org.wfanet.measurement.common.grpc.buildMutualTlsChannel import org.wfanet.measurement.common.grpc.withDefaultDeadline import org.wfanet.measurement.common.k8s.KubernetesClient +import org.wfanet.measurement.common.k8s.KubernetesClientImpl import org.wfanet.measurement.common.k8s.testing.PortForwarder import org.wfanet.measurement.common.k8s.testing.Processes import org.wfanet.measurement.common.testing.chainRulesSequentially @@ -515,7 +516,7 @@ class EmptyClusterPanelMatchCorrectnessTest : AbstractPanelMatchCorrectnessTest( private val localSystem = LocalSystem( - lazy { KubernetesClient(ClientBuilder.defaultClient()) }, + lazy { KubernetesClientImpl(ClientBuilder.defaultClient()) }, lazy { tempDir }, lazy { UUID.randomUUID().toString() }, )