From 953da9b0c00a9cd77de8b76eb2eaadae3f3309f2 Mon Sep 17 00:00:00 2001 From: Pearson White Date: Sun, 30 Nov 2025 16:34:59 -0500 Subject: [PATCH] Refactor to use configs and server - Use configs to define actions - Run a server to allow running commands interactively --- Dockerfile | 31 +- README.md | 82 ++++-- api_requester.py | 327 +++++++++++++++++++++ bind.yaml | 24 ++ build.sh | 10 +- client.py | 136 +++++++++ config.yaml | 62 ++++ configs.py | 222 ++++++++++++++ deployment.yaml | 43 +++ publisher-service.yaml | 14 + requirements.txt | 7 + store_msg_retriever.py | 639 ----------------------------------------- utils.py | 303 +++++++++++++++++++ 13 files changed, 1230 insertions(+), 670 deletions(-) create mode 100644 api_requester.py create mode 100644 bind.yaml create mode 100644 client.py create mode 100644 config.yaml create mode 100644 configs.py create mode 100644 deployment.yaml create mode 100644 publisher-service.yaml create mode 100644 requirements.txt delete mode 100644 store_msg_retriever.py create mode 100644 utils.py diff --git a/Dockerfile b/Dockerfile index cd47c3d..1428ed9 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,5 +1,30 @@ -FROM python:3.11.9-alpine +FROM python:3.11.9-alpine AS base +WORKDIR /app -ADD store_msg_retriever.py /app/store_msg_retriever.py +COPY requirements.txt . +RUN pip install --no-cache-dir --break-system-packages -r requirements.txt +COPY api_requester.py utils.py configs.py /app/ -RUN pip install requests pydantic +FROM base AS debug +WORKDIR /app +RUN apk add --no-cache \ + bash \ + bind-tools \ + curl \ + ethtool \ + iputils \ + jq \ + net-tools \ + tcpdump \ + vim \ + wget \ + ws \ + nodejs \ + npm \ + && npm install -g wscat + +ENTRYPOINT ["sleep", "infinity"] + +FROM base AS production +WORKDIR /app +ENTRYPOINT ["python", "./api_requester.py", "--mode", "server", "--config", "/mount/config.yaml"] \ No newline at end of file diff --git a/README.md b/README.md index 5abe0db..0a734a3 100644 --- a/README.md +++ b/README.md @@ -1,40 +1,70 @@ ## Waku Storage Retriever -This Python script retrieves messages -from a Waku storage service using an HTTP API. -It supports pagination and resolves DNS for the service host. +This Python script facilitates arbitrary GET and POST requests +to pods in a Kubernetes cluster. The script is designed to run inside a Docker container. ### Usage -Run the script with: +The script can be run in one of two modes: `server` or `batch`. + +#### Batch Mode +Simply runs all actions sequentially + +#### Server Mode + +`python ./api_requester.py --mode server --config /mount/config.yaml` + +Runs a server, allowing scripts to call API endpoints, causing this pod to make API requests to other pods. + +See endpoints under `def create_app` in `api_requester.py` for usage details. + +### Config Format + +The ConfigMap in `config.yaml` defines the config objects. +Class definitions are in `configs.py`. + +The idea is that a user can define various pieces of the config, +then combine them as needed. Each object stands independently, but +they work together when running an action or making a request. + +Each config object has a name by which it can be referenced. +Some fields are optional. + +Endpoints - Defines an API endpoint for a request. +Targets - Defines a set of filters to use to determine if pods on a cluster are part of the target. +Requests - Contains an Endpoint and some additional information for retries and delays. +Actions - Combines Targets and Requests into a defined action, representing a series of requests. + +#### How an Action is performed + +1. For each ConfigTarget, add all pods to the list "all_pods". Note: No deduplication is done. +2. Sort the list of pods according to `order`. +3. Starting at `pod_start_index`, take `pod_count` pods. Note: Loop through the list if needed to get `pod_count` items. +4. According to `loop_order`, make every request in `requests` to every pod in the remaining list. + +See docstrings in the `ConfigAction` class for more details. + +### Files + ``` -python script.py [-c CONTENT_TOPIC] [-p PUBSUB_TOPIC] [-ps PAGE_SIZE] [-cs CURSOR] +api_requester.py Main code that will run on the pod-api-requester pod +bind.yaml Necessary Kubernetes Role + RoleBinding to give permission to list pods +config.yaml Kubernetes config containing the definitions for Targets, Endpoints, Requests, and Actions + +build.sh Sample commands to build the Docker container +Dockerfile File to build the pod-api-requester image + +deployment.yaml Sample pod for development/testing +client.py Sample code to make API requests directed at a pod running this code ``` -### Arguments -- `-c`, `--contentTopics` (default: `/my-app/1/dst/proto`): Content topic to query. -- `-p`, `--pubsubTopic` (default: `/waku/2/rs/2/0`): Pubsub topic. -- `-ps`, `--pageSize` (default: 60): Number of messages per request. -- `-cs`, `--cursor` (optional): Cursor for pagination. - -### Example in Kubernetes yaml -``` -containers: - - name: container - image: /get_store_messages:v1.0.0 - imagePullPolicy: IfNotPresent - command: - - sh - - -c - - python /app/store_msg_retriever.py --contentTopics=/my-app/1/dst/proto -``` - -### How It Works -Queries a random waku node by selecting a random ip from `"zerotesting-service:8645"`. -It keeps querying that node until all messages are retrieved. ### Changelog +- `v2.0.0`: + - Changed to using a ConfigMap to define Targets, Endpoints, Requests, and Actions + - Added server capability + - Removed --debug mode logic - `v1.0.1`: - Added `--debug` mode. Makes multiple API requests to each IP - Added `--select-types` mode diff --git a/api_requester.py b/api_requester.py new file mode 100644 index 0000000..1e49d59 --- /dev/null +++ b/api_requester.py @@ -0,0 +1,327 @@ +import argparse +import random +import traceback +from argparse import Namespace +from collections import defaultdict +from pathlib import Path +from typing import Dict, List + +import requests +import uvicorn +import yaml +from fastapi import Depends, FastAPI, HTTPException +from kubernetes import client, config +from kubernetes.client.models.v1_pod import V1Pod +from pydantic import BaseModel, ConfigDict + +from configs import ConfigAction, ConfigEndpoint, ConfigRequest, ConfigTarget +from utils import paged_request, setup_logger + +logger = setup_logger(__file__) + + +app = FastAPI() + + +class TargetPodInfo(BaseModel): + model_config = ConfigDict(arbitrary_types_allowed=True) + + config_target: ConfigTarget + pod: V1Pod + + @property + def pod_name(self) -> str: + return self.pod.metadata.name + + +def do_request(request: ConfigRequest, pod_info: TargetPodInfo): + raise NotImplementedError() + + +def call_endpoint(endpoint: ConfigEndpoint, pod_info: TargetPodInfo) -> dict: + result_data = {"request": {"configEndpoint": endpoint}} + request_data = {"params": endpoint.params, "headers": endpoint.headers} + + try: + request_data["url"] = endpoint.url.format( + node=pod_info.pod.status.pod_ip, port=pod_info.config_target.port + ) + request_data["pod"] = f"{pod_info.pod.metadata.name}" + logger.info(f"request_data: {request_data}") + + if endpoint.paged: + if endpoint.type != "GET": + raise NotImplementedError("Paged requests only implemented for GET requests.") + result = paged_request(request=request_data, max_attempts=1, page_request_delay=0) + else: + if endpoint.type == "POST": + result = requests.post( + request_data["url"], + json=request_data["params"], + headers=request_data["headers"], + ) + elif endpoint.type == "GET": + result = requests.post( + request_data["url"], + json=request_data["params"], + headers=request_data["headers"], + ) + else: + raise AttributeError(f"Unknown request type. request: `{endpoint}`") + + result_data["request"].update(request_data) + result_data["response"] = {"status_code": result.status_code, "text": result.text} + except Exception as e: + error = traceback.format_exc() + logger.error( + f"Exception attempting API request. endpoint: `{endpoint}`, exception: `{e}`, error: `{error}`" + ) + result_data["exception"] = error + + logger.info(result_data) + return result_data + + +def assert_unique_attr(objects: List[object], attribute: str): + names = [getattr(endpoint, attribute) for endpoint in objects] + duplicates = set() + seen = [] + + for name in names: + if any(name == item for item in seen): + duplicates.add(name) + else: + seen.append(name) + + assert not duplicates, ( + f"At least one object has the same attribute as another. " + f"Attribute name: `{attribute}`. " + f"Duplicate attributes: `{duplicates}`. " + f"Objects: `{objects}`" + ) + + +def get_pods_for_target(target: ConfigTarget) -> List[str]: + config.load_incluster_config() + v1 = client.CoreV1Api() + namespace = open("/var/run/secrets/kubernetes.io/serviceaccount/namespace").read() or "default" + + if target.service is not None: + pods = v1.list_namespaced_pod(namespace) + else: + service = v1.read_namespaced_service(target.service, namespace) + selector = service.spec.selector + selector_str = ",".join([f"{k}={v}" for k, v in selector.items()]) + pods = v1.list_namespaced_pod(namespace, label_selector=selector_str) + + return list(filter(lambda pod: target.matches(pod), pods.items)) + + +def parse_config(config: Dict[str, List[object]]) -> Dict[str, Dict[str, object]]: + targets = [ConfigTarget.model_validate(targ) for targ in config.get("targets", [])] + targets_dict = {target.name: target for target in targets} + assert_unique_attr(targets, "name") + + endpoints = [ConfigEndpoint.model_validate(endpoint) for endpoint in config["endpoints"]] + endpoints_dict = {endpoint.name: endpoint for endpoint in endpoints} + assert_unique_attr(endpoints, "name") + + requests = [] + for request_dict in config["requests"]: + request_dict["endpoint"] = endpoints_dict[request_dict["endpoint"]] + requests.append(ConfigRequest.model_validate(request_dict)) + requests_dict = {request.name: request for request in requests} + assert_unique_attr(requests, "name") + + actions: List[ConfigAction] = [] + for action in config["actions"]: + try: + action["requests"] = [requests_dict[req] for req in action["requests"]] + except KeyError as e: + raise ValueError( + f"Action contains unknown request. action: `{action}` requests: `{requests_dict}`" + ) from e + try: + action["targets"] = [targets_dict[targ] for targ in action["targets"]] + except KeyError as e: + raise ValueError( + f"Action contains unknown target. action: `{action}` targets: `{targets_dict}`" + ) from e + actions.append(ConfigAction.model_validate(action)) + actions_dict = {action.name: action for action in actions} + assert_unique_attr(actions, "name") + + return { + "targets": targets_dict, + "endpoints": endpoints_dict, + "requests": requests_dict, + "actions": actions_dict, + } + + +def load_configs(config_files: List[str]) -> Dict[str, Dict[str, object]]: + logger.info(f"Loading configs: {config_files}") + full_config = defaultdict(list) + for config_file in config_files: + with open(config_file, "r") as file: + config = yaml.safe_load(file) + for key, value in config.items(): + full_config[key].extend(value) + return parse_config(full_config) + + +def do_action( + action: ConfigAction, + pods: List[TargetPodInfo], +) -> List[TargetPodInfo]: + target_names = [target.name for target in action.targets] + possible_pods = [pod for pod in pods if pod.config_target.name in target_names] + + if action.order == "random": + random.shuffle(possible_pods) + elif action.order == "ascending": + possible_pods.sort(key=lambda pod: pod.pod_name) + elif action.order == "descending": + possible_pods.sort(key=lambda pod: pod.pod_name, reverse=True) + else: + raise ValueError(f"Unknown order for action: {action.order}") + + pods = [] + count = len(possible_pods) if action.pod_count == "all" else action.pod_count + index = action.pod_start_index + for _ in range(count): + pods.append(possible_pods[index]) + index = (index + 1) % len(possible_pods) + + if action.loop_order == "foreach_pod_make_all_requests": + for pod in pods: + for request in action.requests: + # time.sleep(delay_between_requests) TODO + call_endpoint(request, pod) + elif action.loop_order == "foreach_request_target_each_pod": + for request in action.requests: + # TODO: ensure time between requests has elapsed + for pod in pods: + call_endpoint(request, pod) + else: + raise ValueError(f"Unknown loop_order for action: {action}") + + +def get_pod_infos(targets: List[ConfigTarget]) -> List[TargetPodInfo]: + pods_info: List[TargetPodInfo] = [] + for target in targets: + pods = get_pods_for_target(target) + for pod in pods: + pods_info.append(TargetPodInfo(config_target=target, pod=pod)) + return pods_info + + +def create_app(config) -> FastAPI: + app = FastAPI() + + async def get_config(): + logger.info(f"todo get config: {config}") + return config + + class InvokeRequestData(BaseModel): + target: ConfigTarget | str + endpoint: ConfigEndpoint | str + + @app.post("/process") + # TODO: Implement try/catch return error in decorator. It will be the same for all endpoints. + def process_data(data: InvokeRequestData, config=Depends(get_config)): + """ + Performs an API request to the given endpoint on the given target. + + :param data: Contains target and endpoint. + For each, the argument may either the name from the config, + or a custom object passed in as a dict. + + Sample usage (from outside the cluster): + data = { + "target": { + "name": "dummy", + "service": "zerotesting-lightpush-client", + "name_template": "lpclient-0-0", + }, + "endpoint": "lightpush-publish-static-sharding", + } + url = f"http://{external_ip}:{node_port}/process" + response = requests.post(url, json=data) + """ + try: + try: + # Treat target as the name of a preset target from config. + target = config["targets"][data.target] + except TypeError: + # If no target with that name exists, treat as custom target. + target = data.target + try: + # Treat endpoint as the name of preset endpoint from config. + endpoint = config["endpoints"][data.endpoint] + except TypeError: + # If no endpoint with that name exists, treat as custom endpoint. + endpoint = data.endpoint + request = ConfigRequest( + name="dummy_request", endpoint=endpoint, retries=0, retry_delay=0 + ) + pod_info = next(iter(get_pod_infos([target]))) + result = call_endpoint(request.endpoint, pod_info) + return result + except Exception as e: + # TODO: Add hints to errors (eg. Action doesn't exist, etc) + logger.error(HTTPException(status_code=500, detail=f"{e!r}\n{traceback.format_exc()}")) + raise HTTPException(status_code=500, detail=f"{e!r}\n{traceback.format_exc()}") + + return app + + +def main(args: Namespace): + config = load_configs(args.config_files) + available_endpoints = [endpoint.name for endpoint in config["endpoints"].values()] + logger.debug(f"Loaded config. Available endpoints: {available_endpoints}") + if args.mode == "server": + app = create_app(config) + uvicorn.run(app, host="0.0.0.0", port=args.port, log_config=None) + else: + pods_info = get_pod_infos(config["targets"]) + for action in config["actions"]: + do_action(action, pods_info) + + +def mode_type(value): + if value not in ["batch", "server"]: + raise argparse.ArgumentTypeError(f"Invalid mode: {value}. Must be 'batch' or 'server'.") + return value + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Waku storage retriever") + parser.add_argument( + "--config", + type=Path, + action="append", + dest="config_files", + required=True, + help="Paths to config files. Can be passed multiple times.", + ) + parser.add_argument( + "--mode", + type=mode_type, + default="server", + help="Batch: Run actions immediately. Server: Wait for API calls to /action/ to run.", + ) + parser.add_argument( + "--port", + type=int, + default=8645, + help="Port for the action HTTP server (default 8000)", + ) + + args = parser.parse_args() + return args + + +if __name__ == "__main__": + args = parse_args() + main(args) diff --git a/bind.yaml b/bind.yaml new file mode 100644 index 0000000..5e3da67 --- /dev/null +++ b/bind.yaml @@ -0,0 +1,24 @@ +apiVersion: rbac.authorization.k8s.io/v1 +kind: Role +metadata: + namespace: zerotesting + name: pod-service-reader +rules: +- apiGroups: [""] + resources: ["pods", "services"] + verbs: ["get", "list", "watch"] + +--- +apiVersion: rbac.authorization.k8s.io/v1 +kind: RoleBinding +metadata: + name: pod-service-reader-binding + namespace: zerotesting +subjects: +- kind: ServiceAccount + name: default # service account your pod uses + namespace: zerotesting +roleRef: + kind: Role + name: pod-service-reader + apiGroup: rbac.authorization.k8s.io diff --git a/build.sh b/build.sh index af265c9..69fae89 100644 --- a/build.sh +++ b/build.sh @@ -1,2 +1,8 @@ -docker build -t /get_store_messages:v1.0.0 . -docker push /get_store_messages:v1.0.0 \ No newline at end of file +docker build -t /pod-api-requester: --target debug . +# or +docker build -t /pod-api-requester: --target production . +# The default build is production +docker build -t /pod-api-requester: + + +docker push /pod-api-requester: \ No newline at end of file diff --git a/client.py b/client.py new file mode 100644 index 0000000..2ec127e --- /dev/null +++ b/client.py @@ -0,0 +1,136 @@ +import asyncio +import json +from typing import Literal, Optional + +import requests +from kubernetes import client, config +from pydantic import NonNegativeInt + +from utils import setup_logger + +logger = setup_logger(__file__) + + +async def main(): + raise NotImplementedError("Choose your Kubernetes config path and remove this.") + config.load_kube_config("/path_to_kube_config.yaml") + config.load_kube_config() # WARNING! LOCAL + publish_message( + namespace="zerotesting", + message_type="lightpush", + pod_name_template="lpclient-0-0", + service="zerotesting-lightpush-client", + ) + + +PublishType = Literal["lightpush", "relay"] + + +async def publish_message( + namespace: str, + message_type: PublishType, + *, + pod_name_template: Optional[str] = None, + service: Optional[str] = None, + stateful_set_name: Optional[str] = None, + port: NonNegativeInt = 80, +): + if message_type == "lightpush": + endpoint = "lightpush-publish-static-sharding" + elif message_type == "relay": + raise NotImplementedError() + else: + raise ValueError("Unknown message type") + + data = { + "target": { + "name": "dummy", + "service": service, + "name_template": pod_name_template, + "stateful_set": stateful_set_name, + "port": port, + }, + "endpoint": endpoint, + } + + return await pod_api_request( + namespace=namespace, + service_name="zerotesting-publisher", + app="zerotenkay-publisher", + data=data, + ) + + +class PodApiRequestError(Exception): + pass + + +async def pod_api_request( + namespace: str, + service_name: str, + app: str, + data: dict, + *, + publisher_pod: str | NonNegativeInt = 0, +) -> dict: + v1 = client.CoreV1Api() + + try: + pods = v1.list_namespaced_pod(namespace=namespace, label_selector=f"app={app}") + if isinstance(publisher_pod, str): + pod = next(pod for pod in pods.items if pod.metadata.name == publisher_pod) + else: + pod = pods.items[publisher_pod] + except IndexError as e: + logger.error(f"No pod found. app: `{app}` pod_index: `{publisher_pod}`") + raise ValueError() from e + except StopIteration as e: + logger.error(f"No pod found. app: `{app}` pod_name: `{publisher_pod}`") + raise ValueError() from e + + # Get publisher IP. + node = v1.read_node(name=pod.spec.node_name) + target_ip = kube_utils.get_node_ip(node) + + # Get publisher port. + service = v1.read_namespaced_service(service_name, namespace) + node_port = service.spec.ports[0].node_port + if node_port is None: + raise ValueError(f"Failed to find port for service. Service: `{service.metadata.name}`") + + url = f"http://{target_ip}:{node_port}/process" + + logger.info(f"publishing message. url: `{url}` data: `{data}`") + response = requests.post(url, json=data) + response_obj = json.loads(response.text) + if response.status_code != 200: + err = response_obj["detail"].replace("\n", "\n") + logger.error(err) + raise PodApiRequestError(response_obj) + + try: + # Assuming that the pod we made the API request to returns a response with a JSON object. + inner_response_obj = json.loads(response_obj["response"]["text"]) + response_obj["inner_response"] = inner_response_obj + if response_obj["response"]["status_code"] != 200: + # JsWaku puts the error under the key "error". + try: + err = inner_response_obj["error"].replace("\n", "\n") + except KeyError as e: + err = "" + logger.error(f"Publisher request returned failure. inner_error: `{err}`") + raise PodApiRequestError(response_obj) + except json.JSONDecodeError as e: + # Response was not a Json object. + pass + except KeyError as e: + err = response_obj["exception"].replace("\n", "\n") + logger.error(f"The publisher's API request attempt failed. Exception: `{err}`") + raise PodApiRequestError(response_obj) from e + + logger.info(f"Response: `{response_obj}`") + return response_obj + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/config.yaml b/config.yaml new file mode 100644 index 0000000..ae75b14 --- /dev/null +++ b/config.yaml @@ -0,0 +1,62 @@ +apiVersion: v1 +kind: ConfigMap +metadata: + name: api-requester-config + namespace: zerotesting +data: + config.yaml: | + targets: + - name: firstClients + service: zerotesting-lightpush-client + name_template: "^client-0-(1[0-2]|[0-9])$" + stateful_set: "client-0" + + - name: jswakuClients + service: zerotesting-lightpush-client + stateful_set: "client-0" + port: 8080 + + - name: nwakuClients + service: zerotesting-lightpush-client + stateful_set: "client-0" + port: 8645 + + + endpoints: + - name: set-debug + headers: {"accept": "text/plain"} + params: {"logLevel": "DEBUG"} + url: "http://{node}/admin/v1/log-level/DEBUG" + type: "POST" + paged: False + + - name: lightpush-publish-static-sharding + url: "http://{node}:{port}/lightpush/v3/message" + headers: {"Content-Type": "application/json"} + params: {"pubsubTopic": "/waku/2/rs/2/0", "message": {"contentTopic": "/test/1/cross-network/proto", "payload": "W1EsIDIsIDNd"}} + type: "POST" + paged: False + + - name: lightpush-publish-auto-sharding + url: "http://{node}:{port}/lightpush/v3/message" + headers: {"Content-Type": "application/json"} + params: { "pubsubTopic":"", "message": {"contentTopic": "/test/1/cross-network/proto", "payload" : "W1EsIDIsIDNd"}} + type: "POST" + paged: False + + + requests: + - name: publish-request + endpoint: lightpush-publish-static-sharding + retries: 0 + retry_delay: 0.3 + + actions: + - name: publish-to-clients-random + requests: ["publish-request"] + targets: ["nwakuClients"] + pod_start_index: 0 + pod_count: 91 + order: ascending + loop_order: foreach_request_target_each_pod + diff --git a/configs.py b/configs.py new file mode 100644 index 0000000..2ddb3f8 --- /dev/null +++ b/configs.py @@ -0,0 +1,222 @@ +import datetime +import logging +import re +from typing import List, Literal, Optional + +from kubernetes import client +from kubernetes.client.models.v1_pod import V1Pod +from pydantic import BaseModel, NonNegativeFloat, NonNegativeInt, PositiveInt + + +class UTCFormatter(logging.Formatter): + def formatTime(self, record, datefmt=None): + # Get UTC time and format with milliseconds + utc_dt = datetime.datetime.utcfromtimestamp(record.created) + if datefmt: + s = utc_dt.strftime(datefmt) + # Add milliseconds + s = s + f".{int(record.msecs):03d}" + return s + else: + t = utc_dt.strftime("%Y-%m-%d %H:%M:%S") + s = f"{t}.{int(record.msecs):03d}" + return s + + +logfmt = "%(asctime)s %(levelname)s [%(name)s] %(message)s" +datefmt = "%Y-%m-%d %H:%M:%S" + +handler = logging.StreamHandler() +handler.setFormatter(UTCFormatter(logfmt, datefmt=datefmt)) + +logging.basicConfig(level=logging.INFO, handlers=[handler]) + +logger = logging.getLogger(__file__) + + +class ConfigEndpoint(BaseModel): + """Describes an endpoint on a pod in the cluster. + This endpoint may exist on multiple pods, or just a single pod, or no pod at all. + + It is the responsibility of the caller to combine a defined endpoint with a proper pod.""" + + name: str + """The name of this config object.""" + + headers: dict + """ + The header to send with the request. + + Typically either + headers: {"accept": "text/plain"} + or + {"Content-Type": "application/json"} + """ + + params: dict + """HTTP POST data to include with request.""" + + url: str + """Url for the endpoint. + Instances of `{node}` and `{port}` will be replaced with + the pod IP and the target port respectively. + + For example, when calling the following endpoint: + `http://{node}:{port}/lightpush/v3/message` + on the node at index `2` of the a target with: + `stateful_set: "client", port: 8645`, then the following url will be used: + `http://client-2:8645/lightpush/v3/message` + """ + + type: Literal["POST", "GET"] + """Specifies the method of the request. Either `POST` or `GET`.""" + + paged: bool + """Use `True` if the request returns paged data. Otherwise, use `False`.""" + + +class ConfigRequest(BaseModel): + """A request to be made to a pod. + Contains the `endpoint` and some additional data for retries/delays.""" + + name: str + """The name of this config object.""" + + endpoint: ConfigEndpoint + """The Endpoint to use for this request.""" + + retries: NonNegativeInt + """Number of times to retry the request if it fails.""" + + retry_delay: NonNegativeFloat + """The delay between each retry attempt for this request.""" + + +class ConfigTarget(BaseModel): + """A config describing pods. + This is a list of filters to apply to any pod + to see if that pod is part of the target group. + """ + + name: str + """The name of this config object. Not the pod name.""" + + service: Optional[str] = None + """The name of the service that any target pod must belong. + Example: zerotesting-bootstrap""" + + name_template: Optional[str] = None + """Regex describing the pod names. Example: ^client-([0-9])$""" + + stateful_set: Optional[str] = None + """Name of the StatefulSet that any target pod must belong to. + Example: "bootstrap" + """ + + port: NonNegativeInt = 80 + """Port to use for requests to endpoints with this target. + Default is 80.""" + + def matches(self, pod: V1Pod) -> bool: + """Check if pod is a valid target of self""" + + if self.stateful_set is not None: + if pod.metadata.owner_references is None: + return False + if not all( + [ + owner.kind == "StatefulSet" and owner.name == self.stateful_set + for owner in pod.metadata.owner_references + ] + ): + return False + + if self.name_template is not None: + if not re.search(self.name_template, pod.metadata.name): + return False + + if self.service is not None: + v1 = client.CoreV1Api() + namespace = ( + open("/var/run/secrets/kubernetes.io/serviceaccount/namespace").read() or "default" + ) + service = v1.read_namespaced_service(self.service, namespace) + selector = service.spec.selector + if not all([pod.metadata.labels.get(key) == value for key, value in selector.items()]): + return False + + return True + + +class ConfigAction(BaseModel): + """Description of an action to take. Here is how an action is performed: + 1. For each ConfigTarget, add all pods to the list. + 2. Sort the list of pods according to `order`. + 3. Starting at `pod_start_index`, take `pod_count` pods. + 4. According to `loop_order`, make every request in `requests` to every pod in the remaining list. + """ + + name: str + """The name of this config object.""" + + loop_order: Literal["foreach_pod_make_all_requests", "foreach_request_target_each_pod"] + """Which algorithm to use to determine how requests should be made to pods. + + `foreach_pod_make_all_requests`: Loop through the list of pods. + At each pod, make all the requests in `requests` + + `foreach_request_target_each_pod`: Loop through the `requests` list. + For each `ConfigRequest`, execute that request on all pods in the + list of pods derived from the algorithm described above. + """ + + pod_start_index: NonNegativeInt = 0 + """Allows a user to "skip" a certain amount of pods. + This is applied to a list created by combining the lists of pods from `targets`, + and sorting the list according to `order`. + + Assumes that `pod_start_index < len(all_pods)`. + """ + + pod_count: PositiveInt | Literal["all"] = "all" + """The number of pods for this action. + + This can be used to limit the total number of pods considered for requests. + Like, `pod_start_index`, this applies to the list of pods created via combining + pods from `targets` and sorting them. + + If `pod_count` is `"all"`, then all pods will be used. This will not deduplicate any pods in the list. + If `pod_count > len(all_pods)`, then the cursor will loop back to the beginning of the list and continue + adding pods until the list of pods to use has exactly `pod_count` elements in it. + """ + + order: Literal["ascending", "descending", "random"] | None + """Once the list of possible pods is gathered by combining the lists of pods for each `ConfigTarget`, + they will be sorted by this ordering before applying `pod_start_index` and `pod_count`. + """ + + targets: List[ConfigTarget] + """A list of all `ConfigTarget`s used to gather the list of pods. + For each target, all pods will be added to as potential targets. + Then the list will be sorted according to `order`, and spliced + according to `pod_start_index` and `pod_count`. + + Note: A pod may match multiple `ConfigTarget`s. In this case, + the pod will be added to the list as many times as it matches. + + For example, with `ConfigTarget`s {"stateful_set": "some_pod"} and {"name_template": "^some_pod-[1-2]$"}, + where the StatefulSet of some_pod has `replicas: 4`, the list of pods to use would be: + + ["some_pod-0", "some_pod-1", "some_pod-2", "some_pod-3", "some_pod-1", "some_pod-2"] + + which then may be sorted by `ascending` to look like: + + ["some_pod-0", "some_pod-1", "some_pod-1", "some_pod-2", "some_pod-2", "some_pod-3"] + + then, the list would be spliced using `all_pods[pod_start_index:pod_start_index+pod_count]`, + assuming that `pod_start_index+pod_count < len(all_pods)`. + """ + + requests: List[ConfigRequest] + """The list of requests to do for each pod that ends up in the list of pods to request to. + Every request will be executed, but it may not be on all pods matching every `ConfigTarget` in `targets`. See `targets`.""" diff --git a/deployment.yaml b/deployment.yaml new file mode 100644 index 0000000..246bb55 --- /dev/null +++ b/deployment.yaml @@ -0,0 +1,43 @@ +apiVersion: v1 +kind: Pod +metadata: + name: publisher + namespace: zerotesting + labels: + app: zerotenkay-publisher +spec: + restartPolicy: Never + dnsConfig: + searches: + - zerotesting-publisher.zerotesting.svc.cluster.local + volumes: + - name: api-requester-config-volume + configMap: + name: api-requester-config + containers: + - name: publisher-container + image: pearsonwhite/pod-api-requester:1e161cdf41478000bea17f9332f3624e9aad0829 + imagePullPolicy: Always + command: + [ + "python", + "/app/api_requester.py", + "--mode", + "server", + "--config", + "/mount/config.yaml", + ] + ports: + - containerPort: 8645 + - containerPort: 8008 + - containerPort: 8080 + volumeMounts: + - name: api-requester-config-volume + mountPath: /mount + resources: + requests: + memory: 64Mi + cpu: 150m + limits: + memory: 600Mi + cpu: 400m diff --git a/publisher-service.yaml b/publisher-service.yaml new file mode 100644 index 0000000..429ef6c --- /dev/null +++ b/publisher-service.yaml @@ -0,0 +1,14 @@ +apiVersion: v1 +kind: Service +metadata: + name: zerotesting-publisher + namespace: zerotesting-pwhite +spec: + type: NodePort + selector: + app: zerotenkay-publisher + ports: + - protocol: TCP + port: 8000 + targetPort: 8645 + nodePort: 30080 \ No newline at end of file diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..2153e81 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,7 @@ +aiohttp==3.9.3 +fastapi==0.124.0 +kubernetes==27.2.0 +pydantic==2.12.5 +PyYAML==6.0.3 +Requests==2.32.5 +uvicorn==0.38.0 diff --git a/store_msg_retriever.py b/store_msg_retriever.py deleted file mode 100644 index bbb3170..0000000 --- a/store_msg_retriever.py +++ /dev/null @@ -1,639 +0,0 @@ -# Python Imports -import argparse -import datetime -import json -import logging -import socket -import time -import traceback -from argparse import Namespace -from copy import deepcopy -from pathlib import Path -from typing import Any, Dict, List, Tuple - -import requests -from pydantic import BaseModel, Field, PositiveInt - - -class UTCFormatter(logging.Formatter): - def formatTime(self, record, datefmt=None): - # Get UTC time and format with milliseconds - utc_dt = datetime.datetime.utcfromtimestamp(record.created) - if datefmt: - s = utc_dt.strftime(datefmt) - # Add milliseconds - s = s + f".{int(record.msecs):03d}" - return s - else: - t = utc_dt.strftime("%Y-%m-%d %H:%M:%S") - s = f"{t}.{int(record.msecs):03d}" - return s - - -# Usage -logfmt = "%(asctime)s %(levelname)s [%(name)s] %(message)s" -datefmt = "%Y-%m-%d %H:%M:%S" - -handler = logging.StreamHandler() -handler.setFormatter(UTCFormatter(logfmt, datefmt=datefmt)) - -logging.basicConfig(level=logging.INFO, handlers=[handler]) - -logger = logging.getLogger(__file__) - - -def next_cursor(data: Dict) -> str | None: - cursor = data.get("paginationCursor") - if not cursor: - logger.info("No more messages") - return None - - return cursor - - -def fetch_all_messages(base_url: str, initial_params: Dict, headers: Dict) -> List[str]: - all_messages = [] - params = initial_params.copy() - - while True: - logger.info( - f"requests.get: url: `{base_url}` init_params: `{initial_params}` params: `{params}`" - ) - response = requests.get(base_url, headers=headers, params=params) - logger.info(f"response: `{response.text}`") - if response.status_code != 200: - logger.error(f"Error fetching data: {response.status_code}") - logger.error(response.text) - break - - data = response.json() - logger.info(data) - if data["statusCode"] != 200: - logger.info(f"failed. statusCode: `{data['statusCode']}`") - paged_messages = [message["messageHash"] for message in data["messages"]] - logger.info(f"Retrieved {len(paged_messages)} messages") - all_messages.extend([message["messageHash"] for message in data["messages"]]) - - cursor = next_cursor(data) - if not cursor: - break - params["cursor"] = cursor - return all_messages - - -def dict_extract(obj: dict, path: Path): - def extract(obj: Any, parts: list, is_list=False): - if isinstance(obj, list): - results = [] - for item in obj: - results.extend(extract(item, parts, is_list=True)) - return results - if not parts: - return [obj] if is_list else obj - next_obj = obj[parts[0]] - return extract(next_obj, parts[1:], is_list) - - return extract(obj, path.parts) - - -def paged_request(request: dict, max_attempts: PositiveInt, page_request_delay: float) -> dict: - """ - GET request with a "paged" param. - - :param request: Must contain "params":dict. - """ - attempt_num = 1 - - url = request["url"] - all_messages = [] - pages_data = [] - params = request["params"] - status_codes = [] - inner_status_codes = [] - while True: - time.sleep(page_request_delay) - - logger.info(f"Making paged request. request: `{request}`, params=`{params}`") - response = requests.get(url, headers=request["headers"], params=params) - - try: - data = response.json() - except requests.exceptions.JSONDecodeError: - data = response.text - - status_codes.append(response.status_code) - pages_data.append(data) - - logger.info(f"response to paged request: `{response}`") - if response.status_code != 200: - logger.error( - f"Error fetching paged data. status_code: `{response.status_code}` data: `{data}`" - ) - break - - inner_status_codes.append(data["statusCode"]) - logger.info(f"Response data: `{data}`") - - if data["statusCode"] != 200: - logger.info( - f"inner_status_code != 200: status_code: `{data['statusCode']}`, attempt: `{attempt_num}`" - ) - - if attempt_num >= max_attempts: - logger.info(f"Exhausted all attempts: `{attempt_num}`") - break - attempt_num += 1 - continue - - logger.info(f"inner_status_code == 200: attempt: `{attempt_num}`") - if attempt_num > 1: - logger.info("A previous attempt failed, but now it worked.") - - paged_data = dict_extract(data, request.get("extract_keys", Path())) - logger.info(f"Retrieved {len(paged_data)} messages on attempt `{attempt_num}`") - all_messages.extend(paged_data) - - cursor = next_cursor(data) - if not cursor: - logger.info(f"page request finished with !cursor on attempt `{attempt_num}`") - break - params["cursor"] = cursor - - attempt_num = 1 - - logger.info("finished page request") - return { - "request": request, - "response": { - "statusCodes": status_codes, - "inner_statusCodes": inner_status_codes, - "messages": all_messages, - "pages": pages_data, - "attempt_num": attempt_num, - }, - } - - -def api_request(action, request) -> dict: - url = request["url"] - response = action(url, request["headers"], request.get("params")) - - try: - data = response.json() - except requests.exceptions.JSONDecodeError: - data = response.text - - if response.status_code != 200: - logger.error(f"Error fetching data: {response.status_code}") - logger.error(data) - - return { - "request": request, - "response": { - "statusCode": response.status_code, - "contents": data, - }, - } - - -def get_node_info( - name: str, node: str, api_args: dict, delay_between_requests=0.3 -) -> Dict[str, dict]: - all_requests = { - "debug": { - "headers": {"accept": "text/plain"}, - "params": {"logLevel": "DEBUG"}, - "url": "http://{node}/admin/v1/log-level/DEBUG", - "type": "POST", - }, - "info": { - "url": "http://{node}/debug/v1/info", - "headers": {"accept": "application/json"}, - "type": "GET", - }, - "peers": { - "url": "http://{node}/admin/v1/peers", - "headers": {"accept": "application/json"}, - "type": "GET", - }, - "mesh": { - "url": "http://{node}/admin/v1/peers/mesh", - "headers": {"accept": "application/json"}, - "type": "GET", - }, - "stats": { - "url": "http://{node}/admin/v1/peers/stats", - "headers": {"accept": "application/json"}, - "type": "GET", - }, - "connected": { - "url": "http://{node}/admin/v1/peers/connected", - "headers": {"accept": "application/json"}, - "type": "GET", - }, - "service": { - "url": "http://{node}/admin/v1/peers/service", - "headers": {"accept": "application/json"}, - "type": "GET", - }, - "store_messages": { - "url": f"http://{node}/store/v3/messages", - "headers": {"accept": "application/json"}, - "paged": True, - "params": api_args, - "extract_keys": Path("messages", "messageHash"), - }, - } - - request_data = {} - for key, node_request in all_requests.items(): - request = deepcopy(node_request) - request["url"] = request["url"].format(node=node) - request["node"] = name - if request.get("type") == "POST": - action = lambda url, headers, params: requests.post(url, data=params, headers=headers) - elif request.get("type") == "GET": - action = lambda url, headers, params: requests.get(url, headers=headers, params=params) - - try: - if request.get("paged"): - result = paged_request(request=request, max_attempts=1, page_request_delay=0) - else: - result = api_request(action, request) - request_data[key] = result - except Exception as e: - error = traceback.format_exc() - logger.error( - f"Exception attempting API request. request: `{request}`, exception: `{e}`, error: `{error}`" - ) - request_data[key] = { - "request": request, - "exception": error, - } - - time.sleep(delay_between_requests) - - return request_data - - -def serializer(obj): - if isinstance(obj, Path): - return obj.as_posix() - if isinstance(obj, set): - return list(obj) - raise TypeError(f"Object of type {obj.__class__.__name__} is not JSON serializable") - - -def resolve_dns(node: str) -> Tuple[str, str]: - start_time = time.time() - name, port = node.split(":") - ip_address = socket.gethostbyname(name) - entire_hostname = socket.gethostbyaddr(ip_address) - hostname = entire_hostname[0].split(".")[0] - elapsed = (time.time() - start_time) * 1000 - logger.info(f"{node} DNS Response took {elapsed} ms") - logger.info(f"Talking with {hostname}, ip address: {ip_address}") - - return (entire_hostname, f"{ip_address}:{port}") - - -class NodeType(BaseModel): - name_template: str - """Format string for node name. Eg. fserver-0-{index}""" - service: str - count_key: str - namespace: str = Field(default="zerotesting") - - def dns_name(self, index: PositiveInt) -> str: - """Return name for DNS lookup. - . - """ - return f"{self.get_node_name(index)}.{self.service}" - - def get_node_name(self, index: PositiveInt) -> str: - return self.name_template.format(index=index) - - -node_types = [ - NodeType( - name_template="store-0-{index}", - service="zerotesting-store", - count_key="store", - ), - NodeType( - # Note the plural "nodes" with an 's'! - # This is to match the name used in regression tests. - name_template="nodes-0-{index}", - service="zerotesting-service", - count_key="relay", - ), - NodeType( - name_template="fserver-0-{index}", - service="zerotesting-filter", - count_key="filter_server", - ), - NodeType( - name_template="fclient-0-{index}", - service="zerotesting-filter", - count_key="filter_client", - ), - NodeType( - name_template="lpserver-0-{index}", - service="zerotesting-lightpush-server", - count_key="lightpush_server", - ), - NodeType( - name_template="lpclient-0-{index}", - service="zerotesting-lightpush-client", - count_key="lightpush_client", - ), - NodeType( - name_template="bootstrap-{index}", - service="zerotesting-bootstrap", - count_key="bootstrap", - ), -] - - -def get_ips_by_type(args: dict, *, namespace=None) -> List[Tuple[str, str]]: - """ - Get node ips based on type flags (--store, --relay, etc) starting at start_index for each node type. - - :return: (name, ip) tuples for node specified. - :rtype: List[str, str] - """ - # TODO: Handle multiple shards. - - results = [] - for node_type in node_types: - start_index = args.get("start_index", 0) - if args[node_type.count_key] == "all": - try: - _, _, ip_list = socket.gethostbyname_ex(node_type.service) - count = len(ip_list) - start_index - except socket.gaierror: - # This happens when either: - # 1. The service doesn't exist. - # 2. No pods with the matching app selector exist, thus though the service exists, it isn't running on any pod. - count = 0 - # TODO: Check at the end if all `count` ips have been found. - # TODO: Add "unknown-{index}" for ips not in {nodetype}-0-{index} - # Note that if node types share the same service, count will be set to the total. - # for example fserver/fclient both use zerotesting-filter. - else: - try: - count = int(args[node_type.count_key]) - except (KeyError, TypeError): - logger.info(f"No count for nodetype specified. `{node_type}`") - continue - - logger.info( - f"Getting {count} IPs from nodes of type `{node_type.name_template}` starting at index {start_index}" - ) - for index in range(start_index, start_index + count): - dns = node_type.dns_name(index) - try: - _, _, ips = socket.gethostbyname_ex(dns) - results.append((node_type.get_node_name(index), ips[0])) - except Exception as e: - error = traceback.format_exc() - logger.error( - f"Failed to resolve dns. dns: `{dns}`, node_type: `{node_type}`, exception: `{e}`, error: {error}" - ) - - return results - - -def get_api_args(args_dict: dict) -> dict: - """These are the arguments that should be passed on to the GET request for store messages.""" - return { - key: value - for key, value in args_dict.items() - if key - in [ - "contentTopics", - "pubsubTopic", - "pageSize", - "cursor", - ] - } - - -def positive_int_or_all(value): - if value == "all": - return value - try: - int_value = int(value) - assert int_value >= 0 - return int_value - except (ValueError, AssertionError): - raise argparse.ArgumentTypeError(f"{value} is not an integer or 'all'.") - - -def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description="Waku storage retriever") - parser.add_argument( - "-c", "--contentTopics", type=str, help="Content topic", default="/my-app/1/dst/proto" - ) - parser.add_argument( - "-p", "--pubsubTopic", type=str, help="Pubsub topic", default="/waku/2/rs/2/0" - ) - parser.add_argument( - "-ps", "--pageSize", type=int, help="Number of messages to retrieve per page", default=60 - ) - parser.add_argument( - "-cs", - "--cursor", - type=str, - help="Cursor field intended for pagination purposes. ", - default="", - ) - - parser.add_argument( - "-d", - "--debug", - action="store_true", - help="", - dest="debug", - ) - - parser.add_argument( - "-rd", - "--request-delay", - type=float, - default=0.3, - help="Delay between each REST API call on a node. Only applicable in --debug mode.", - dest="delay_between_requests", - ) - - parser.add_argument( - "-t", - "--select-types", - action="store_true", - help="If specified, gathers ips for nodes of the types indicated by additional flags (e.g., --store, --relay). If not specified, selects a random node of any type.", - dest="select_types", - ) - - parser.add_argument( - "-s", - "--store", - type=positive_int_or_all, - help="Number of store nodes", - dest="store", - ) - parser.add_argument( - "-r", - "--relay", - type=positive_int_or_all, - help="Number of plain relay nodes", - dest="relay", - ) - parser.add_argument( - "-fs", - "--filter-server", - type=positive_int_or_all, - help="Number of fserver nodes", - dest="filter_server", - ) - parser.add_argument( - "-fc", - "--filter-client", - type=positive_int_or_all, - help="Number of fclient nodes", - dest="filter_client", - ) - parser.add_argument( - "-lps", - "--lightpush-server", - type=positive_int_or_all, - help="Number of lpserver nodes", - dest="lightpush_server", - ) - parser.add_argument( - "-lpc", - "--lightpush-client", - type=positive_int_or_all, - help="Number of lpclient nodes", - dest="lightpush_client", - ) - parser.add_argument( - "-bn", - "--bootstrap", - type=positive_int_or_all, - help="Number of bootstrap nodes", - dest="bootstrap", - ) - - parser.add_argument( - "-si", - "--start-index", - type=int, - default=0, - help="Start looking for at index: {nodetype}-0-{index}", - dest="start_index", - ) - - args = parser.parse_args() - assert args.select_types == any( - [ - args.relay, - args.store, - args.filter_server, - args.filter_client, - args.lightpush_server, - args.lightpush_client, - args.bootstrap, - ] - ), "--select-types should be True if any node types have been specified and False otherwise." - - return args - - -def main(args: Namespace): - args_dict = vars(args) - api_args = get_api_args(args_dict) - - logger.info(f"Arguments: {args_dict}") - - nodes = get_ips(args) - - messages = [] - for index, (name, node) in enumerate(nodes): - try: - logger.info( - f"fetching messages. name: `{name}` url: `{node}` index: {index+1}/{len(nodes)} ({100* (index+1) / len(nodes):.2f}%)" - ) - - url = f"http://{node}/store/v3/messages" - logger.info(f"Query to {url}") - headers = {"accept": "application/json"} - new_messages = fetch_all_messages(url, api_args, headers) - messages.extend(new_messages) - - except Exception as e: - error = traceback.format_exc() - print(f"exception while fetching messages. exception: `{e}`, error: `{error}`") - - logger.info("List of messages") - # # We do a print here, so it is easier to parse when reading from victoria logs - print(messages) - - -def get_ips(args) -> Tuple[str, str]: - port = 8645 - if args.select_types: - ips = get_ips_by_type(vars(args)) - logger.info(f"ips: ({len(ips)}): ```{ips}```") - return [(name, f"{ip}:{port}") for name, ip in ips] - else: - service = f"zerotesting-service:{port}" - return [resolve_dns(service)] - - -def main_debug(args: Namespace): - args_dict = vars(args) - api_args = get_api_args(args_dict) - - logger.info(f"Arguments: {args_dict}") - - nodes = get_ips(args) - for name, node in nodes: - attempt = 1 - max_attempts = 10 - delay = 0.5 - while True: - time.sleep(delay) - try: - logger.info(f"fetching messages. name: `{name}` url: `{node}` attempt: `{attempt}`") - logger.info(f"fetching messages. name: `{name}` url: `{node}`") - node_info = get_node_info(name, node, api_args, args.delay_between_requests) - node_info["attempt"] = attempt - logger.info( - f"store_msg_retriever::node_info: ```{json.dumps(node_info, default=serializer)}```" - ) - if all( - code == 200 - for code in node_info["store_messages"]["response"]["inner_statusCodes"] - ): - logging.info("No inner status failures") - if attempt > 1: - logger.info("main::A previous attempt failed, but now it worked.") - break - logging.info("Inner status failures detected") - if attempt >= max_attempts: - break - attempt += 1 - - except Exception as e: - error = traceback.format_exc() - logging.error( - f"exception while fetching messages. exception: `{e}`, error: `{error}`" - ) - - -if __name__ == "__main__": - args = parse_args() - if args.debug: - main_debug(args) - else: - main(args) diff --git a/utils.py b/utils.py new file mode 100644 index 0000000..17786b4 --- /dev/null +++ b/utils.py @@ -0,0 +1,303 @@ +import datetime +import logging +import socket +import time +import traceback +from pathlib import Path +from typing import Any, Dict, List, Tuple + +import requests +from pydantic import BaseModel, Field, PositiveInt + +LOGFMT = "%(asctime)s %(levelname)s [%(name)s] %(message)s" +DATEFMT = "%Y-%m-%d %H:%M:%S" + + +class UTCFormatter(logging.Formatter): + """Formatter that outputs UTC timestamps with milliseconds.""" + + def formatTime(self, record, datefmt=None): + dt = datetime.datetime.fromtimestamp(record.created, datetime.timezone.utc) + if datefmt is None: + datefmt = DATEFMT + base = dt.strftime(datefmt) + return f"{base}.{int(record.msecs):03d}" + + +def setup_logger(name: str) -> logging.Logger: + handler = logging.StreamHandler() + handler.setFormatter(UTCFormatter(LOGFMT, datefmt=DATEFMT)) + + logging.basicConfig(level=logging.INFO, handlers=[handler], force=True) + return logging.getLogger(name) + + +logger = setup_logger(__file__) + + +def get_ips_by_service(service: str) -> List[str]: + try: + _, _, ips = socket.gethostbyname_ex(service) + return ips[0] + except Exception as e: + error = traceback.format_exc() + logger.error( + f"Failed to resolve dns. service: `{service}`, exception: `{e}`, error: {error}" + ) + raise + + +class Target(BaseModel): + pod_name: str + ip: str + service: str + dns_name: str + + +class NodeType(BaseModel): + name_template: str + """Format string for node name. Eg. fserver-0-{index}""" + service: str + count_key: str + namespace: str = Field(default="zerotesting") + + def dns_name(self, index: PositiveInt) -> str: + """Return name for DNS lookup. + . + """ + return f"{self.get_node_name(index)}.{self.service}" + + def get_node_name(self, index: PositiveInt) -> str: + return self.name_template.format(index=index) + + +node_types = [ + NodeType( + name_template="store-0-{index}", + service="zerotesting-store", + count_key="store", + ), + NodeType( + # Note the plural "nodes" with an 's'! + # This is to match the name used in regression tests. + name_template="nodes-0-{index}", + service="zerotesting-service", + count_key="relay", + ), + NodeType( + name_template="fserver-0-{index}", + service="zerotesting-filter", + count_key="filter_server", + ), + NodeType( + name_template="fclient-0-{index}", + service="zerotesting-filter", + count_key="filter_client", + ), + NodeType( + name_template="lpserver-0-{index}", + service="zerotesting-lightpush-server", + count_key="lightpush_server", + ), + NodeType( + name_template="lpclient-0-{index}", + service="zerotesting-lightpush-client", + count_key="lightpush_client", + ), + NodeType( + name_template="bootstrap-{index}", + service="zerotesting-bootstrap", + count_key="bootstrap", + ), +] + + +def get_ips_by_type(args: dict, *, namespace=None) -> List[Tuple[str, str]]: + """ + Get node ips based on type flags (--store, --relay, etc) starting at start_index for each node type. + + :return: (name, ip) tuples for node specified. + :rtype: List[str, str] + """ + # TODO: Handle multiple shards. + + results = [] + for node_type in node_types: + start_index = args.get("start_index", 0) + if args[node_type.count_key] == "all": + try: + _, _, ip_list = socket.gethostbyname_ex(node_type.service) + count = len(ip_list) - start_index + except socket.gaierror: + # This happens when either: + # 1. The service doesn't exist. + # 2. No pods with the matching app selector exist, thus though the service exists, it isn't running on any pod. + count = 0 + # TODO: Check at the end if all `count` ips have been found. + # TODO: Add "unknown-{index}" for ips not in {nodetype}-0-{index} + # Note that if node types share the same service, count will be set to the total. + # for example fserver/fclient both use zerotesting-filter. + else: + try: + count = int(args[node_type.count_key]) + except (KeyError, TypeError): + logger.info(f"No count for nodetype specified. `{node_type}`") + continue + + logger.info( + f"Getting {count} IPs from nodes of type `{node_type.name_template}` starting at index {start_index}" + ) + for index in range(start_index, start_index + count): + dns = node_type.dns_name(index) + try: + _, _, ips = socket.gethostbyname_ex(dns) + results.append((node_type.get_node_name(index), ips[0])) + except Exception as e: + error = traceback.format_exc() + logger.error( + f"Failed to resolve dns. dns: `{dns}`, node_type: `{node_type}`, exception: `{e}`, error: {error}" + ) + + return results + + +def resolve_dns(node: str) -> Tuple[str, str]: + start_time = time.time() + name, port = node.split(":") + ip_address = socket.gethostbyname(name) + entire_hostname = socket.gethostbyaddr(ip_address) + hostname = entire_hostname[0].split(".")[0] + elapsed = (time.time() - start_time) * 1000 + logger.info(f"{node} DNS Response took {elapsed} ms") + logger.info(f"Talking with {hostname}, ip address: {ip_address}") + + return (entire_hostname, f"{ip_address}:{port}") + + +def get_ips(args) -> Tuple[str, str]: + port = 8645 + if args.select_types: + ips = get_ips_by_type(vars(args)) + logger.info(f"ips: ({len(ips)}): ```{ips}```") + return [(name, f"{ip}:{port}") for name, ip in ips] + else: + service = f"zerotesting-service:{port}" + return [resolve_dns(service)] + + +# TODO: Extraneous code? (unused) +def get_api_args(args_dict: dict) -> dict: + """These are the arguments that should be passed on to the GET request for store messages.""" + return { + key: value + for key, value in args_dict.items() + if key + in [ + "contentTopics", + "pubsubTopic", + "pageSize", + "cursor", + ] + } + + +def dict_extract(obj: dict, path: Path): + def extract(obj: Any, parts: list, is_list=False): + if isinstance(obj, list): + results = [] + for item in obj: + results.extend(extract(item, parts, is_list=True)) + return results + if not parts: + return [obj] if is_list else obj + next_obj = obj[parts[0]] + return extract(next_obj, parts[1:], is_list) + + return extract(obj, path.parts) + + +def next_cursor(data: Dict) -> str | None: + cursor = data.get("paginationCursor") + if not cursor: + logger.info("No more messages") + return None + + return cursor + + +def paged_request(request: dict, max_attempts: PositiveInt, page_request_delay: float) -> dict: + """ + GET request with a "paged" param. + + :param request: Must contain "params":dict. + """ + attempt_num = 1 + + url = request["url"] + all_messages = [] + pages_data = [] + params = request["params"] + status_codes = [] + inner_status_codes = [] + while True: + time.sleep(page_request_delay) + + logger.info(f"Making paged request. request: `{request}`, params=`{params}`") + response = requests.get(url, headers=request["headers"], params=params) + + try: + data = response.json() + except requests.exceptions.JSONDecodeError: + data = response.text + + status_codes.append(response.status_code) + pages_data.append(data) + + logger.info(f"response to paged request: `{response}`") + if response.status_code != 200: + logger.error( + f"Error fetching paged data. status_code: `{response.status_code}` data: `{data}`" + ) + break + + inner_status_codes.append(data["statusCode"]) + logger.info(f"Response data: `{data}`") + + if data["statusCode"] != 200: + logger.info( + f"inner_status_code != 200: status_code: `{data['statusCode']}`, attempt: `{attempt_num}`" + ) + + if attempt_num >= max_attempts: + logger.info(f"Exhausted all attempts: `{attempt_num}`") + break + attempt_num += 1 + continue + + logger.info(f"inner_status_code == 200: attempt: `{attempt_num}`") + if attempt_num > 1: + logger.info("A previous attempt failed, but now it worked.") + + paged_data = dict_extract(data, request.get("extract_keys", Path())) + logger.info(f"Retrieved {len(paged_data)} messages on attempt `{attempt_num}`") + all_messages.extend(paged_data) + + cursor = next_cursor(data) + if not cursor: + logger.info(f"page request finished with !cursor on attempt `{attempt_num}`") + break + params["cursor"] = cursor + + attempt_num = 1 + + logger.info("finished page request") + return { + "request": request, + "response": { + "statusCodes": status_codes, + "inner_statusCodes": inner_status_codes, + "messages": all_messages, + "pages": pages_data, + "attempt_num": attempt_num, + }, + }