[INF-187] Add elastic search to discovery provider (#3800)

AudiusProject · Sep 16, 2022 · eac4bb7 · eac4bb7
1 parent 31a81c1
commit eac4bb7
Show file tree

Hide file tree

Showing 3 changed files with 82 additions and 12 deletions.
diff --git a/dev-tools/audius-compose b/dev-tools/audius-compose
@@ -10,13 +10,22 @@ import eth_account
 import solana.keypair
 
 
-def generate_env(protocol_dir, creator_node_replicas):
+def generate_env(
+    protocol_dir,
+    creator_node_replicas,
+    discovery_provider_replicas,
+    elasticsearch_replicas,
+):
     config = {}
     if (protocol_dir / "dev-tools/config.json").exists():
         config = json.load((protocol_dir / "dev-tools/config.json").open())
 
     env = {}
 
+    env["CREATOR_NODE_REPLICAS"] = str(creator_node_replicas)
+
+    env["DISCOVERY_PROVIDER_REPLICAS"] = str(discovery_provider_replicas)
+
     env["CONTENT_NODE_VERSION"] = json.loads(
         (protocol_dir / "creator-node/.version.json").read_text(),
     )["version"]
@@ -25,6 +34,12 @@ def generate_env(protocol_dir, creator_node_replicas):
         (protocol_dir / "discovery-provider/.version.json").read_text(),
     )["version"]
 
+    env["ELASTICSEARCH_REPLICAS"] = str(elasticsearch_replicas)
+    if elasticsearch_replicas:
+        env["ELASTICSEARCH_CONDITION"] = "service_healthy"
+    else:  # exists to prevent discovery provider from waiting for non-existent elasticsearch instances
+        env["ELASTICSEARCH_CONDITION"] = "service_started"
+
     for name, secret_key in config.get("solana-accounts", {}).items():
         keypair = solana.keypair.Keypair.from_secret_key(bytes(secret_key))
         env[f"{name}_SECRET_KEY"] = json.dumps(list(keypair.secret_key))
@@ -104,9 +119,21 @@ def cli(ctx, protocol_dir):
 
 @cli.command()
 @click.option("-c", "--creator-node-replicas", default=3, type=int)
+@click.option("-d", "--discovery-provider-replicas", default=3, type=int)
+@click.option("-e", "--elasticsearch-replicas", default=0, type=int)
 @click.pass_obj
