ES生成测试数据并写入阿里云ES

weixin_457827192024-10-24 13:33

利用go语言的并发能力并发生成测试数据写入es

ES生成测试数据并写入阿里云ES

go 复制代码 
package main

import (
	"bufio"
	"context"
	"encoding/json"
	"fmt"
	"log"
	"math"
	"math/rand"
	"net/http" // 重新添加 net/http 包
	"os"
	"path/filepath"
	"strconv"
	"sync"
	"time"

	"github.com/brianvoe/gofakeit/v6"
	"github.com/olivere/elastic/v7"
)

// 定义数据结构
type Address struct {
	Street  string `json:"street"`
	City    string `json:"city"`
	Zipcode string `json:"zipcode"`
}

type Record struct {
	UserID     string  `json:"user_id"`
	Name       string  `json:"name"`
	Email      string  `json:"email"`
	Age        int     `json:"age"`
	SignupDate string  `json:"signup_date"`
	Address    Address `json:"address"`
}

// 配置常量
const (
	TargetSizeGB          = 10                  // 目标数据大小（GB）
	OutputDir             = "output_data"       // 输出文件目录
	Workers               = 8                   // 数据生成器数量
	BufferChannelSize     = 10000               // 通道缓冲区大小
	MaxFileSizeBytes      = 500 * 1024 * 1024   // 每个文件最大500MB
	AvgRecordSizeBytes    = 500                 // 估计每条记录500字节
	ProgressPrintMB       = 100                 // 每100MB打印一次进度
	ProgressGenInterval   = 100000              // 每10万条记录生成器打印一次进度
	BulkSize              = 500                 // 每个 Bulk 请求的记录数，减半以提高成功率
	WorkerCount           = 8                   // 并发的工作线程数
	MaxRetries            = 3                   // 最大重试次数
	InitialBackoff        = 2 * time.Second     // 初始退避时间
	ElasticsearchURL      = "http://es-cn-jtexxxxx54r31.public.elasticsearch.aliyuncs.com:9200" // 使用 HTTP 和 9200 端口
	ElasticsearchUsername = "elastic"           // 替换为您的用户名
	ElasticsearchPassword = "xxxxxx56789"      // 替换为您的密码
	IndexName             = "es_index"          // 替换为您的索引名称
)

func init() {
	gofakeit.Seed(0)
	rand.Seed(time.Now().UnixNano())
}

// 生成一条记录
func generateRecord() Record {
	return Record{
		UserID:     gofakeit.UUID(),
		Name:       gofakeit.Name(),
		Email:      gofakeit.Email(),
		Age:        rand.Intn(63) + 18, // 18-80岁
		SignupDate: gofakeit.Date().Format("2006-01-02"),
		Address: Address{
			Street:  gofakeit.Street(),
			City:    gofakeit.City(),
			Zipcode: gofakeit.Zip(),
		},
	}
}

// 创建 Elasticsearch 客户端（使用 HTTP 和 9200 端口）
func createESClient() (*elastic.Client, error) {
	// 创建默认的 HTTP 客户端，增加超时时间至60秒
	httpClient := &http.Client{
		Timeout: 60 * time.Second, // 增加超时时间至60秒
	}

	// 创建 Elasticsearch 客户端
	client, err := elastic.NewClient(
		elastic.SetURL(ElasticsearchURL),
		elastic.SetBasicAuth(ElasticsearchUsername, ElasticsearchPassword),
		elastic.SetSniff(false), // 在云服务中通常禁用 sniff
		elastic.SetHttpClient(httpClient), // 使用自定义的 http.Client
	)

	if err != nil {
		return nil, fmt.Errorf("创建 Elasticsearch 客户端失败: %v", err)
	}

	// 检查连接
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	_, _, err = client.Ping(ElasticsearchURL).Do(ctx)
	if err != nil {
		return nil, fmt.Errorf("无法连接 Elasticsearch: %v", err)
	}

	fmt.Println("Elasticsearch 客户端连接成功")
	return client, nil
}

// 获取所有 JSON 文件路径
func getJSONFiles(dir string) ([]string, error) {
	var files []string
	err := filepath.Walk(dir, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		if !info.IsDir() && filepath.Ext(path) == ".json" {
			files = append(files, path)
		}
		return nil
	})
	return files, err
}

