diff --git a/.github/workflows/README.md b/.github/workflows/README.md index 41287872ee25..3ff1aa5a10b4 100644 --- a/.github/workflows/README.md +++ b/.github/workflows/README.md @@ -556,4 +556,3 @@ PostCommit Jobs run in a schedule against master branch and generally do not get | [ Modify the GCP User Roles according to the infra/users.yml file ](https://github.com/apache/beam/actions/workflows/beam_Infrastructure_UsersPermissions.yml) | N/A | [![.github/workflows/beam_Infrastructure_UsersPermissions.yml](https://github.com/apache/beam/actions/workflows/beam_Infrastructure_UsersPermissions.yml/badge.svg?event=schedule)](https://github.com/apache/beam/actions/workflows/beam_Infrastructure_UsersPermissions.yml?query=event%3Aschedule) | | [ Service Account Keys Management ](https://github.com/apache/beam/actions/workflows/beam_Infrastructure_ServiceAccountKeys.yml) | N/A | [![.github/workflows/beam_Infrastructure_ServiceAccountKeys.yml](https://github.com/apache/beam/actions/workflows/beam_Infrastructure_ServiceAccountKeys.yml/badge.svg?event=schedule)](https://github.com/apache/beam/actions/workflows/beam_Infrastructure_ServiceAccountKeys.yml?query=event%3Aschedule) | | [ Upgrade GCP Libraries BOM ](https://github.com/apache/beam/actions/workflows/beam_Upgrade_GCP_BOM.yml) | N/A | [![.github/workflows/beam_Upgrade_GCP_BOM.yml](https://github.com/apache/beam/actions/workflows/beam_Upgrade_GCP_BOM.yml/badge.svg?event=schedule)](https://github.com/apache/beam/actions/workflows/beam_Upgrade_GCP_BOM.yml?query=event%3Aschedule) | -| [ Unmanaged Service Accounts Keys Audit ](https://github.com/apache/beam/actions/workflows/beam_Infrastructure_AuditUnmanagedKeys.yml) | N/A | [![.github/workflows/beam_Infrastructure_AuditUnmanagedKeys.yml](https://github.com/apache/beam/actions/workflows/beam_Infrastructure_AuditUnmanagedKeys.yml/badge.svg?event=schedule)](https://github.com/apache/beam/actions/workflows/beam_Infrastructure_AuditUnmanagedKeys.yml?query=event%3Aschedule) | diff --git a/.github/workflows/beam_Infrastructure_AuditUnmanagedKeys.yml b/.github/workflows/beam_Infrastructure_AuditUnmanagedKeys.yml deleted file mode 100644 index 37d6ecbd362c..000000000000 --- a/.github/workflows/beam_Infrastructure_AuditUnmanagedKeys.yml +++ /dev/null @@ -1,76 +0,0 @@ -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you 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. - -# This workflow works with the GCP security log analyzer to -# generate weekly security reports and initialize log sinks - -name: Unmanaged Service Accounts Keys Audit - -on: - workflow_dispatch: - schedule: - # Every day at 00:00 UTC - - cron: '0 0 * * *' - -concurrency: - group: ${{ github.workflow }} - cancel-in-progress: false - -permissions: - contents: read - issues: write - id-token: write - -jobs: - beam_UnmanagedKeysAudit: - name: Audit Unmanaged Service Account Keys - runs-on: [self-hosted, ubuntu-24.04, main] - timeout-minutes: 30 - steps: - - uses: actions/checkout@v7 - with: - persist-credentials: true - - - name: Setup gcloud - uses: google-github-actions/setup-gcloud@aa5489c8933f4cc7a4f7d45035b3b1440c9c10db - - - name: Setup Python - uses: actions/setup-python@v7 - with: - python-version: '3.13' - - - name: Install Python dependencies - working-directory: ./infra/enforcement - run: | - python -m pip install --upgrade pip - pip install -r requirements.txt - - - name: Run Unmanaged Service Account Keys Audit - working-directory: ./infra/enforcement - env: - GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} - GITHUB_REPOSITORY: ${{ github.repository }} - SMTP_SERVER: smtp.gmail.com - SMTP_PORT: 465 - EMAIL_ADDRESS: ${{ secrets.ISSUE_REPORT_SENDER_EMAIL_ADDRESS }} - EMAIL_PASSWORD: ${{ secrets.ISSUE_REPORT_SENDER_EMAIL_PASSWORD }} - EMAIL_RECIPIENT: "dev@beam.apache.org" - run: python account_keys.py --action announce - - - - diff --git a/.github/workflows/beam_Infrastructure_PolicyEnforcer.yml b/.github/workflows/beam_Infrastructure_PolicyEnforcer.yml index 52f851d2cdaa..730d45cf992f 100644 --- a/.github/workflows/beam_Infrastructure_PolicyEnforcer.yml +++ b/.github/workflows/beam_Infrastructure_PolicyEnforcer.yml @@ -35,6 +35,7 @@ concurrency: permissions: contents: read issues: write + id-token: write jobs: beam_Infrastructure_PolicyEnforcer: @@ -45,21 +46,21 @@ jobs: - uses: actions/checkout@v7 with: persist-credentials: false - + - name: Setup Python uses: actions/setup-python@v7 with: python-version: '3.13' - + - name: Install Python dependencies working-directory: ./infra/enforcement run: | python -m pip install --upgrade pip pip install -r requirements.txt - + - name: Setup gcloud uses: google-github-actions/setup-gcloud@aa5489c8933f4cc7a4f7d45035b3b1440c9c10db - + - name: Run IAM Policy Enforcement working-directory: ./infra/enforcement env: @@ -70,7 +71,7 @@ jobs: EMAIL_ADDRESS: ${{ secrets.ISSUE_REPORT_SENDER_EMAIL_ADDRESS }} EMAIL_PASSWORD: ${{ secrets.ISSUE_REPORT_SENDER_EMAIL_PASSWORD }} EMAIL_RECIPIENT: "dev@beam.apache.org" - run: python iam.py --action print + run: python iam.py --action announce - name: Run Account Keys Policy Enforcement working-directory: ./infra/enforcement @@ -82,4 +83,4 @@ jobs: EMAIL_ADDRESS: ${{ secrets.ISSUE_REPORT_SENDER_EMAIL_ADDRESS }} EMAIL_PASSWORD: ${{ secrets.ISSUE_REPORT_SENDER_EMAIL_PASSWORD }} EMAIL_RECIPIENT: "dev@beam.apache.org" - run: python account_keys.py --action print + run: python account_keys.py --action announce diff --git a/infra/enforcement/README.md b/infra/enforcement/README.md index 6d883f7e6806..b92e5f7e1802 100644 --- a/infra/enforcement/README.md +++ b/infra/enforcement/README.md @@ -134,16 +134,18 @@ The enforcement tools are integrated with GitHub Actions to provide automated co ### Workflow Configuration -The repository includes workflows for different security domains: -- **IAM Policy Enforcer** (`.github/workflows/beam_Infrastructure_PolicyEnforcer.yml`): Runs weekly on Mondays at 9:00 AM UTC. -- **Unmanaged Keys Audit** (`.github/workflows/beam_Infrastructure_AuditUnmanagedKeys.yml`): Runs daily at 00:00 UTC. It manages the continuous execution of the `account_keys.py` script to swiftly detect rogue service account keys generated outside the official rotation system. -- **Manual trigger**: Can be triggered manually via `workflow_dispatch` -- **Actions**: Runs both IAM and Account Keys enforcement with the `announce` action +The enforcement tools are consolidated into a single daily workflow (`.github/workflows/beam_Infrastructure_PolicyEnforcer.yml`) that runs automatically at 00:00 UTC. + +This unified workflow executes both security domains sequentially: +- **IAM Policy Enforcement:** Validates user bindings against the defined policies. +- **Unmanaged Keys Audit:** Detects rogue service account keys generated outside the official rotation system. **Note**: -- The email service is configured to use gmail -- The recipient email is set to `dev@beam.apache.org` for Apache Beam project notifications -- The `GITHUB_TOKEN` is automatically provided by GitHub Actions and doesn't need to be configured manually +- **Manual trigger**: The workflow can also be triggered manually via `workflow_dispatch`. +- **Actions**: It executes the respective Python scripts using the `announce` action. +- The email service is configured to use gmail. +- The recipient email is set to `dev@beam.apache.org` for Apache Beam project notifications. +- The `GITHUB_TOKEN` is automatically provided by GitHub Actions and doesn't need to be configured manually. ## Account Keys diff --git a/infra/enforcement/account_keys.py b/infra/enforcement/account_keys.py index 56ccbf654b03..a1248e3ee09c 100644 --- a/infra/enforcement/account_keys.py +++ b/infra/enforcement/account_keys.py @@ -19,6 +19,7 @@ import yaml import argparse import os +from datetime import datetime, timezone from typing import List, Dict, TypedDict, Optional from google.cloud import secretmanager from google.cloud import iam_admin_v1 @@ -26,6 +27,7 @@ from sending import SendingClient SECRET_MANAGER_LABEL = "beam-infra-secret-manager" +IAC_DRIFT_SA_KEY = "IAC_DRIFT_SA_KEY" class AuthorizedUser(TypedDict): email: str @@ -312,22 +314,24 @@ def check_compliance(self) -> List[str]: # Check that all service accounts that exist are declared for service_account in live_service_accounts: + if self._denormalize_account_email(service_account) not in [account["account_id"] for account in file_service_accounts]: msg = f"Service account '{service_account}' is not declared in the service account keys file." compliance_issues.append(msg) self.logger.warning(msg) - else: - iam_keys = self._get_user_managed_keys_from_iam(service_account) - if iam_keys: - secret_name = f"{self._denormalize_account_email(service_account)}-key" - legal_keys = [] - if secret_name in managed_secrets: - legal_keys = self._get_verified_keys_from_secret_manager(secret_name) - unmanaged_keys = set(iam_keys) - set(legal_keys) - for unmanaged_key in unmanaged_keys: - msg = f"SECURITY ALERT: Unmanaged key '{unmanaged_key}' detected on account '{service_account}'. This key was created outside of Beam's service account management system. " - compliance_issues.append(msg) - self.logger.warning(msg) + + iam_keys = self._get_user_managed_keys_from_iam(service_account) + + if iam_keys: + secret_name = f"{self._denormalize_account_email(service_account)}-key" + legal_keys = [] + if secret_name in managed_secrets: + legal_keys = self._get_verified_keys_from_secret_manager(secret_name) + unmanaged_keys = set(iam_keys) - set(legal_keys) + for unmanaged_key in unmanaged_keys: + msg = f"IAC_DRIFT_SA_KEY: Unmanaged key '{unmanaged_key}' detected on account '{service_account}'. This key was created outside of Beam's service account management system. " + compliance_issues.append(msg) + self.logger.warning(msg) extracted_secrets = [f"{self._denormalize_account_email(account['account_id'])}-key" for account in file_service_accounts] @@ -376,13 +380,14 @@ def create_announcement(self, recipient: str) -> None: self.logger.info("No compliance issues found, no announcement will be created.") return - unmanaged_keys_issues = [issue for issue in diff if "SECURITY ALERT" in issue] - general_issues = [issue for issue in diff if "SECURITY ALERT" not in issue] + unmanaged_keys_issues = [issue for issue in diff if IAC_DRIFT_SA_KEY in issue] + general_issues = [issue for issue in diff if IAC_DRIFT_SA_KEY not in issue] if general_issues: self.logger.info(f"Found {len(general_issues)} general compliance issues. Triggering announcement...") - title = f"Account Keys Compliance Issue Detected" - body = f"Account keys for project {self.project_id} are not compliant with the defined policies on {self.service_account_keys_file}\n\n" + title = f"[{IAC_DRIFT_SA_KEY}] Action Required: Unauthorized Service Accounts Detected" + body = f"Unauthorized Service Accounts Report\n\n" + body += f"Account keys for project {self.project_id} are not compliant with the defined policies on {self.service_account_keys_file}\n\n" for issue in general_issues: body += f"- {issue}\n" @@ -406,23 +411,44 @@ def print_announcement(self, recipient: str) -> None: """ if not self.sending_client: raise ValueError("SendingClient is required for printing announcements") - + diff = self.check_compliance() if not diff: self.logger.info("No compliance issues found, no announcement will be printed.") return - title = f"Account Keys Compliance Issue Detected" - body = f"Account keys for project {self.project_id} are not compliant with the defined policies on {self.service_account_keys_file}\n\n" - for issue in diff: - body += f"- {issue}\n" + unmanaged_keys_issues = [issue for issue in diff if IAC_DRIFT_SA_KEY in issue] + general_issues = [issue for issue in diff if IAC_DRIFT_SA_KEY not in issue] - announcement = f"Dear team,\n\nThis is an automated notification about compliance issues detected in the Account Keys policy for project {self.project_id}.\n\n" - announcement += f"We found {len(diff)} compliance issue(s) that need your attention.\n" - announcement += f"\nPlease check the GitHub issue for detailed information and take appropriate action to resolve these compliance violations." + if general_issues: + self.logger.info("Printing general compliance announcement...") + title = f"[IAC_DRIFT_SA_KEY] Action Required: Unauthorized Service Accounts Detected" + body = f"Unauthorized Service Accounts Report\n\n" + body += f"Account keys for project {self.project_id} are not compliant with the defined policies on {self.service_account_keys_file}\n\n" + for issue in general_issues: + body += f"- {issue}\n" - self.sending_client.print_announcement(title, body, recipient, announcement) + announcement = f"Dear team,\n\nThis is an automated notification about compliance issues detected in the Account Keys policy for project {self.project_id}.\n\n" + announcement += f"We found {len(general_issues)} compliance issue(s) that need your attention.\n" + announcement += f"\nPlease check the GitHub issue for detailed information and take appropriate action to resolve these compliance violations." + + self.sending_client.print_announcement(title, body, recipient, announcement) + + if unmanaged_keys_issues: + self.logger.info("Printing security dashboard update for unmanaged keys...") + timestamp = datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S UTC") + print("\n" + "="*60) + print("SIMULATING GITHUB SECURITY ISSUE CREATION/UPDATE") + print("="*60) + print(f"Title: [{IAC_DRIFT_SA_KEY}] Action Required: Unmanaged Service Account Keys Detected\n") + print(f"Body:\n### Unmanaged Keys Audit Report ({timestamp})") + print(f"The following unauthorized or unmanaged keys were detected in `{self.project_id}`:\n") + for issue in unmanaged_keys_issues: + print(f"- {issue}") + print("\n*Please investigate and revoke these keys if they are not part of the official rotation system.*\n") + print("### History\n
\nClick to expand\n\n[... Previous reports would be collapsed here ...]\n
") + print("="*60 + "\n") def generate_compliance(self) -> None: """ diff --git a/infra/enforcement/iam.py b/infra/enforcement/iam.py index c4c65c7c679c..14dbf455a256 100644 --- a/infra/enforcement/iam.py +++ b/infra/enforcement/iam.py @@ -21,10 +21,12 @@ import yaml from google.api_core import exceptions from google.cloud import resourcemanager_v3 -from typing import Optional, List, Dict, Tuple +from typing import Optional, List, Dict from sending import SendingClient +from datetime import datetime, timezone CONFIG_FILE = "config.yml" +IAC_DRIFT_IAM_USER = "IAC_DRIFT_IAM_USER" class IAMPolicyComplianceChecker: @@ -214,7 +216,7 @@ def check_compliance(self) -> List[str]: self.logger.info(error_msg) raise RuntimeError(error_msg) - differences = [] + differences = [] all_emails = set(current_users.keys()) | set(existing_users.keys()) @@ -223,7 +225,7 @@ def check_compliance(self) -> List[str]: existing_user = existing_users.get(email) if current_user and not existing_user: - differences.append(f"User {email} not found in existing policy.") + differences.append(f"IAC_DRIFT_IAM_USER: Unauthorized user '{email}' detected in GCP but not found in existing policy.") elif not current_user and existing_user: differences.append(f"User {email} found in policy file but not in GCP.") elif current_user and existing_user: @@ -247,51 +249,86 @@ def create_announcement(self, recipient: str) -> None: """ if not self.sending_client: raise ValueError("SendingClient is required for creating announcements") - diff = self.check_compliance() if not diff: self.logger.info("No compliance issues found, no announcement will be created.") return - title = f"IAM Policy Non-Compliance Detected" - body = f"IAM policy for project {self.project_id} is not compliant with the defined policies on {self.users_file}\n\n" - for issue in diff: - body += f"- {issue}\n" + iam_drift_issues = [issue for issue in diff if IAC_DRIFT_IAM_USER in issue] + general_issues = [issue for issue in diff if IAC_DRIFT_IAM_USER not in issue] + + if general_issues: + self.logger.info(f"Found {len(general_issues)} general IAM compliance issues. Triggering announcement...") + title = f"IAM Policy Non-Compliance Detected" + body = f"IAM policy for project {self.project_id} is not compliant with the defined policies on {self.users_file}\n\n" + for issue in general_issues: + body += f"- {issue}\n" + + announcement = f"Dear team,\n\nThis is an automated notification about compliance issues detected in the IAM policy for project {self.project_id}.\n\n" + announcement += f"We found {len(general_issues)} compliance issue(s) that need your attention.\n" + announcement += f"\nPlease check the GitHub issue for detailed information and take appropriate action to resolve these compliance violations." + + self.sending_client.create_announcement(title, body, recipient, announcement) - announcement = f"Dear team,\n\nThis is an automated notification about compliance issues detected in the IAM policy for project {self.project_id}.\n\n" - announcement += f"We found {len(diff)} compliance issue(s) that need your attention.\n" - announcement += f"\nPlease check the GitHub issue for detailed information and take appropriate action to resolve these compliance violations." + if iam_drift_issues: + self.logger.info(f"Found {len(iam_drift_issues)} critical IAM security alerts. Dispatching to GitHub security issue...") + title = f"[{IAC_DRIFT_IAM_USER}] Action Required: Unauthorized IAM Users Detected" + body = f"Critical security violations detected in IAM policies for project {self.project_id}:\n\n" + for issue in iam_drift_issues: + body += f"- {issue}\n" - self.sending_client.create_announcement(title, body, recipient, announcement) + announcement = f"URGENT: Dear team,\n\nThis is an automated security alert regarding unauthorized IAM access in project {self.project_id}.\n\n" + announcement += f"We found {len(iam_drift_issues)} critical security alert(s) that require IMMEDIATE attention.\n" + announcement += f"\nPlease check the GitHub issue for detailed information and revoke unauthorized access immediately." + + self.sending_client.create_announcement(title, body, recipient, announcement) def print_announcement(self, recipient: str) -> None: """ Prints announcement details instead of sending them (for testing purposes). - + Args: recipient (str): The email address of the announcement recipient. """ if not self.sending_client: raise ValueError("SendingClient is required for printing announcements") - + diff = self.check_compliance() if not diff: self.logger.info("No compliance issues found, no announcement will be printed.") return - title = f"IAM Policy Non-Compliance Detected" - body = f"IAM policy for project {self.project_id} is not compliant with the defined policies on {self.users_file}\n\n" - for issue in diff: - body += f"- {issue}\n" + iam_drift_issues = [issue for issue in diff if IAC_DRIFT_IAM_USER in issue] + general_issues = [issue for issue in diff if IAC_DRIFT_IAM_USER not in issue] - announcement = f"Dear team,\n\nThis is an automated notification about compliance issues detected in the IAM policy for project {self.project_id}.\n\n" - announcement += f"We found {len(diff)} compliance issue(s) that need your attention.\n" - announcement += f"\nPlease check the GitHub issue for detailed information and take appropriate action to resolve these compliance violations." + if general_issues: + self.logger.info(f"Found {len(general_issues)} general IAM compliance issues. Printing announcement...") + title = f"IAM Policy Non-Compliance Detected" + body = f"IAM policy for project {self.project_id} is not compliant with the defined policies on {self.users_file}\n\n" + for issue in general_issues: + body += f"- {issue}\n" + + announcement = f"Dear team,\n\nThis is an automated notification about compliance issues detected in the IAM policy for project {self.project_id}.\n\n" + announcement += f"We found {len(general_issues)} compliance issue(s) that need your attention.\n" + announcement += f"\nPlease check the GitHub issue for detailed information and take appropriate action to resolve these compliance violations." + + self.sending_client.print_announcement(title, body, recipient, announcement) + + if iam_drift_issues: + self.logger.info("Printing security dashboard update for IAM vulnerabilities...") + title = f"[{IAC_DRIFT_IAM_USER}] Action Required: Unauthorized IAM Users Detected" + body = f"Critical security violations detected in IAM policies for project {self.project_id}:\n\n" + for issue in iam_drift_issues: + body += f"- {issue}\n" + + announcement = f"URGENT: Dear team,\n\nThis is an automated security alert regarding unauthorized IAM access in project {self.project_id}.\n\n" + announcement += f"We found {len(iam_drift_issues)} critical security alert(s) that require IMMEDIATE attention.\n" + announcement += f"\nPlease check the GitHub issue for detailed information and revoke unauthorized access immediately." + + self.sending_client.print_announcement(title, body, recipient, announcement) - self.sending_client.print_announcement(title, body, recipient, announcement) - def generate_compliance(self) -> None: """ Modifies the users file to match the current IAM policy. diff --git a/infra/enforcement/sending.py b/infra/enforcement/sending.py index bd9787b6ce87..9d24a816fbc3 100644 --- a/infra/enforcement/sending.py +++ b/infra/enforcement/sending.py @@ -18,6 +18,7 @@ import smtplib, ssl from typing import List, Optional from dataclasses import dataclass +from datetime import datetime, timezone @dataclass class GitHubIssue: @@ -229,7 +230,7 @@ def report_unmanaged_keys(self, project_id: str, compilance_issues: List[str]) - issue_title = "[SECURITY] Action Required: Unmanaged Service Account Keys Detected" #markdown body - timestamp = __import__("datetime").datetime.now(__import__("datetime").timezone.utc).strftime("%Y-%m-%d %H:%M:%S UTC") + timestamp = datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S UTC") new_report = f"### Unmanaged Keys Audit Report ({timestamp})\n" new_report += f"The following unauthorized or unmanaged keys were detected in `{project_id}`:\n\n" @@ -248,10 +249,13 @@ def report_unmanaged_keys(self, project_id: str, compilance_issues: List[str]) - if history_marker in old_body: # If history already exists, append the new report to it - headed = old_body.split(history_marker) + headed = old_body.split(history_marker, 1) last_report = headed[0].strip() old_history = headed[1].replace("", "").strip() + if old_history.endswith(""): + old_history = old_history[:-10].rstrip() + combined_history = f"{last_report}\n\n---\n\n{old_history}" else: combined_history = old_body.strip() @@ -294,18 +298,40 @@ def create_announcement(self, title: str, body: str, recipient: str, announcemen """ open_issues = self._get_open_issues(title) open_issues.sort(key=lambda x: x.updated_at, reverse=True) + + timestamp = datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S UTC") + new_report = f"### Compliance Audit Report ({timestamp})\n{body}" + if open_issues: - self.logger.info(f"Issue with title '{title}' already exists: #{open_issues[0].number}") - announcement += f"\n\nRelated GitHub Issue: {open_issues[0].html_url}" + target_issue = open_issues[0] + self.logger.info(f"Issue with title '{title}' already exists: #{target_issue.number}") + announcement += f"\n\nRelated GitHub Issue: {target_issue.html_url}" + + old_body = target_issue.body or "" + history_marker = "### History\n
\nClick to expand\n\n" - if open_issues[0].body != body: - self.logger.info(f"Updating body of issue #{open_issues[0].number}") - self.update_issue_body(open_issues[0].number, body) + if history_marker in old_body: + # If history already exists, append the new report to it + headed = old_body.split(history_marker, 1) + last_report = headed[0].strip() + old_history = headed[1].rstrip() + + if old_history.endswith("
"): + old_history = old_history[:-10].rstrip() + + combined_history = f"{last_report}\n\n---\n\n{old_history}" else: - self.logger.info(f"No changes detected for issue #{open_issues[0].number}") + # First time updating, turn the entire old body into history + combined_history = old_body.strip() + + final_body = f"{new_report}\n\n{history_marker}{combined_history}\n" + + self.logger.info(f"Appending report and archiving history to existing issue #{target_issue.number}") + self.update_issue_body(target_issue.number, final_body) self._send_email(title, announcement, recipient) else: - new_issue = self.create_issue(title, body) + self.logger.info(f"Creating new compliance issue for: {title}") + new_issue = self.create_issue(title, new_report) announcement += f"\n\nRelated GitHub Issue: {new_issue.html_url}" self._send_email(title, announcement, recipient) @@ -319,6 +345,13 @@ def print_announcement(self, title: str, body: str, recipient: str, announcement print(f"Recipient: {recipient}") print(f"Announcement: {announcement}") - print("\nSimulating GitHub issue creation...") - print(f"Title: {title}") - print(f"Body: {body}") + timestamp = datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S UTC") + + print("\n" + "="*60) + print("SIMULATING GITHUB GENERAL ISSUE CREATION/UPDATE") + print("="*60) + print(f"Title: {title}\n") + print(f"Body:\n### Compliance Audit Report ({timestamp})") + print(body) + print("### History\n
\nClick to expand\n\n[... Previous reports would be collapsed here ...]\n
") + print("="*60 + "\n") \ No newline at end of file diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java index 365ebbdc1f9d..7e9c3eca13f9 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java @@ -55,6 +55,7 @@ import org.apache.beam.runners.dataflow.worker.streaming.ExecutableWork; import org.apache.beam.runners.dataflow.worker.streaming.FailedWorkHandler; import org.apache.beam.runners.dataflow.worker.streaming.KeyCommitTooLargeException; +import org.apache.beam.runners.dataflow.worker.streaming.MultiKeyCommitValidationException; import org.apache.beam.runners.dataflow.worker.streaming.Watermarks; import org.apache.beam.runners.dataflow.worker.streaming.Work; import org.apache.beam.runners.dataflow.worker.streaming.config.StreamingGlobalConfig; @@ -193,7 +194,6 @@ public interface KeyTransitionListener { private @Nullable KeyTransitionListener keyTransitionListener; private @Nullable FailedWorkHandler onFailedWorkHandler; - private List executedWorks = Collections.emptyList(); private List outputBuilders = Collections.emptyList(); // Map> @@ -319,7 +319,6 @@ public byte[] getCurrentRecordOffset() { public void reset() { // these lists and maps are returned to callers after processing // don't clear and reuse, instead reset the reference. - this.executedWorks = Collections.emptyList(); this.outputBuilders = Collections.emptyList(); this.finalizationCallbacks = Collections.emptyMap(); // Work from prior bundles might have a reference to the old workBatchFailed. @@ -353,7 +352,6 @@ public void start( FailedWorkHandler onFailedWorkHandler) throws CoderException { reset(); - this.executedWorks = new ArrayList<>(); this.outputBuilders = new ArrayList<>(); this.finalizationCallbacks = new HashMap<>(); this.keyCoder = keyCoder; @@ -578,11 +576,13 @@ public void setActiveReader(UnboundedReader reader) { /** Invalidate the state and reader caches for this computation and key. */ public void invalidateCache() { - for (Work w : executedWorks) { - WindmillComputationKey compKey = - WindmillComputationKey.create(computationId, w.getShardedKey()); - readerCache.invalidateReader(compKey); - stateCache.invalidate(w.getShardedKey()); + if (budgetHandle != null) { + for (Work w : budgetHandle.getWorkBatch()) { + WindmillComputationKey compKey = + WindmillComputationKey.create(computationId, w.getShardedKey()); + readerCache.invalidateReader(compKey); + stateCache.invalidate(w.getShardedKey()); + } } if (activeReader != null) { try { @@ -718,6 +718,23 @@ private void validateCommitRequestSize() { return; } + // If this is a multi-key work item, then we need to retry all of the individual work items + // without merging so that we can identify large commits to truncate. + // TODO: Can we request truncation without retrying if the first commit exceed the limits? + BoundedQueueExecutorWorkHandle handle = checkNotNull(budgetHandle); + List currentBatch = handle.getWorkBatch(); + checkState(!currentBatch.isEmpty()); + if (currentBatch.size() > 1) { + LOG.warn( + "Windmill Commit limit exceeded on a multi key bundle. Retrying without batching. Batch size: {}", + currentBatch.size()); + for (Work w : currentBatch) { + w.setMultiKeyBatchingDisabled(true); + } + throw new MultiKeyCommitValidationException( + "Commit size validation failed for batch. Retrying individually."); + } + KeyCommitTooLargeException e = KeyCommitTooLargeException.causedBy( systemName, byteLimit, commitRequest, key, hotKeyLoggingEnabled); @@ -731,11 +748,6 @@ private void validateCommitRequestSize() { buildWorkItemTruncationRequestBuilder(currentWork, estimatedCommitSize); currentBuilder.clear(); currentBuilder.mergeFrom(truncationBuilder.build()); - - // TODO: throw and retry when truncation is not on a single key bundle. - checkState( - !multiKeyBundleOptions.multiKeyBundleEnabled(), - "Commit truncation not implemented for multikey bundles"); } private Windmill.WorkItemCommitRequest.Builder buildWorkItemTruncationRequestBuilder( @@ -774,7 +786,9 @@ public boolean advance() throws CoderException { throw new WorkItemCancelledException(activeWork.getWorkItem().getShardingKey()); } - if (activeWork.getKeyGroup().equals(Work.KeyGroup.DEFAULT) || shouldStopBatching()) { + if (activeWork.getKeyGroup().equals(Work.KeyGroup.DEFAULT) + || activeWork.isMultiKeyBatchingDisabled() + || shouldStopBatching()) { return false; } @@ -797,7 +811,6 @@ public boolean advance() throws CoderException { } private boolean shouldStopBatching() { - // TODO: stop batching if the previous work item requested truncation if (workItemsPolled >= multiKeyBundleOptions.maxKeyGroupBatchSize()) { return true; } @@ -821,7 +834,6 @@ private void startForNewKey(Work newWork) throws CoderException { this.outputBuilder = createOutputBuilder(newWork); this.outputBuilders.add(this.outputBuilder); newWork.setOnFailureListener(this.workBatchFailed); - this.executedWorks.add(newWork); logHotKeyIfDetected(newWork, this.key); @@ -862,11 +874,6 @@ public List getWorkItemCommits() { return commits; } - // Returns list of Work that was executed in the bundle - public List getExecutedWorks() { - return executedWorks; - } - // Returns finalization callbacks recorded during the bundle execution public Map> getFinalizationCallbacks() { return finalizationCallbacks; diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/BoundedQueueExecutorWorkHandle.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/BoundedQueueExecutorWorkHandle.java index 20661aae0a04..d7a61562bc58 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/BoundedQueueExecutorWorkHandle.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/BoundedQueueExecutorWorkHandle.java @@ -17,13 +17,15 @@ */ package org.apache.beam.runners.dataflow.worker.streaming; -import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; +import java.util.List; /** * A handle to use when requesting pulling more work from @BoundedQueueExecutor * via @BoundedQueueExecutor.pollWork */ public interface BoundedQueueExecutorWorkHandle { - // Returns all work that are tracked by the handle - ImmutableList getWorkBatch(); + // Returns all work that are tracked by the handle. + // Returned list cannot be modified. Copying the list is fine. + // Don't keep reference to the returned list after the processing exits the harness threads. + List getWorkBatch(); } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/ExecutableWork.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/ExecutableWork.java index 7748a554f0fc..4a992e872a4c 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/ExecutableWork.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/ExecutableWork.java @@ -82,4 +82,12 @@ public String getComputationId() { public Work.KeyGroup getKeyGroup() { return work().getKeyGroup(); } + + /** + * Returns true if multi-key batching is disabled for this work item (e.g. after a prior batch + * commit size validation failure). + */ + public boolean isMultiKeyBatchingDisabled() { + return work().isMultiKeyBatchingDisabled(); + } } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/MultiKeyCommitValidationException.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/MultiKeyCommitValidationException.java new file mode 100644 index 000000000000..f147d380d073 --- /dev/null +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/MultiKeyCommitValidationException.java @@ -0,0 +1,28 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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.apache.beam.runners.dataflow.worker.streaming; + +/** + * Thrown when a multi-key bundle exceeds commit size limits, triggering unbatching and local retry + * of individual work items. + */ +public final class MultiKeyCommitValidationException extends RuntimeException { + public MultiKeyCommitValidationException(String message) { + super(message); + } +} diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/Work.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/Work.java index 4541a1c313a2..2acee9410fa3 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/Work.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/Work.java @@ -83,6 +83,10 @@ public final class Work implements RefreshableWork { private final long serializedWorkItemSize; private volatile TimedState currentState; private volatile boolean isFailed; + // If true, this work item will not be batched with other work items in a multi-key bundle. + // This is used to isolate work items that failed validation (e.g. commit size limit exceeded) + // so they can be retried individually and potentially truncated. + private volatile boolean disableMultiKeyBatching = false; private volatile String processingThreadName = ""; private final AtomicReference<@Nullable AtomicBoolean> onFailureListener = new AtomicReference<>(null); @@ -399,6 +403,22 @@ public boolean isFailed() { return isFailed; } + /** + * Sets whether multi-key batching should be disabled for this work item. When true, this work + * item will not be batched with other work items upon local retry. + */ + public void setMultiKeyBatchingDisabled(boolean disableMultiKeyBatching) { + this.disableMultiKeyBatching = disableMultiKeyBatching; + } + + /** + * Returns true if multi-key batching is disabled for this work item (e.g. after a prior batch + * commit size validation failure). + */ + public boolean isMultiKeyBatchingDisabled() { + return disableMultiKeyBatching; + } + boolean isStuckCommittingAt(Instant stuckCommitDeadline) { return currentState.state() == Work.State.COMMITTING && currentState.startTime().isBefore(stuckCommitDeadline); diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/BoundedQueueExecutor.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/BoundedQueueExecutor.java index 2dd0f971168e..046d8cae9f9d 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/BoundedQueueExecutor.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/BoundedQueueExecutor.java @@ -22,6 +22,7 @@ import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument; import java.util.ArrayList; +import java.util.Collections; import java.util.List; import java.util.concurrent.ConcurrentLinkedQueue; import java.util.concurrent.LinkedBlockingQueue; @@ -36,7 +37,6 @@ import org.apache.beam.runners.dataflow.worker.streaming.Work; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions; -import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.Monitor; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.Monitor.Guard; import org.checkerframework.checker.nullness.qual.Nullable; @@ -306,8 +306,13 @@ public synchronized boolean isClosed() { } @Override - public synchronized ImmutableList getWorkBatch() { - return ImmutableList.copyOf(workBatch); + /* + * Returns an unmodifiable view over the underlying list. + * It is unsafe to use the returned list with concurrent calls to mutating methods + * like merge/close + */ + public synchronized List getWorkBatch() { + return Collections.unmodifiableList(workBatch); } @VisibleForTesting diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/KeyGroupWorkQueue.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/KeyGroupWorkQueue.java index d151157ec68f..dd409616ab9d 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/KeyGroupWorkQueue.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/KeyGroupWorkQueue.java @@ -20,6 +20,7 @@ import static org.apache.beam.sdk.util.Preconditions.checkArgumentNotNull; import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument; +import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState; import java.util.AbstractQueue; import java.util.Collection; @@ -67,9 +68,14 @@ static class Node { @Nullable Node prevKeyGroupNode; @Nullable Node nextKeyGroupNode; + private static boolean isMultiKeyBatchingDisabled(Runnable task) { + return !(task instanceof QueuedWork) + || ((QueuedWork) task).getWork().isMultiKeyBatchingDisabled(); + } + Node(Runnable task) { this.task = task; - if (task instanceof QueuedWork) { + if (!isMultiKeyBatchingDisabled(task)) { this.computationId = ((QueuedWork) task).getWork().getComputationId(); this.keyGroup = ((QueuedWork) task).getWork().getKeyGroup(); } else { @@ -193,6 +199,10 @@ private void unlinkNode(Node node) { if (firstNode == keyGroupWorkList.tail) { return null; } + + // MultiKeyBatchingDisabled items should not be in keyGroupWorkList + checkState(!Node.isMultiKeyBatchingDisabled(firstNode.task)); + unlinkNode(firstNode); return (QueuedWork) firstNode.task; diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/StreamingWorkScheduler.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/StreamingWorkScheduler.java index 63cfad5a9a6f..958cd62f5eb3 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/StreamingWorkScheduler.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/StreamingWorkScheduler.java @@ -242,7 +242,6 @@ private void processWork( long processingStartTimeNanos = System.nanoTime(); StageInfo stageInfo = getStageInfo(computationState); - @Nullable List workBatch = null; try { if (work.isFailed()) { throw new WorkItemCancelledException(workItem.getShardingKey()); @@ -251,7 +250,7 @@ private void processWork( // Execute the user code for the Work batch. ExecuteWorkResult executeWorkResult = executeWork(work, stageInfo, computationState, handle, keyTransitionListener); - workBatch = executeWorkResult.workBatch(); + List workBatch = handle.getWorkBatch(); List workItemCommits = executeWorkResult.workItemCommits(); commitFinalizer.cacheCommitFinalizers(executeWorkResult.finalizationCallbacks()); @@ -264,7 +263,7 @@ private void processWork( handleProcessWorkFailure( computationState, handle.getWorkBatch(), computationId, systemName, work, t); } finally { - List processedWorkBatch = workBatch != null ? workBatch : ImmutableList.of(work); + List processedWorkBatch = handle.getWorkBatch(); // Update total processing time counters. Updating in finally clause ensures that // work items causing exceptions are also accounted in time spent. recordProcessingTime(stageInfo, processedWorkBatch, processingStartTimeNanos); @@ -328,7 +327,6 @@ private ExecuteWorkResult executeWork( computationWorkExecutor.executeWork( work, workExecutor, handle, keyTransitionListener, onFailedWorkHandler); - List workBatch; List workItemCommits; Map> finalizationCallbacks; long stateBytesRead; @@ -338,9 +336,6 @@ private ExecuteWorkResult executeWork( } context.flushState(); - // Retrieve executed works, work item commits, and accumulated callbacks from execution - // context - workBatch = context.getExecutedWorks(); workItemCommits = context.getWorkItemCommits(); finalizationCallbacks = context.getFinalizationCallbacks(); stateBytesRead = context.getStateBytesRead(); @@ -351,8 +346,7 @@ private ExecuteWorkResult executeWork( computationState.releaseComputationWorkExecutor(computationWorkExecutor); computationWorkExecutor = null; - return ExecuteWorkResult.create( - workBatch, workItemCommits, finalizationCallbacks, stateBytesRead); + return ExecuteWorkResult.create(workItemCommits, finalizationCallbacks, stateBytesRead); } catch (Throwable t) { if (computationWorkExecutor != null) { // If processing failed due to a thrown exception, close the executionState. Do not @@ -419,10 +413,6 @@ private void commitMultiKeyWorkBatch( } for (int i = 0; i < workBatch.size(); i++) { Windmill.WorkItemCommitRequest commit = workItemCommits.get(i); - // TODO: Retry on commit truncations - checkState( - !commit.getExceedsMaxWorkItemCommitBytes(), - "Commit truncation with multikey bundles not implemented"); Work w = workBatch.get(i); multiKeyBuilder.addRequests( commit @@ -523,16 +513,13 @@ private KeyTransitionListener createKeyTransitionListener() { @AutoValue abstract static class ExecuteWorkResult { static ExecuteWorkResult create( - List workBatch, List workItemCommits, Map> finalizationCallbacks, long stateBytesRead) { return new AutoValue_StreamingWorkScheduler_ExecuteWorkResult( - workBatch, workItemCommits, finalizationCallbacks, stateBytesRead); + workItemCommits, finalizationCallbacks, stateBytesRead); } - abstract List workBatch(); - abstract List workItemCommits(); // Map> diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/failures/WorkFailureProcessor.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/failures/WorkFailureProcessor.java index d23c870178e0..b635bde7e08a 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/failures/WorkFailureProcessor.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/failures/WorkFailureProcessor.java @@ -25,6 +25,7 @@ import org.apache.beam.runners.dataflow.worker.status.LastExceptionDataProvider; import org.apache.beam.runners.dataflow.worker.streaming.ExecutableWork; import org.apache.beam.runners.dataflow.worker.streaming.FailedWorkHandler; +import org.apache.beam.runners.dataflow.worker.streaming.MultiKeyCommitValidationException; import org.apache.beam.runners.dataflow.worker.streaming.Work; import org.apache.beam.runners.dataflow.worker.util.BoundedQueueExecutor; import org.apache.beam.sdk.annotations.Internal; @@ -162,6 +163,16 @@ private RetryEvaluation evaluateRetry( @Nullable final Throwable cause = t.getCause(); Throwable parsedException = (t instanceof UserCodeException && cause != null) ? cause : t; + if (parsedException instanceof MultiKeyCommitValidationException) { + LOG.info( + "Execution of work for computation '{}' on sharding key '{}' for work token '{}' exceeded commit size limits. " + + "Work will be retried locally in smaller batches.", + computationId, + work.getWorkItem().getShardingKey(), + work.getWorkItem().getWorkToken()); + return RetryEvaluation.RETRY_LOCALLY; + } + LastExceptionDataProvider.reportException(parsedException); LOG.debug("Failed work: {}", work); Duration elapsedTimeSinceStart = new Duration(work.getStartTime(), clock.get()); diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorkerTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorkerTest.java index c48b30ecf640..d8063ae66d44 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorkerTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorkerTest.java @@ -151,9 +151,11 @@ import org.apache.beam.sdk.coders.VarIntCoder; import org.apache.beam.sdk.extensions.gcp.util.Transport; import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.apache.beam.sdk.state.BagState; import org.apache.beam.sdk.state.StateSpec; import org.apache.beam.sdk.state.StateSpecs; import org.apache.beam.sdk.state.ValueState; +import org.apache.beam.sdk.testing.ExpectedLogs; import org.apache.beam.sdk.transforms.DoFn; import org.apache.beam.sdk.transforms.DoFnSchemaInformation; import org.apache.beam.sdk.transforms.windowing.AfterPane; @@ -301,6 +303,11 @@ public Long get() { }; @Rule public transient Timeout globalTimeout = Timeout.seconds(600); + + @Rule + public ExpectedLogs expectedStreamingModeExecutionContextLogs = + ExpectedLogs.none(StreamingModeExecutionContext.class); + @Rule public BlockingFn blockingFn = new BlockingFn(); @Rule public TestRule restoreMDC = new RestoreDataflowLoggingMDC(); @Rule public final GrpcCleanupRule grpcCleanup = new GrpcCleanupRule(); @@ -346,6 +353,8 @@ private Iterable buildCounters() { @Before public void setUp() { + FixedSizeBagCommitFn.SEEN_ELEMENTS.set(0); + LargeBagCommitFn.SEEN_ELEMENTS.set(0); server.clearCommitsReceived(); streamingCounters = StreamingCounters.create(); } @@ -4884,6 +4893,479 @@ public void testSkipInputElementsWithDecodingExceptions() throws Exception { "12345", commit.getOutputMessages(0).getBundles(0).getMessages(0).getData().toStringUtf8()); } + // TODO: Add similar tests with productions after changing WindmillSink to flush in finishKey. + + @Test + public void testMultiKeyCommit_batchCommitSizeExceededUnBatchSucceeds() throws Exception { + if (!streamingEngine) { + return; + } + KvCoder kvCoder = KvCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of()); + + List instructions = + Arrays.asList( + makeSourceInstruction(kvCoder), + makeDoFnInstruction(new FixedSizeBagCommitFn(500), 0, kvCoder), + makeSinkInstruction(kvCoder, 1)); + + StreamingDataflowWorker worker = + makeWorker( + defaultWorkerParams( + "--experiments=unstable_enable_multi_key_bundle,windmill_max_key_group_batch_time_ms=50000", + "--numberOfWorkerHarnessThreads=1") + .setLocalRetryTimeoutMs(100) + .setInstructions(instructions) + .setStreamingGlobalConfig( + StreamingGlobalConfig.builder() + .setOperationalLimits( + OperationalLimits.builder().setMaxWorkItemCommitBytes(1000).build()) + .build()) + .build()); + worker.start(); + + String batchInputText = + "work {" + + " computation_id: \"" + + DEFAULT_COMPUTATION_ID + + "\"" + + " input_data_watermark: 0" + + " work {" + + " key: \"key1\"" + + " sharding_key: 1" + + " work_token: 1" + + " cache_token: 1" + + " key_group { high: 0 low: 1 }" + + " message_bundles {" + + " source_computation_id: \"" + + DEFAULT_SOURCE_COMPUTATION_ID + + "\"" + + " messages {" + + " timestamp: 0" + + " data: \"data1\"" + + " }" + + " }" + + " }" + + " work {" + + " key: \"key2\"" + + " sharding_key: 2" + + " work_token: 2" + + " cache_token: 2" + + " key_group { high: 0 low: 1 }" + + " message_bundles {" + + " source_computation_id: \"" + + DEFAULT_SOURCE_COMPUTATION_ID + + "\"" + + " messages {" + + " timestamp: 0" + + " data: \"data2\"" + + " }" + + " }" + + " }" + + " work {" + + " key: \"key3\"" + + " sharding_key: 3" + + " work_token: 3" + + " cache_token: 3" + + " key_group { high: 0 low: 1 }" + + " message_bundles {" + + " source_computation_id: \"" + + DEFAULT_SOURCE_COMPUTATION_ID + + "\"" + + " messages {" + + " timestamp: 0" + + " data: \"data3\"" + + " }" + + " }" + + " }" + + "}"; + Windmill.GetWorkResponse batchInput = + buildInput( + batchInputText, + CoderUtils.encodeToByteArray( + CollectionCoder.of(IntervalWindow.getCoder()), + Collections.singletonList(DEFAULT_WINDOW))); + + server.whenGetWorkCalled().thenReturn(batchInput); + + Map result = server.waitForAndGetCommits(3); + + assertEquals(3, result.size()); + assertTrue(result.containsKey(1L)); + assertTrue(result.containsKey(2L)); + assertTrue(result.containsKey(3L)); + for (Windmill.WorkItemCommitRequest commitRequest : result.values()) { + assertFalse(commitRequest.getExceedsMaxWorkItemCommitBytes()); + } + + List multiKeyCommits = + server.getMultiKeyCommitsReceived(); + assertEquals(3, multiKeyCommits.size()); + assertEquals(1, multiKeyCommits.get(0).getRequestsCount()); + assertEquals(1, multiKeyCommits.get(1).getRequestsCount()); + assertEquals(1, multiKeyCommits.get(2).getRequestsCount()); + // 2 in initial batch (item 1 succeeds with 500 bytes, item 2 fails after accumulating 1000 + // bytes) + 3 unbatched retries + assertEquals(5, FixedSizeBagCommitFn.SEEN_ELEMENTS.get()); + expectedStreamingModeExecutionContextLogs.verifyWarn( + "Windmill Commit limit exceeded on a multi key bundle"); + + worker.stop(); + } + + @Test + public void testMultiKeyCommit_batchCommitSizeExceededUnBatchTruncates() throws Exception { + if (!streamingEngine) { + return; + } + KvCoder kvCoder = KvCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of()); + + List instructions = + Arrays.asList( + makeSourceInstruction(kvCoder), + makeDoFnInstruction(new FixedSizeBagCommitFn(500), 0, kvCoder), + makeSinkInstruction(kvCoder, 1)); + + StreamingDataflowWorker worker = + makeWorker( + defaultWorkerParams( + "--experiments=unstable_enable_multi_key_bundle,windmill_max_key_group_batch_time_ms=50000", + "--numberOfWorkerHarnessThreads=1") + .setLocalRetryTimeoutMs(100) + .setInstructions(instructions) + .setStreamingGlobalConfig( + StreamingGlobalConfig.builder() + .setOperationalLimits( + // All workitems exceed commit limits + OperationalLimits.builder().setMaxWorkItemCommitBytes(400).build()) + .build()) + .build()); + worker.start(); + + String batchInputText = + "work {" + + " computation_id: \"" + + DEFAULT_COMPUTATION_ID + + "\"" + + " input_data_watermark: 0" + + " work {" + + " key: \"key1\"" + + " sharding_key: 1" + + " work_token: 1" + + " cache_token: 1" + + " key_group { high: 0 low: 1 }" + + " message_bundles {" + + " source_computation_id: \"" + + DEFAULT_SOURCE_COMPUTATION_ID + + "\"" + + " messages {" + + " timestamp: 0" + + " data: \"data1\"" + + " }" + + " }" + + " }" + + " work {" + + " key: \"key2\"" + + " sharding_key: 2" + + " work_token: 2" + + " cache_token: 2" + + " key_group { high: 0 low: 1 }" + + " message_bundles {" + + " source_computation_id: \"" + + DEFAULT_SOURCE_COMPUTATION_ID + + "\"" + + " messages {" + + " timestamp: 0" + + " data: \"data2\"" + + " }" + + " }" + + " }" + + " work {" + + " key: \"key3\"" + + " sharding_key: 3" + + " work_token: 3" + + " cache_token: 3" + + " key_group { high: 0 low: 1 }" + + " message_bundles {" + + " source_computation_id: \"" + + DEFAULT_SOURCE_COMPUTATION_ID + + "\"" + + " messages {" + + " timestamp: 0" + + " data: \"data3\"" + + " }" + + " }" + + " }" + + "}"; + Windmill.GetWorkResponse batchInput = + buildInput( + batchInputText, + CoderUtils.encodeToByteArray( + CollectionCoder.of(IntervalWindow.getCoder()), + Collections.singletonList(DEFAULT_WINDOW))); + + server.whenGetWorkCalled().thenReturn(batchInput); + + Map result = server.waitForAndGetCommits(3); + + assertEquals(3, result.size()); + assertTrue(result.containsKey(1L)); + assertTrue(result.containsKey(2L)); + assertTrue(result.containsKey(3L)); + for (Windmill.WorkItemCommitRequest commitRequest : result.values()) { + assertTrue(commitRequest.getExceedsMaxWorkItemCommitBytes()); + } + + List multiKeyCommits = + server.getMultiKeyCommitsReceived(); + assertEquals(3, multiKeyCommits.size()); + assertEquals(1, multiKeyCommits.get(0).getRequestsCount()); + assertEquals(1, multiKeyCommits.get(1).getRequestsCount()); + assertEquals(1, multiKeyCommits.get(2).getRequestsCount()); + // 1 in initial batch (fails after first item's bag write exceeds limit) + 3 unbatched retries + assertEquals(4, FixedSizeBagCommitFn.SEEN_ELEMENTS.get()); + expectedStreamingModeExecutionContextLogs.verifyWarn( + "Windmill Commit limit exceeded on a multi key bundle"); + + worker.stop(); + } + + @Test + public void testMultiKeyCommit_batchCommitSizeExceededUnBatchFirstItemTruncates() + throws Exception { + if (!streamingEngine) { + return; + } + KvCoder kvCoder = KvCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of()); + + List instructions = + Arrays.asList( + makeSourceInstruction(kvCoder), + makeDoFnInstruction(new LargeBagCommitFn(), 0, kvCoder), + makeSinkInstruction(kvCoder, 1)); + + StreamingDataflowWorker worker = + makeWorker( + defaultWorkerParams( + "--experiments=unstable_enable_multi_key_bundle,windmill_max_key_group_batch_time_ms=50000", + "--numberOfWorkerHarnessThreads=1") + .setLocalRetryTimeoutMs(100) + .setInstructions(instructions) + .setStreamingGlobalConfig( + StreamingGlobalConfig.builder() + .setOperationalLimits( + OperationalLimits.builder().setMaxWorkItemCommitBytes(1000).build()) + .build()) + .build()); + worker.start(); + + String batchInputText = + "work {" + + " computation_id: \"" + + DEFAULT_COMPUTATION_ID + + "\"" + + " input_data_watermark: 0" + + " work {" + + " key: \"large_key\"" + + " sharding_key: 1" + + " work_token: 1" + + " cache_token: 1" + + " key_group { high: 0 low: 1 }" + + " message_bundles {" + + " source_computation_id: \"" + + DEFAULT_SOURCE_COMPUTATION_ID + + "\"" + + " messages {" + + " timestamp: 0" + + " data: \"data1\"" + + " }" + + " }" + + " }" + + " work {" + + " key: \"small_key\"" + + " sharding_key: 2" + + " work_token: 2" + + " cache_token: 2" + + " key_group { high: 0 low: 1 }" + + " message_bundles {" + + " source_computation_id: \"" + + DEFAULT_SOURCE_COMPUTATION_ID + + "\"" + + " messages {" + + " timestamp: 0" + + " data: \"data2\"" + + " }" + + " }" + + " }" + + " work {" + + " key: \"small_key\"" + + " sharding_key: 3" + + " work_token: 3" + + " cache_token: 3" + + " key_group { high: 0 low: 1 }" + + " message_bundles {" + + " source_computation_id: \"" + + DEFAULT_SOURCE_COMPUTATION_ID + + "\"" + + " messages {" + + " timestamp: 0" + + " data: \"data3\"" + + " }" + + " }" + + " }" + + "}"; + Windmill.GetWorkResponse batchInput = + buildInput( + batchInputText, + CoderUtils.encodeToByteArray( + CollectionCoder.of(IntervalWindow.getCoder()), + Collections.singletonList(DEFAULT_WINDOW))); + + server.whenGetWorkCalled().thenReturn(batchInput); + + Map result = server.waitForAndGetCommits(3); + + assertEquals(3, result.size()); + assertTrue(result.containsKey(1L)); + assertTrue(result.containsKey(2L)); + assertTrue(result.containsKey(3L)); + assertTrue(result.get(1L).getExceedsMaxWorkItemCommitBytes()); + assertFalse(result.get(2L).getExceedsMaxWorkItemCommitBytes()); + assertFalse(result.get(3L).getExceedsMaxWorkItemCommitBytes()); + + List multiKeyCommits = + server.getMultiKeyCommitsReceived(); + assertEquals(3, multiKeyCommits.size()); + assertEquals(1, multiKeyCommits.get(0).getRequestsCount()); + assertEquals(1, multiKeyCommits.get(1).getRequestsCount()); + assertEquals(1, multiKeyCommits.get(2).getRequestsCount()); + // 1 in initial batch (fails after first item's bag write exceeds limit) + 3 unbatched retries + assertEquals(4, LargeBagCommitFn.SEEN_ELEMENTS.get()); + + expectedStreamingModeExecutionContextLogs.verifyWarn( + "Windmill Commit limit exceeded on a multi key bundle"); + + worker.stop(); + } + + @Test + public void testMultiKeyCommit_batchCommitSizeExceededUnBatchSecondItemTruncates() + throws Exception { + if (!streamingEngine) { + return; + } + KvCoder kvCoder = KvCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of()); + + List instructions = + Arrays.asList( + makeSourceInstruction(kvCoder), + makeDoFnInstruction(new LargeBagCommitFn(), 0, kvCoder), + makeSinkInstruction(kvCoder, 1)); + + StreamingDataflowWorker worker = + makeWorker( + defaultWorkerParams( + "--experiments=unstable_enable_multi_key_bundle,windmill_max_key_group_batch_time_ms=50000", + "--numberOfWorkerHarnessThreads=1") + .setLocalRetryTimeoutMs(100) + .setInstructions(instructions) + .setStreamingGlobalConfig( + StreamingGlobalConfig.builder() + .setOperationalLimits( + OperationalLimits.builder().setMaxWorkItemCommitBytes(1000).build()) + .build()) + .build()); + worker.start(); + + String batchInputText = + "work {" + + " computation_id: \"" + + DEFAULT_COMPUTATION_ID + + "\"" + + " input_data_watermark: 0" + + " work {" + + " key: \"small_key\"" + + " sharding_key: 1" + + " work_token: 1" + + " cache_token: 1" + + " key_group { high: 0 low: 1 }" + + " message_bundles {" + + " source_computation_id: \"" + + DEFAULT_SOURCE_COMPUTATION_ID + + "\"" + + " messages {" + + " timestamp: 0" + + " data: \"data1\"" + + " }" + + " }" + + " }" + + " work {" + + " key: \"large_key\"" + + " sharding_key: 2" + + " work_token: 2" + + " cache_token: 2" + + " key_group { high: 0 low: 1 }" + + " message_bundles {" + + " source_computation_id: \"" + + DEFAULT_SOURCE_COMPUTATION_ID + + "\"" + + " messages {" + + " timestamp: 0" + + " data: \"data2\"" + + " }" + + " }" + + " }" + + " work {" + + " key: \"small_key\"" + + " sharding_key: 3" + + " work_token: 3" + + " cache_token: 3" + + " key_group { high: 0 low: 1 }" + + " message_bundles {" + + " source_computation_id: \"" + + DEFAULT_SOURCE_COMPUTATION_ID + + "\"" + + " messages {" + + " timestamp: 0" + + " data: \"data3\"" + + " }" + + " }" + + " }" + + "}"; + Windmill.GetWorkResponse batchInput = + buildInput( + batchInputText, + CoderUtils.encodeToByteArray( + CollectionCoder.of(IntervalWindow.getCoder()), + Collections.singletonList(DEFAULT_WINDOW))); + + server.whenGetWorkCalled().thenReturn(batchInput); + + Map result = server.waitForAndGetCommits(3); + + assertEquals(3, result.size()); + assertTrue(result.containsKey(1L)); + assertTrue(result.containsKey(2L)); + assertTrue(result.containsKey(3L)); + assertFalse(result.get(1L).getExceedsMaxWorkItemCommitBytes()); + assertTrue(result.get(2L).getExceedsMaxWorkItemCommitBytes()); + assertFalse(result.get(3L).getExceedsMaxWorkItemCommitBytes()); + + List multiKeyCommits = + server.getMultiKeyCommitsReceived(); + assertEquals(3, multiKeyCommits.size()); + assertEquals(1, multiKeyCommits.get(0).getRequestsCount()); + assertEquals(1, multiKeyCommits.get(1).getRequestsCount()); + assertEquals(1, multiKeyCommits.get(2).getRequestsCount()); + // 2 in initial batch (item 1 succeeds, fails after item 2's bag write exceeds limit) + 3 + // unbatched retries + assertEquals(5, LargeBagCommitFn.SEEN_ELEMENTS.get()); + + expectedStreamingModeExecutionContextLogs.verifyWarn( + "Windmill Commit limit exceeded on a multi key bundle"); + + worker.stop(); + } + static class BlockingFn extends DoFn implements TestRule { public static AtomicReference blocker = @@ -4978,7 +5460,6 @@ public static void reset() { } static class LargeCommitFn extends DoFn, KV> { - @ProcessElement public void processElement(ProcessContext c) { if (c.element().getKey().equals("large_key")) { @@ -4993,6 +5474,58 @@ public void processElement(ProcessContext c) { } } + static class LargeBagCommitFn extends DoFn, KV> { + @StateId("bag") + private final StateSpec> bagSpec = StateSpecs.bag(StringUtf8Coder.of()); + + public static AtomicInteger SEEN_ELEMENTS = new AtomicInteger(); + + @ProcessElement + public void processElement(ProcessContext c, @StateId("bag") BagState bag) { + SEEN_ELEMENTS.incrementAndGet(); + if (c.element().getKey().equals("large_key")) { + StringBuilder s = new StringBuilder(); + for (int i = 0; i < 100; ++i) { + s.append("large_commit"); + } + bag.add(s.toString()); + } else { + bag.add(c.element().getValue()); + } + } + } + + static class FixedSizeBagCommitFn extends DoFn, KV> { + @StateId("bag") + private final StateSpec> bagSpec = StateSpecs.bag(StringUtf8Coder.of()); + + private final int size; + public static AtomicInteger SEEN_ELEMENTS = new AtomicInteger(); + private List bundleElements = new ArrayList<>(); + + FixedSizeBagCommitFn(int size) { + this.size = size; + } + + @StartBundle + public void startBundle() { + bundleElements = new ArrayList<>(); + } + + @ProcessElement + public void processElement(ProcessContext c, @StateId("bag") BagState bag) { + SEEN_ELEMENTS.incrementAndGet(); + StringBuilder s = new StringBuilder(); + for (int i = 0; i < size; ++i) { + s.append("a"); + } + bundleElements.add(s.toString()); + for (String elem : bundleElements) { + bag.add(elem); + } + } + } + static class ExceptionCatchingFn extends DoFn, KV> { @ProcessElement diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContextTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContextTest.java index 5aceb0ca9564..53dd96620a55 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContextTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContextTest.java @@ -78,7 +78,7 @@ import org.apache.beam.runners.dataflow.worker.windmill.state.WindmillStateCache; import org.apache.beam.runners.dataflow.worker.windmill.state.WindmillTagEncodingV1; import org.apache.beam.runners.dataflow.worker.windmill.state.WindmillTagEncodingV2; -import org.apache.beam.runners.dataflow.worker.windmill.work.processing.failures.FailureTracker; +import org.apache.beam.runners.dataflow.worker.windmill.work.processing.failures.StreamingEngineFailureTracker; import org.apache.beam.runners.dataflow.worker.windmill.work.refresh.HeartbeatSender; import org.apache.beam.sdk.Pipeline; import org.apache.beam.sdk.coders.Coder; @@ -164,7 +164,7 @@ private StreamingModeExecutionContext createExecutionContext( /*stepName=*/ "stepName", /*systemName=*/ "systemName", StreamingCounters.create(), - mock(FailureTracker.class), + StreamingEngineFailureTracker.create(10, 10), "sourceBytesProcessCounterName", MultiKeyBundleOptions.fromOptions(options), SideInputStateFetcherFactory.fromOptions(options)); @@ -900,4 +900,44 @@ public void testInternalsPoisonedAfterFlushState() throws Exception { assertThat(e.getMessage(), Matchers.containsString("poisoned")); } } + + @Test + public void testAdvance_stopsWhenCurrentWorkBatchingDisabled() throws Exception { + DataflowWorkerHarnessOptions optionsMultiKey = + PipelineOptionsFactory.as(DataflowWorkerHarnessOptions.class); + optionsMultiKey + .as(ExperimentalOptions.class) + .setExperiments(Arrays.asList("unstable_enable_multi_key_bundle")); + StreamingModeExecutionContext context = + createExecutionContext(optionsMultiKey, globalConfigHandle); + + BoundedQueueExecutor mockExecutor = mock(BoundedQueueExecutor.class); + BoundedQueueExecutorWorkHandle mockHandle = mock(BoundedQueueExecutorWorkHandle.class); + Windmill.Uint128Proto keyGroup = + Windmill.Uint128Proto.newBuilder().setHigh(1).setLow(2).build(); + + Work work1 = + createMockWork( + Windmill.WorkItem.newBuilder() + .setKey(ByteString.copyFromUtf8("key1")) + .setWorkToken(1L) + .setKeyGroup(keyGroup) + .build(), + Watermarks.builder().setInputDataWatermark(Instant.EPOCH).build()); + work1.setMultiKeyBatchingDisabled(true); + + AtomicBoolean transitionListenerCalled = new AtomicBoolean(false); + context.start( + work1, + workExecutor, + mockExecutor, + mockHandle, + null, + (oldWork, newWork) -> transitionListenerCalled.set(true), + FAILING_FAILED_WORK_HANDLER); + + assertFalse(context.advance()); + assertFalse(transitionListenerCalled.get()); + verifyNoInteractions(mockExecutor); + } } diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/KeyGroupWorkQueueTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/KeyGroupWorkQueueTest.java index 77fcb0597586..c7be44525502 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/KeyGroupWorkQueueTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/util/KeyGroupWorkQueueTest.java @@ -489,6 +489,27 @@ public void testPollWorkWithKeyGroup() { assertTrue(queue.isEmpty()); } + @Test + public void testOffer_multiKeyBatchingDisabled_notInsertedInKeyGroupQueue() { + KeyGroupWorkQueue queue = new KeyGroupWorkQueue(fairQueue); + QueuedWork workDisabled = createQueuedWork("compA", 100); + workDisabled.getWork().work().setMultiKeyBatchingDisabled(true); + QueuedWork workEnabled = createQueuedWork("compA", 200); + + queue.offer(workDisabled); + queue.offer(workEnabled); + assertEquals(2, queue.size()); + + QueuedWork polledWork = queue.pollWork("compA", TEST_KEY_GROUP); + assertNotNull(polledWork); + assertEquals(workEnabled, polledWork); + assertEquals(1, queue.size()); + + assertNull(queue.pollWork("compA", TEST_KEY_GROUP)); + assertEquals(workDisabled, queue.poll()); + assertTrue(queue.isEmpty()); + } + private void waitForThreadState(Thread t, State state) throws InterruptedException { long timeoutMs = 30000; long start = System.currentTimeMillis(); diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/failures/WorkFailureProcessorTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/failures/WorkFailureProcessorTest.java index 741cc35376fc..f1cc33c963f1 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/failures/WorkFailureProcessorTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/failures/WorkFailureProcessorTest.java @@ -30,6 +30,8 @@ import java.util.function.Consumer; import java.util.function.Supplier; import org.apache.beam.runners.dataflow.worker.streaming.ExecutableWork; +import org.apache.beam.runners.dataflow.worker.streaming.FailedWorkHandler; +import org.apache.beam.runners.dataflow.worker.streaming.MultiKeyCommitValidationException; import org.apache.beam.runners.dataflow.worker.streaming.Watermarks; import org.apache.beam.runners.dataflow.worker.streaming.Work; import org.apache.beam.runners.dataflow.worker.util.BoundedQueueExecutor; @@ -272,4 +274,25 @@ public void logAndProcessFailureBatch_mixRetryAndAbort() throws Throwable { assertThat(executedWork2).isEmpty(); assertThat(invalidWork).containsExactly(work2.work()); } + + @Test + public void logAndProcessFailureBatch_retriesOnMultiKeyCommitValidationException() + throws Throwable { + CountDownLatch runWork = new CountDownLatch(1); + ExecutableWork work = createWork(ignored -> runWork.countDown()); + FailureTracker failureTracker = streamingEngineFailureReporter(); + WorkFailureProcessor workFailureProcessor = createWorkFailureProcessor(failureTracker); + Set invalidWork = new HashSet<>(); + + workFailureProcessor.logAndProcessFailureBatch( + DEFAULT_COMPUTATION_ID, + DEFAULT_COMPUTATION_ID, + List.of(work), + new MultiKeyCommitValidationException("test"), + (FailedWorkHandler) invalidWork::add); + + runWork.await(); + assertThat(invalidWork).isEmpty(); + assertThat(failureTracker.drainPendingFailuresToReport()).isEmpty(); + } }