-def build(protocol_dir, creator_node_replicas):
-    generate_env(protocol_dir, creator_node_replicas)
+def build(
+    protocol_dir,
+    creator_node_replicas,
+    discovery_provider_replicas,
+    elasticsearch_replicas,
+):
+    generate_env(
+        protocol_dir,
+        creator_node_replicas,
+        discovery_provider_replicas,
+        elasticsearch_replicas,
+    )
 
     subprocess.run(
         [
@@ -122,9 +149,20 @@ def build(protocol_dir, creator_node_replicas):
 @cli.command()
 @click.option("-c", "--creator-node-replicas", default=3, type=int)
 @click.option("-d", "--discovery-provider-replicas", default=3, type=int)
+@click.option("-e", "--elasticsearch-replicas", default=0, type=int)
 @click.pass_obj
-def up(protocol_dir, creator_node_replicas, discovery_provider_replicas):
-    generate_env(protocol_dir, creator_node_replicas)
+def up(
+    protocol_dir,
+    creator_node_replicas,
+    discovery_provider_replicas,
+    elasticsearch_replicas,
+):
+    generate_env(
+        protocol_dir,
+        creator_node_replicas,
+        discovery_provider_replicas,
+        elasticsearch_replicas,
+    )
 
     subprocess.run(
         [
@@ -134,10 +172,6 @@ def up(protocol_dir, creator_node_replicas, discovery_provider_replicas):
             protocol_dir,
             "up",
             "--build",
-            "--scale",
-            f"creator-node={creator_node_replicas}",
-            "--scale",
-            f"discovery-provider={discovery_provider_replicas}",
             "-d",
         ],
     )
@@ -212,7 +246,7 @@ def ps(protocol_dir):
 
     print(
         "CONTAINER ID".ljust(13),
-        "NAME".ljust(25),
+        "NAME".ljust(35),
         "STATUS".ljust(10),
         "PORTS",
     )
@@ -237,13 +271,15 @@ def ps(protocol_dir):
         if service["Service"] == "discovery-provider":
             name = f"{service['Service']}-{replica}"
             ports[5000 + replica - 1] = 5000
+        if service["Service"] == "discovery-provider-elasticsearch":
+            name = f"{service['Service']}-{replica}"
 
         ports = sorted(ports.items())
         ports = ", ".join(f"{target}->{published}" for target, published in ports)
 
         print(
             service["ID"][:12].ljust(13),
-            name.ljust(25),
+            name.ljust(35),
             status.ljust(10),
             ports,
         )

diff --git a/dev-tools/startup/discovery-provider.sh b/dev-tools/startup/discovery-provider.sh
@@ -7,5 +7,11 @@ export audius_discprov_url="http://audius-protocol-discovery-provider-${replica}
 export audius_delegate_owner_wallet=$(printenv "DP${replica}_DELEGATE_OWNER_ADDRESS")
 export audius_delegate_private_key=$(printenv "DP${replica}_DELEGATE_OWNER_PRIVATE_KEY")
 
+elasticsearch_host=$(nslookup $(hostname -i) | grep -o "name = .*" | grep -o "[^ ]\+$" | sed 's/discovery-provider/discovery-provider-elasticsearch/')
+if nslookup "$elasticsearch_host" >/dev/null 2>&1; then
+    export audius_elasticsearch_url="http://$elasticsearch_host:9200"
+    export audius_elasticsearch_run_indexer="true"
+fi
+
 # Run register script in background as it waits for the node to be healthy
 ./scripts/register.py &
diff --git a/docker-compose.yml b/docker-compose.yml
@@ -292,6 +292,26 @@ services:
 
   # discovery-provider
 
+  discovery-provider-elasticsearch:
+    image: docker.elastic.co/elasticsearch/elasticsearch:8.1.0
+    environment:
+      discovery.type: single-node
+      xpack.security.enabled: false
+      ES_JAVA_OPTS: -Xms512m -Xmx512m
+    healthcheck:
+      test:
+        [
+          "CMD-SHELL",
+          "curl --fail http://localhost:9200/_cluster/health || exit 1"
+        ]
+      interval: 10s
+      timeout: 5s
+      retries: 15
+    logging: *default-logging
+    deploy:
+      mode: replicated
+      replicas: "${ELASTICSEARCH_REPLICAS}"
+
   discovery-provider:
     build: discovery-provider
     command: sh -c ". /tmp/dev-tools/startup/startup.sh && scripts/start.sh"
@@ -318,7 +338,9 @@ services:
 
       audius_discprov_dev_mode: "true"
     volumes:
-      - ./discovery-provider:/audius-discovery-provider
+      - ./discovery-provider/alembic:/audius-discovery-provider/alembic
+      - ./discovery-provider/solana-tx-parser:/audius-discovery-provider/solana-tx-parser
+      - ./discovery-provider/src:/audius-discovery-provider/src
       - poa-contracts-abis:/audius-discovery-provider/build/contracts
       - eth-contracts-abis:/audius-discovery-provider/build/eth-contracts
       - solana-programs-idl:/audius-discovery-provider/idl
@@ -330,9 +352,12 @@ services:
         condition: service_healthy
       solana-test-validator:
         condition: service_healthy
+      discovery-provider-elasticsearch:
+        condition: "${ELASTICSEARCH_CONDITION}"
     logging: *default-logging
     deploy:
       mode: replicated
+      replicas: "${DISCOVERY_PROVIDER_REPLICAS}"
 
   # creator-node
 
@@ -380,6 +405,9 @@ services:
     logging: *default-logging
     deploy:
       mode: replicated
+      replicas: "${CREATOR_NODE_REPLICAS}"
+
+  # port forwarder
 
   proxy:
     image: alpine:3.16.0