// 读取单个文件中的所有记录
func readRecords(filePath string, recordsChan chan<- string) error {
	file, err := os.Open(filePath)
	if err != nil {
		return err
	}
	defer file.Close()

	scanner := bufio.NewScanner(file)
	count := 0
	for scanner.Scan() {
		recordsChan <- scanner.Text()
		count++
		if count%100000 == 0 {
			log.Printf("读取文件 %s 的记录数: %d", filePath, count)
		}
	}

	if err := scanner.Err(); err != nil {
		return err
	}
	log.Printf("完成读取文件 %s，共读取 %d 条记录", filePath, count)
	return nil
}

// 执行 Bulk 请求并处理响应
func executeBulk(client *elastic.Client, bulkRequest *elastic.BulkService) error {
	// 使用带有超时的上下文
	ctx, cancel := context.WithTimeout(context.Background(), 120*time.Second) // 设置上下文超时时间
	defer cancel()

	var err error
	backoff := InitialBackoff

	for attempt := 1; attempt <= MaxRetries; attempt++ {
		bulkResponse, err := bulkRequest.Do(ctx)
		if err != nil {
			log.Printf("Bulk 请求第 %d 次失败: %v", attempt, err)
			time.Sleep(backoff)
			backoff *= 2
			continue
		}

		if bulkResponse.Errors {
			hasError := false
			for _, item := range bulkResponse.Items {
				for action, result := range item {
					if result.Error != nil {
						hasError = true
						log.Printf("Bulk 操作失败: %s %s: %v", action, result.Id, result.Error)
					}
				}
			}
			if hasError {
				err = fmt.Errorf("Bulk 请求包含错误")
				log.Printf("Bulk 请求第 %d 次包含错误", attempt)
				time.Sleep(backoff)
				backoff *= 2
				continue
			}
		}

		// 成功执行
		return nil
	}

	return fmt.Errorf("Bulk 请求失败，已达到最大重试次数: %v", err)
}

// 写入器负责将记录写入文件（数据生成部分）
func writer(wg *sync.WaitGroup, recordsChan <-chan Record, filePathChan <-chan string) {
	defer wg.Done()

	var file *os.File
	var writerBuf *bufio.Writer
	var currentSize int64 = 0
	var currentFilePath string
	var nextProgress int64 = ProgressPrintMB * 1024 * 1024

	for record := range recordsChan {
		// 检查是否需要切换文件
		if currentFilePath == "" || currentSize >= MaxFileSizeBytes {
			// 获取新的文件路径
			fp, ok := <-filePathChan
			if !ok {
				log.Println("文件路径通道已关闭，无法获取新的文件路径")
				break
			}

			// 关闭旧文件
			if writerBuf != nil {
				writerBuf.Flush()
			}
			if file != nil {
				file.Close()
			}

			// 打开新文件
			var err error
			file, err = os.Create(fp)
			if err != nil {
				log.Printf("创建文件 %s 失败: %v\n", fp, err)
				continue
			}
			writerBuf = bufio.NewWriter(file)
			currentFilePath = fp
			currentSize = 0
			nextProgress = ProgressPrintMB * 1024 * 1024
			log.Printf("开始写入文件: %s", fp)
		}

		// 序列化记录
		recordBytes, err := json.Marshal(record)
		if err != nil {
			log.Printf("序列化记录失败: %v\n", err)
			continue
		}
		recordStr := string(recordBytes) + "\n"

		// 写入文件
		_, err = writerBuf.WriteString(recordStr)
		if err != nil {
			log.Printf("写入文件失败: %v\n", err)
			continue
		}
		currentSize += int64(len(recordStr))

		// 打印进度
		if currentSize >= nextProgress {
			log.Printf("已写入 %s 大小: %.2f MB", currentFilePath, float64(currentSize)/(1024*1024))
			nextProgress += ProgressPrintMB * 1024 * 1024
		}
	}

	// 写入剩余的数据
	if writerBuf != nil {
		writerBuf.Flush()
	}
	if file != nil {
		file.Close()
	}
	log.Println("写入器完成所有文件的写入")
}

