From 453d1893d8ade9675a6761e27264f6751169a314 Mon Sep 17 00:00:00 2001
From: hyphennn <hyphennn@163.com>
Date: Wed, 17 Jul 2024 18:05:23 +0800
Subject: [PATCH] feature: support multi nodes and basic auth config for es

---
 .../storage/elasticsearch/elasticsearch.go    | 33 +++++++++++++++----
 docs/storage/elasticsearch.md                 |  8 ++++-
 2 files changed, 33 insertions(+), 8 deletions(-)

diff --git a/cmd/internal/storage/elasticsearch/elasticsearch.go b/cmd/internal/storage/elasticsearch/elasticsearch.go
index 40871a0e94..a1e39669e9 100644
--- a/cmd/internal/storage/elasticsearch/elasticsearch.go
+++ b/cmd/internal/storage/elasticsearch/elasticsearch.go
@@ -18,6 +18,7 @@ import (
 	"flag"
 	"fmt"
 	"os"
+	"strings"
 	"sync"
 	"time"
 
@@ -51,6 +52,8 @@ var (
 	argIndexName     = flag.String("storage_driver_es_index", "cadvisor", "ElasticSearch index name")
 	argTypeName      = flag.String("storage_driver_es_type", "stats", "ElasticSearch type name")
 	argEnableSniffer = flag.Bool("storage_driver_es_enable_sniffer", false, "ElasticSearch uses a sniffing process to find all nodes of your cluster by default, automatically")
+	argUserName      = flag.String("storage_driver_es_username", "", "ElasticSearch basic auth username")
+	argPassword      = flag.String("storage_driver_es_password", "", "ElasticSearch basic auth password")
 )
 
 func new() (storage.StorageDriver, error) {
@@ -64,6 +67,8 @@ func new() (storage.StorageDriver, error) {
 		*argTypeName,
 		*argElasticHost,
 		*argEnableSniffer,
+		*argUserName,
+		*argPassword,
 	)
 }
 
@@ -124,7 +129,13 @@ func newStorage(
 	typeName,
 	elasticHost string,
 	enableSniffer bool,
+	username string,
+	password string,
 ) (storage.StorageDriver, error) {
+	// Remove all spaces to help user to configure
+	elasticHost = strings.ReplaceAll(elasticHost, " ", "")
+	hosts := strings.Split(elasticHost, ",")
+
 	// Obtain a client and connect to the default Elasticsearch installation
 	// on 127.0.0.1:9200. Of course you can configure your client to connect
 	// to other hosts and configure it in various other ways.
@@ -132,7 +143,8 @@ func newStorage(
 		elastic.SetHealthcheck(true),
 		elastic.SetSniff(enableSniffer),
 		elastic.SetHealthcheckInterval(30*time.Second),
-		elastic.SetURL(elasticHost),
+		elastic.SetURL(hosts...),
+		elastic.SetBasicAuth(username, password),
 	)
 	if err != nil {
 		// Handle error
@@ -140,13 +152,20 @@ func newStorage(
 	}
 
 	// Ping the Elasticsearch server to get e.g. the version number
-	info, code, err := client.Ping().URL(elasticHost).Do()
-	if err != nil {
-		// Handle error
-		return nil, fmt.Errorf("failed to ping the elasticsearch - %s", err)
-
+	// Just ping anyone of hosts successfully will be ok
+	var res *elastic.PingResult
+	var code int
+	for _, host := range hosts {
+		res, code, err = client.Ping().URL(host).Do()
+		if err == nil {
+			break
+		}
+		fmt.Printf("ping host %s failed, code: %d, err: %s", host, code, err)
+	}
+	if res == nil {
+		return nil, fmt.Errorf("failed to ping any host of the elasticsearch")
 	}
-	fmt.Printf("Elasticsearch returned with code %d and version %s", code, info.Version.Number)
+	fmt.Printf("Elasticsearch returned with code %d and version %s", code, res.Version.Number)
 
 	ret := &elasticStorage{
 		client:      client,
diff --git a/docs/storage/elasticsearch.md b/docs/storage/elasticsearch.md
index fa679cf497..f67a9d1bc1 100644
--- a/docs/storage/elasticsearch.md
+++ b/docs/storage/elasticsearch.md
@@ -1,6 +1,7 @@
 # Exporting cAdvisor Stats to ElasticSearch
 
-cAdvisor supports exporting stats to [ElasticSearch](https://www.elastic.co/). To use ES, you need to provide the additional flags to cAdvisor:
+cAdvisor supports exporting stats to [ElasticSearch](https://www.elastic.co/). To use ES, you need to provide the
+additional flags to cAdvisor:
 
 Set the storage driver as ES:
 
@@ -12,6 +13,8 @@ Specify ES host address:
 
 ```
  -storage_driver_es_host="http://elasticsearch:9200"
+ # If you has several hosts, just use comma to separate it.
+ -storage_driver_es_host="http://elasticsearch1:9200,http://elasticsearch2:9200,http://elasticsearch3:9200"
 ```
 
 There are also optional flags:
@@ -21,6 +24,9 @@ There are also optional flags:
  -storage_driver_es_type="stats"
  # ElasticSearch can use a sniffing process to find all nodes of your cluster automatically. False by default.
  -storage_driver_es_enable_sniffer=false
+ # ElasticSearch basic auth for http request, only works when any one of them is not empty
+ -storage_driver_es_username="xxx"
+ -storage_driver_es_password="xxx"
 ```
 
 # Examples