From 50932b40ae2d671935e062fd7f60f6834a86bdbd Mon Sep 17 00:00:00 2001
From: 9547 <nivefive9547@gmail.com>
Date: Mon, 29 Mar 2021 22:17:24 +0800
Subject: [PATCH] docker: support --nodes (#1251)

---
 docker/docker-compose.yml.tpl | 47 ++++++-----------------------------
 docker/up.sh                  | 24 +++++++++++++++++-
 2 files changed, 31 insertions(+), 40 deletions(-)

diff --git a/docker/docker-compose.yml.tpl b/docker/docker-compose.yml.tpl
index 206f57cfd5..804e80ffe9 100644
--- a/docker/docker-compose.yml.tpl
+++ b/docker/docker-compose.yml.tpl
@@ -19,49 +19,18 @@ services:
       - "8080"
     networks:
       tiops:
-        ipv4_address: __IPPREFIX__.100
-  n1:
+        ipv4_address: {{ipprefix}}.100
+{% for id in range(1, nodes+1) %}
+  n{{id}}:
     <<: *default-node
-    container_name: tiup-cluster-n1
-    hostname: n1
-    # Uncomment for access grafana and prometheus from host when deploy them at n1.
-    # networks:
-    # ports:
-    #   - "3000:3000"
-    #   - "9090:9090"
+    container_name: tiup-cluster-n{{id}}
+    hostname: n{{id}}
     networks:
       tiops:
-        ipv4_address: __IPPREFIX__.101
-  n2:
-    <<: *default-node
-    container_name: tiup-cluster-n2
-    hostname: n2
-    networks:
-      tiops:
-        ipv4_address: __IPPREFIX__.102
-  n3:
-    <<: *default-node
-    container_name: tiup-cluster-n3
-    hostname: n3
-    networks:
-      tiops:
-        ipv4_address: __IPPREFIX__.103
-  n4:
-    <<: *default-node
-    container_name: tiup-cluster-n4
-    hostname: n4
-    networks:
-      tiops:
-        ipv4_address: __IPPREFIX__.104
-  n5:
-    <<: *default-node
-    container_name: tiup-cluster-n5
-    hostname: n5
-    networks:
-      tiops:
-        ipv4_address: __IPPREFIX__.105
+        ipv4_address: {{ipprefix}}.{{id+100}}
+{% endfor %}
 
-# docker network create --gateway __IPPREFIX__.1 --subnet __IPPREFIX__.0/24 tiup-cluster
+# docker network create --gateway {{ipprefix}}.1 --subnet {{ipprefix}}.0/24 tiup-cluster
 networks:
   tiops:
     external: true
diff --git a/docker/up.sh b/docker/up.sh
index 9526d3b18d..71926a6928 100755
--- a/docker/up.sh
+++ b/docker/up.sh
@@ -35,6 +35,7 @@ INIT_ONLY=0
 DEV=""
 COMPOSE=${COMPOSE:-""}
 SUBNET=${SUBNET:-"172.19.0.0/24"}
+NODES=${NODES:-5}
 RUN_AS_DAEMON=0
 POSITIONAL=()
 
@@ -70,6 +71,11 @@ do
             shift # past argument
             shift # past value
             ;;
+        -n|--nodes)
+            NODES="$2"
+            shift # past argument
+            shift # past value
+            ;;
         -d|--daemon)
             INFO "Running docker-compose as daemon"
             RUN_AS_DAEMON=1
@@ -95,6 +101,7 @@ if [ "${HELP}" -eq 1 ]; then
     echo "  --dev                                                 Mounts dir at host's TIUP_CLUSTER_ROOT to /tiup-cluster on tiup-cluster-control container, syncing files for development"
     echo "  --compose PATH                                        Path to an additional docker-compose yml config."
     echo "  --subnet SUBNET                                       Subnet in 24 bit netmask"
+    echo "  --nodes NODES                                         Start how much nodes"
     echo "To provide multiple additional docker-compose args, set the COMPOSE var directly, with the -f flag. Ex: COMPOSE=\"-f FILE_PATH_HERE -f ANOTHER_PATH\" ./up.sh --dev"
     exit 0
 fi
@@ -149,6 +156,14 @@ if [ ${SUBNET##*/} -ne 24 ]; then
     exit 1
 fi
 
+if [ "$NODES" -gt "64" ]; then
+    ERROR "At most 64 nodes is supported"
+    exit 1
+fi
+
+exists python ||
+    { ERROR "Please install python (https://www.python.org/downloads/)";
+      exit 1; }
 exists docker ||
     { ERROR "Please install docker (https://docs.docker.com/engine/installation/)";
       exit 1; }
@@ -156,6 +171,13 @@ exists docker-compose ||
     { ERROR "Please install docker-compose (https://docs.docker.com/compose/install/)";
       exit 1; }
 
+exists pip ||
+    {
+        INFO "Install pip from https://bootstrap.pypa.io/get-pip.py";
+        curl -sSL https://bootstrap.pypa.io/get-pip.py -o get-pip.py || exit 1;
+    }
+pip install -U jinja2
+
 exist_network=$(docker network ls | awk '{if($2 == "tiops") print $1}')
 if [[ "$exist_network" == "" ]]; then
     ipprefix=${SUBNET%.*}
@@ -170,7 +192,7 @@ else
     ipprefix=${SUBNET%.*}
 fi
 
-sed "s/__IPPREFIX__/$ipprefix/g" docker-compose.yml.tpl > docker-compose.yml
+python -c "from jinja2 import Template; print(Template(open('docker-compose.yml.tpl').read()).render(nodes=$NODES, ipprefix='$ipprefix'))" > docker-compose.yml
 sed "s/__IPPREFIX__/$ipprefix/g" docker-compose.dm.yml.tpl > docker-compose.dm.yml
 sed -i '/TIUP_TEST_IP_PREFIX/d' ./secret/control.env
 echo "TIUP_TEST_IP_PREFIX=$ipprefix" >> ./secret/control.env