// 数据生成函数
func generateData() {
	start := time.Now()
	log.Println("开始生成数据...")

	// 创建输出目录
	if _, err := os.Stat(OutputDir); os.IsNotExist(err) {
		err := os.Mkdir(OutputDir, os.ModePerm)
		if err != nil {
			log.Fatalf("创建输出目录失败: %v", err)
		}
		log.Printf("创建输出目录: %s", OutputDir)
	} else {
		log.Printf("输出目录已存在: %s", OutputDir)
	}

	// 计算总记录数
	totalBytes := int64(TargetSizeGB) * 1024 * 1024 * 1024
	totalRecords := totalBytes / AvgRecordSizeBytes
	recordsPerWorker := totalRecords / int64(Workers)
	remainder := totalRecords % int64(Workers)

	// 计算每个文件能容纳的记录数
	recordsPerFile := MaxFileSizeBytes / AvgRecordSizeBytes
	totalFiles := int(math.Ceil(float64(totalRecords) / float64(recordsPerFile)))

	log.Printf("总记录数: %d", totalRecords)
	log.Printf("每个文件记录数: %d", recordsPerFile)
	log.Printf("总文件数: %d", totalFiles)

	// 创建通道
	recordsChan := make(chan Record, BufferChannelSize)
	filePathChan := make(chan string, totalFiles) // 足够的缓冲区存放所有文件路径
	var wg sync.WaitGroup

	// 启动写入器
	wg.Add(1)
	go writer(&wg, recordsChan, filePathChan)

	// 生成所有文件路径
	go func() {
		for i := 1; i <= totalFiles; i++ {
			filePath := filepath.Join(OutputDir, "test_data_part_"+strconv.Itoa(i)+".json")
			filePathChan <- filePath
		}
		close(filePathChan) // 所有文件路径发送完毕后关闭通道
		log.Println("所有文件路径已生成并发送到通道")
	}()

	// 启动生成器
	var genWg sync.WaitGroup
	for i := 0; i < Workers; i++ {
		genWg.Add(1)
		go func(workerID int, numRecords int64) {
			defer genWg.Done()
			for j := int64(0); j < numRecords; j++ {
				record := generateRecord()
				recordsChan <- record
				if (j+1)%ProgressGenInterval == 0 {
					log.Printf("生成器 %d 已生成 %d 条记录", workerID, j+1)
				}
			}
			log.Printf("生成器 %d 完成生成 %d 条记录", workerID, numRecords)
		}(i+1, recordsPerWorker)
	}

	// 处理余数
	if remainder > 0 {
		genWg.Add(1)
		go func() {
			defer genWg.Done()
			for j := int64(0); j < remainder; j++ {
				record := generateRecord()
				recordsChan <- record
			}
			log.Printf("生成器 %d 完成生成余数 %d 条记录", Workers+1, remainder)
		}()
	}

	// 等待生成器完成
	genWg.Wait()
	close(recordsChan)
	log.Println("所有生成器已完成，关闭 recordsChan")

	// 等待写入器完成
	wg.Wait()

	elapsed := time.Since(start)
	log.Printf("数据生成完成，总耗时: %s", elapsed)
}

// 导入 Elasticsearch 的函数
func importToES() {
	start := time.Now()
	log.Println("开始导入数据到 Elasticsearch...")

	client, err := createESClient()
	if err != nil {
		log.Fatalf("无法连接 Elasticsearch: %v", err)
	}

	// 获取所有 JSON 文件
	files, err := getJSONFiles(OutputDir)
	if err != nil {
		log.Fatalf("无法读取 JSON 文件: %v", err)
	}

	if len(files) == 0 {
		log.Fatalf("没有找到任何 JSON 文件在目录: %s", OutputDir)
	}

	log.Printf("找到 %d 个 JSON 文件用于导入", len(files))

	recordsChan := make(chan string, 10000)
	var wg sync.WaitGroup

	// 启动读取文件的 goroutine
	go func() {
		for _, file := range files {
			log.Printf("开始读取文件: %s", file)
			if err := readRecords(file, recordsChan); err != nil {
				log.Printf("读取文件 %s 失败: %v", file, err)
			}
			log.Printf("完成读取文件: %s", file)
		}
		close(recordsChan)
		log.Println("所有文件读取完成，关闭 recordsChan")
	}()

	// 启动 Worker
	for i := 0; i < WorkerCount; i++ {
		wg.Add(1)
		go func(workerID int) {
			defer wg.Done()
			var bulkRequest *elastic.BulkService
			bulkRequest = client.Bulk()
			count := 0
			totalImported := 0

			for record := range recordsChan {
				var doc map[string]interface{}
				if err := json.Unmarshal([]byte(record), &doc); err != nil {
					log.Printf("Worker %d 解析 JSON 失败: %v", workerID, err)
					continue
				}

				// 确保 user_id 字段存在且为 string 类型
				userID, ok := doc["user_id"].(string)
				if !ok {
					log.Printf("Worker %d 无法获取有效的 user_id", workerID)
					continue
				}

				req := elastic.NewBulkIndexRequest().Index(IndexName).Id(userID).Doc(doc)
				bulkRequest = bulkRequest.Add(req)
				count++

				if count >= BulkSize {
					// 执行 Bulk 请求
					if err := executeBulk(client, bulkRequest); err != nil {
						log.Printf("Worker %d 执行 Bulk 请求失败: %v", workerID, err)
					} else {
						totalImported += count
						log.Printf("Worker %d 成功导入 %d 条记录，累计导入 %d 条", workerID, count, totalImported)
					}
					// 重置 Bulk 请求
					bulkRequest = client.Bulk()
					count = 0
				}

				if totalImported%100000 == 0 && totalImported > 0 {
					log.Printf("Worker %d 已累计导入 %d 条记录", workerID, totalImported)
				}
			}

			// 执行剩余的 Bulk 请求
			if count > 0 {
				if err := executeBulk(client, bulkRequest); err != nil {
					log.Printf("Worker %d 执行 Bulk 请求失败: %v", workerID, err)
				} else {
					totalImported += count
					log.Printf("Worker %d 成功导入剩余的 %d 条记录，累计导入 %d 条", workerID, count, totalImported)
				}
			}

			log.Printf("Worker %d 完成数据导入，总导入 %d 条记录", workerID, totalImported)
		}(i + 1)
	}

	// 等待所有 Worker 完成
	wg.Wait()

	elapsed := time.Since(start)
	log.Printf("数据导入完成，总耗时: %s", elapsed)
	log.Println("所有数据已成功导入 Elasticsearch")
}

