From 5425dfbd637e2aa5363198f5b805894ad07ffbb3 Mon Sep 17 00:00:00 2001
From: Klaus Ma <klaus1982.cn@gmail.com>
Date: Sun, 15 Dec 2019 23:17:11 +0800
Subject: [PATCH] Added hosts into environment.

Signed-off-by: Klaus Ma <klaus1982.cn@gmail.com>
---
 pkg/controllers/job/plugins/svc/const.go |  2 ++
 pkg/controllers/job/plugins/svc/svc.go   | 42 ++++++++++++++++++++----
 2 files changed, 38 insertions(+), 6 deletions(-)

diff --git a/pkg/controllers/job/plugins/svc/const.go b/pkg/controllers/job/plugins/svc/const.go
index 526500468ae..8d385f27114 100644
--- a/pkg/controllers/job/plugins/svc/const.go
+++ b/pkg/controllers/job/plugins/svc/const.go
@@ -19,6 +19,8 @@ package svc
 const (
 	// ConfigMapTaskHostFmt key in config map
 	ConfigMapTaskHostFmt = "%s.host"
+	EnvTaskHostFmt       = "VC_%s_HOSTS"
+	EnvTaskNumFmt        = "VC_%s_NUM"
 
 	// ConfigMapMountPath mount path
 	ConfigMapMountPath = "/etc/volcano"
diff --git a/pkg/controllers/job/plugins/svc/svc.go b/pkg/controllers/job/plugins/svc/svc.go
index 54c79a47a81..b1404c27898 100644
--- a/pkg/controllers/job/plugins/svc/svc.go
+++ b/pkg/controllers/job/plugins/svc/svc.go
@@ -19,6 +19,7 @@ package svc
 import (
 	"flag"
 	"fmt"
+	"strconv"
 	"strings"
 
 	"k8s.io/klog"
@@ -43,6 +44,9 @@ type servicePlugin struct {
 
 	// flag parse args
 	publishNotReadyAddresses bool
+
+	// the host info of each task
+	hostEnv map[string]string
 }
 
 // New creates service plugin
@@ -88,6 +92,19 @@ func (sp *servicePlugin) OnPodCreate(pod *v1.Pod, job *batch.Job) error {
 		pod.Spec.Subdomain = job.Name
 	}
 
+	var hostEnv []v1.EnvVar
+	for k, v := range sp.hostEnv {
+		hostEnv = append(hostEnv, v1.EnvVar{Name: k, Value: v})
+	}
+
+	for i := range pod.Spec.Containers {
+		pod.Spec.Containers[i].Env = append(pod.Spec.Containers[i].Env, hostEnv...)
+	}
+
+	for i := range pod.Spec.InitContainers {
+		pod.Spec.InitContainers[i].Env = append(pod.Spec.InitContainers[i].Env, hostEnv...)
+	}
+
 	sp.mountConfigmap(pod, job)
 
 	return nil
@@ -98,9 +115,13 @@ func (sp *servicePlugin) OnJobAdd(job *batch.Job) error {
 		return nil
 	}
 
-	data := generateHost(job)
+	hostFile, hostEnv := generateHosts(job)
 
-	if err := helpers.CreateConfigMapIfNotExist(job, sp.Clientset.KubeClients, data, sp.cmName(job)); err != nil {
+	// Add hosts as environment when creating pods.
+	sp.hostEnv = hostEnv
+
+	// Create ConfigMap of hosts for Pods to mount.
+	if err := helpers.CreateConfigMapIfNotExist(job, sp.Clientset.KubeClients, hostFile, sp.cmName(job)); err != nil {
 		return err
 	}
 
@@ -258,8 +279,9 @@ func (sp *servicePlugin) cmName(job *batch.Job) string {
 	return fmt.Sprintf("%s-%s", job.Name, sp.Name())
 }
 
-func generateHost(job *batch.Job) map[string]string {
-	data := make(map[string]string, len(job.Spec.Tasks))
+func generateHosts(job *batch.Job) (hostFile map[string]string, hostEnv map[string]string) {
+	hostFile = make(map[string]string, len(job.Spec.Tasks))
+	hostEnv = make(map[string]string, len(job.Spec.Tasks))
 
 	for _, ts := range job.Spec.Tasks {
 		hosts := make([]string, 0, ts.Replicas)
@@ -280,8 +302,16 @@ func generateHost(job *batch.Job) map[string]string {
 		}
 
 		key := fmt.Sprintf(ConfigMapTaskHostFmt, ts.Name)
-		data[key] = strings.Join(hosts, "\n")
+		hostFile[key] = strings.Join(hosts, "\n")
+
+		// TODO(k82cn): The splitter and the prefix of env should be configurable.
+		// export hosts as environment
+		key = fmt.Sprintf(EnvTaskHostFmt, strings.ToUpper(ts.Name))
+		hostFile[key] = strings.Join(hosts, ",")
+		// export host number as environment.
+		key = fmt.Sprintf(EnvTaskNumFmt, strings.ToUpper(ts.Name))
+		hostFile[key] = strconv.Itoa(len(hosts))
 	}
 
-	return data
+	return
 }