From 6d828485816d36f89bc179d3d12cb0d8c4c6052a Mon Sep 17 00:00:00 2001
From: Sankalpa Timilsina <sankalpatimilsina12@gmail.com>
Date: Fri, 28 Mar 2025 09:44:40 -0500
Subject: [PATCH] Add trace-collector service

---
 config.yml                                 |  4 +++
 docker-compose.yml                         | 23 ++++++++++++++++
 templates/trace-collector/entrypoint.sh.j2 | 31 ++++++++++++++++++++++
 3 files changed, 58 insertions(+)
 create mode 100644 templates/trace-collector/entrypoint.sh.j2

diff --git a/config.yml b/config.yml
index a0e4ed0..06689b3 100644
--- a/config.yml
+++ b/config.yml
@@ -51,6 +51,10 @@ services:
     template_path: templates/file-server
     render_path: dist/file-server
 
+  trace-collector:
+    template_path: templates/trace-collector
+    render_path: dist/trace-collector
+
 # Global variables that are available in all templates
 globals:
   ndnNetwork: /ndn
diff --git a/docker-compose.yml b/docker-compose.yml
index 2d2561c..2b9eece 100644
--- a/docker-compose.yml
+++ b/docker-compose.yml
@@ -220,3 +220,26 @@ services:
       nfd: { condition: service_healthy }
     logging:
       driver: local
+
+  trace-collector:
+    image: ghcr.io/sankalpatimilsina12/trace-collector:20250328
+    init: true
+    network_mode: service:nfd
+    volumes:
+      - trace-collector-ssh:/root/.ssh # SSH keys
+      - trace-collector-dump:/root/dump # Dumps
+      - ${PWD}/dist/trace-collector:/config  # Entrypoint
+      - ${PWD}/dist/file-server/trace-collector-pubkey:/file-server/trace-collector-pubkey  # Publish public key
+    entrypoint: /bin/bash
+    command: /config/entrypoint.sh
+    restart: unless-stopped
+    depends_on:
+      nfd: { condition: service_healthy }
+      file-server: { condition: service_healthy }
+    logging:
+      driver: local
+
+# Persist trace-collector's SSH keys and dump files
+volumes:
+  trace-collector-ssh:
+  trace-collector-dump:
diff --git a/templates/trace-collector/entrypoint.sh.j2 b/templates/trace-collector/entrypoint.sh.j2
new file mode 100644
index 0000000..1411d11
--- /dev/null
+++ b/templates/trace-collector/entrypoint.sh.j2
@@ -0,0 +1,31 @@
+#!/bin/bash
+set -e
+
+# Variables
+SITE="{{ site }}"
+SITE_CLEAN=$(echo "$SITE" | sed 's#^/##; s#/#_#g') # Eg. /edu/memphis get converted to edu_memphis
+SSH_DIR="/root/.ssh"
+PRIVATE_KEY="$SSH_DIR/id_rsa"
+PUBLIC_KEY="$SSH_DIR/id_rsa.pub"
+PUBKEY_EXPORT_PATH="/file-server/trace-collector-pubkey/id_rsa.pub"
+
+# Ensure permissions
+mkdir -p "$SSH_DIR"
+chmod 700 "$SSH_DIR"
+
+# Generate SSH key pair if not already present
+if [[ ! -f "$PRIVATE_KEY" || ! -f "$PUBLIC_KEY" ]]; then
+    echo "[trace-collector] No SSH key pair found. Generating a new one..."
+    ssh-keygen -t rsa -b 4096 -N "" -f "$PRIVATE_KEY"
+else
+    echo "[trace-collector] SSH key pair already exists."
+fi
+
+# Always export the public key
+mkdir -p "$(dirname "$PUBKEY_EXPORT_PATH")"
+cp "$PUBLIC_KEY" "$PUBKEY_EXPORT_PATH"
+echo "[trace-collector] Public key copied to $PUBKEY_EXPORT_PATH"
+
+# Run
+echo "[trace-collector] Starting trace collector..."
+exec python3.10 /app/scheduler.py --site "$SITE_CLEAN"
\ No newline at end of file