func main() {
	// 配置日志输出格式
	log.SetFlags(log.LstdFlags | log.Lshortfile)

	// 根据功能分离，先生成数据，再导入ES
	generateData()
	importToES()
}

主要更新和优化

  1. 增加详细的日志输出：

    • 在数据读取、写入和导入的各个关键步骤添加了日志输出，帮助跟踪进度和排查问题。
    • 使用 log.Printf 记录关键事件，如开始和完成读取文件、生成器生成记录、写入文件以及每个 Worker 导入的进度。

  2. 调整 Bulk 请求的超时时间：

    • 将 executeBulk 函数中的上下文超时时间从 60 秒增加到 120 秒，以防止由于大量数据导致的请求超时。

  3. 校验和报错加强：

    • 在 readRecords 函数中记录每个文件读取的进度和完成情况。
    • 在 importToES 的 Worker 中记录每次 Bulk 请求的成功和失败情况，以及累计导入的记录数。

  4. 优化并发和通道处理：

    • 确保 filePathChan 在所有文件路径发送完毕后关闭，以防止写入器阻塞。
    • 在生成数据的过程中，确保所有生成器完成后才关闭 recordsChan。

  5. 日志格式配置：

    • 在 main 函数中配置日志输出格式，以包括时间戳和文件行号，便于调试。

运行和监控

运行此程序时，您将看到类似以下的日志输出：

2024/04/27 10:00:00 main.go:123: 开始生成数据...
2024/04/27 10:00:00 main.go:133: 创建输出目录: output_data
2024/04/27 10:00:00 main.go:142: 总记录数: 21474836480
2024/04/27 10:00:00 main.go:143: 每个文件记录数: 1048576
2024/04/27 10:00:00 main.go:144: 总文件数: 20480
2024/04/27 10:00:00 main.go:156: 开始写入文件: output_data/test_data_part_1.json
2024/04/27 10:00:05 main.go:172: 生成器 1 已生成 100000 条记录
...
2024/04/27 10:02:00 main.go:198: 数据生成完成，总耗时: 2m0s
2024/04/27 10:02:00 main.go:201: 开始导入数据到 Elasticsearch...
2024/04/27 10:02:00 main.go:61: Elasticsearch 客户端连接成功
2024/04/27 10:02:00 main.go:217: 找到 20480 个 JSON 文件用于导入
2024/04/27 10:02:00 main.go:227: 开始读取文件: output_data/test_data_part_1.json
2024/04/27 10:02:30 main.go:237: 完成读取文件: output_data/test_data_part_1.json，共读取 1000000 条记录
...
2024/04/27 10:05:00 main.go:259: Worker 1 成功导入 500 条记录，累计导入 500 条
...
2024/04/27 10:10:00 main.go:291: Worker 1 完成数据导入，总导入 500000 条记录
...
2024/04/27 10:20:00 main.go:303: 数据导入完成，总耗时: 18m20s
2024/04/27 10:20:00 main.go:304: 所有数据已成功导入 Elasticsearch

通过这些日志，您可以：

  • 跟踪数据生成过程： 了解每个生成器的进度和完成情况。
  • 监控文件写入过程： 确认每个文件的写入进度和完成状态。
  • 监测数据导入过程： 追踪每个 Worker 导入的记录数，以及每次 Bulk 请求的处理情况。

进一步优化建议

  1. 调整 Bulk 请求大小和 Worker 数量：

    • 根据 Elasticsearch 集群的性能，您可以尝试增大或减小 BulkSize，以及调整 WorkerCount 以优化导入速度。
    • 例如，将 BulkSize 增加到 1000 或 2000，观察其对导入性能的影响。

  2. 监控 Elasticsearch 性能：

    • 在导入过程中，监控 Elasticsearch 集群的 CPU、内存、磁盘 I/O 和网络状况，确保集群能够处理高负载。
    • 使用 Elasticsearch 提供的监控工具，比如 Kibana 的监控插件，来实时查看集群的健康状态和性能指标。

  3. 处理导入过程中可能出现的错误：

    • 如果遇到频繁的 Bulk 请求错误，可以根据错误日志分析原因，比如索引模板不匹配、映射冲突等，并进行相应调整。

  4. 优化索引设置：

    • 临时关闭刷新间隔： 在大规模导入时，可以暂时关闭自动刷新，以提高导入性能。导入完成后，再手动触发刷新。

      go 复制代码 
      // 在导入前临时调整索引设置
_, err := client.IndexPutSettings(IndexName).BodyJson(map[string]interface{}{
    "index": map[string]interface{}{
        "refresh_interval": "-1",
    },
}).Do(context.Background())
if err != nil {
    log.Fatalf("无法更新索引设置: %v", err)
}

// 导入完成后恢复刷新间隔
_, err = client.IndexPutSettings(IndexName).BodyJson(map[string]interface{}{
    "index": map[string]interface{}{
        "refresh_interval": "1s",
    },
}).Do(context.Background())
if err != nil {
    log.Fatalf("无法恢复索引设置: %v", err)
}

    • 增加主分片数量： 根据数据量和查询需求，调整索引的主分片数量，以优化性能。

  5. 分批导入：

    • 如果数据量过大，可以分批次导入，每次导入一部分数据，观察其对集群性能的影响。

总结

通过上述更新和优化，能够更好地监控和管理数据生成与导入的过程，及时发现并解决可能的问题。如果导入过程仍然出现卡顿，请查看日志中的详细信息，以便进一步分析和定位问题。

相关推荐
Pandaconda14 分钟前
【Golang 面试题】每日 3 题（三十九）
开发语言·经验分享·笔记·后端·面试·golang·go
加油，旭杏18 分钟前
【go语言】变量和常量
服务器·开发语言·golang
行路见知18 分钟前
3.3 Go 返回值详解
开发语言·golang
编程小筑1 小时前
R语言的编程范式
开发语言·后端·golang
技术的探险家1 小时前
Elixir语言的文件操作
开发语言·后端·golang
Ai 编码助手1 小时前
Golang 中强大的重试机制，解决瞬态错误
开发语言·后端·golang
齐雅彤2 小时前
Lisp语言的区块链
开发语言·后端·golang
齐雅彤2 小时前
Lisp语言的循环实现
开发语言·后端·golang
梁雨珈2 小时前
Lisp语言的物联网
开发语言·后端·golang
邓熙榆3 小时前
Logo语言的网络编程
开发语言·后端·golang
热门推荐
01xgboost: Why not implement distributed XGBoost on top of spark02Dell服务器升级ubuntu 22.04失败解决03半导体应用系统一些小知识收集(strip&wafer mapping,EAP&scada)04密码学原理技术-第六章-introduction to pulibc-key cryptography05(欧拉)openEuler系统添加网卡文件配置流程、(欧拉)openEuler系统手动配置ipv6地址流程、(欧拉)openEuler系统网络管理说明06优化手机性能，解决卡顿问题：关闭这3个微信开关，释放内存空间07ChatGPT + Stable Diffusion + 百度AI + MoviePy 实现文字生成视频，小说转视频，自媒体神器！(二)08Windows10安装PCL1.14.0及点云配准09渗透测试之SQLMAP工具详解 kali自带SQLmap解释 重点sqlmap --tamper 使用方式详解 搞完你就很nice了10JS设计模式之中介者模式