Update
This commit is contained in:
zun.yang
276
test/test.go
Normal file
276
test/test.go
Normal file
@ -0,0 +1,276 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"archive/zip"
|
||||
"context"
|
||||
"embed"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"log"
|
||||
"math/rand"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/bitfield/script"
|
||||
"github.com/bramvdbogaerde/go-scp"
|
||||
pcmd "github.com/elulcao/progress-bar/cmd"
|
||||
"github.com/sirupsen/logrus"
|
||||
"golang.org/x/crypto/ssh"
|
||||
)
|
||||
|
||||
type envConfig struct {
|
||||
Host string `json:"Host"`
|
||||
Port int `json:"Port"`
|
||||
Username string `json:"Username"`
|
||||
Password string `json:"Password"`
|
||||
Geekbench_License string `json:"Geekbench_License"`
|
||||
}
|
||||
|
||||
// 在本地启动一个 iperf3 服务端
|
||||
func iperfServer() {
|
||||
cmd := exec.Command("benchmark/tools/iperf3", "-s")
|
||||
cmd.Run()
|
||||
}
|
||||
|
||||
func pBar(hsotNum int) {
|
||||
pb := pcmd.NewPBar()
|
||||
pb.SignalHandler()
|
||||
pb.Total = uint16(150 * hsotNum)
|
||||
|
||||
// pb.RenderPBar(i)
|
||||
for i := 1; uint16(i) <= pb.Total; i++ {
|
||||
pb.RenderPBar(i)
|
||||
if uint16(i) == pb.Total-1 {
|
||||
i -= 2
|
||||
}
|
||||
time.Sleep(1 * time.Second)
|
||||
}
|
||||
}
|
||||
|
||||
func Unzip(src string, dest string) error {
|
||||
r, err := zip.OpenReader(src)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer r.Close()
|
||||
|
||||
for _, f := range r.File {
|
||||
rc, err := f.Open()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer rc.Close()
|
||||
|
||||
fpath := filepath.Join(dest, f.Name)
|
||||
if f.FileInfo().IsDir() {
|
||||
os.MkdirAll(fpath, os.ModePerm)
|
||||
} else {
|
||||
var dir string
|
||||
if lastIndex := strings.LastIndex(fpath, string(os.PathSeparator)); lastIndex > -1 {
|
||||
dir = fpath[:lastIndex]
|
||||
}
|
||||
err = os.MkdirAll(dir, os.ModePerm)
|
||||
if err != nil {
|
||||
log.Fatal(err)
|
||||
return err
|
||||
}
|
||||
f, err := os.OpenFile(
|
||||
fpath, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, f.Mode())
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer f.Close()
|
||||
|
||||
_, err = io.Copy(f, rc)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func randomString(length int) string {
|
||||
const charset = "abcdefghijklmnopqrstuvwxyz" +
|
||||
"ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
|
||||
|
||||
var seededRand *rand.Rand = rand.New(
|
||||
rand.NewSource(time.Now().UnixNano()))
|
||||
|
||||
b := make([]byte, length)
|
||||
for i := range b {
|
||||
b[i] = charset[seededRand.Intn(len(charset))]
|
||||
}
|
||||
return string(b)
|
||||
}
|
||||
|
||||
// 封装 远程执行 shell 命令
|
||||
func remoteShellExec(rClient ssh.Client, commands []string, logOut bool) error {
|
||||
outputStr := ""
|
||||
for _, command := range commands {
|
||||
session, err := rClient.NewSession()
|
||||
if err != nil {
|
||||
logrus.Errorf("Failed to create session: %s", err.Error())
|
||||
}
|
||||
defer session.Close()
|
||||
output, err := session.CombinedOutput(command)
|
||||
if err != nil {
|
||||
logrus.Errorf("Failed to run: %s", err.Error())
|
||||
}
|
||||
outputStr = string(output)
|
||||
if logOut {
|
||||
fmt.Println(" ")
|
||||
fmt.Println(outputStr)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// 测试 SSH 目标主机是否可联通
|
||||
func testSSH(rClient ssh.Client) bool {
|
||||
|
||||
session, err := rClient.NewSession()
|
||||
if err != nil {
|
||||
logrus.Error("Failed to dial: ", err)
|
||||
return false
|
||||
}
|
||||
defer session.Close()
|
||||
return true
|
||||
}
|
||||
|
||||
//go:embed benchmark.zip
|
||||
var f embed.FS
|
||||
|
||||
func main() {
|
||||
srcDir := "benchmark"
|
||||
dstDir := "/tmp/" + randomString(10) + "/"
|
||||
|
||||
data, _ := f.ReadFile("benchmark.zip")
|
||||
|
||||
// 将字节切片写入临时zip文件
|
||||
tmpfile, err := os.CreateTemp("", "example")
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
tmpfile.Write(data)
|
||||
tmpfile.Close()
|
||||
|
||||
// 解压文件
|
||||
Unzip(tmpfile.Name(), ".")
|
||||
|
||||
file, err := os.Open("config.json")
|
||||
if err != nil {
|
||||
logrus.Error("Error opening file:", err)
|
||||
return
|
||||
}
|
||||
defer file.Close()
|
||||
envConfig := &envConfig{}
|
||||
decoder := json.NewDecoder(file)
|
||||
err = decoder.Decode(&envConfig)
|
||||
if err != nil {
|
||||
fmt.Println("Error decoding file:", err)
|
||||
return
|
||||
}
|
||||
// 读取本地 config.json 文件, 注入成变量
|
||||
username := envConfig.Username
|
||||
password := envConfig.Password
|
||||
port := envConfig.Port
|
||||
geekbenchLicense := envConfig.Geekbench_License
|
||||
HostExecOut, _ := script.Exec("hostname -I").String()
|
||||
hostIPS := strings.Split(HostExecOut, " ")
|
||||
|
||||
config := &ssh.ClientConfig{
|
||||
User: username,
|
||||
Auth: []ssh.AuthMethod{
|
||||
ssh.Password(password),
|
||||
},
|
||||
HostKeyCallback: ssh.InsecureIgnoreHostKey(),
|
||||
}
|
||||
|
||||
// 打印开始 motto
|
||||
logrus.Info("Benchmarking...")
|
||||
// 获取当前时间
|
||||
start := time.Now() // 获取当前时间
|
||||
|
||||
// 初始化 bar cmd
|
||||
go pBar(len(strings.Split(envConfig.Host, ",")))
|
||||
|
||||
// 创建 iperfServer 协程
|
||||
go iperfServer()
|
||||
|
||||
// 自动关闭 iperfServer 协程
|
||||
defer func() {
|
||||
cmd := exec.Command("killall", "iperf3", " || ", "rm -rf", srcDir)
|
||||
_, err := cmd.Output()
|
||||
if err != nil {
|
||||
logrus.Error("Failed to kill iperf3")
|
||||
}
|
||||
}()
|
||||
|
||||
// 使用 testSSH 函数测试目标主机是否可联通
|
||||
for _, host := range strings.Split(envConfig.Host, ",") {
|
||||
serverURL := fmt.Sprintf("%s:%d", host, port)
|
||||
client, err := ssh.Dial("tcp", serverURL, config)
|
||||
if err != nil {
|
||||
logrus.Error("Failed to dial: ", err)
|
||||
}
|
||||
if testSSH(*client) {
|
||||
logrus.Infof("Host: %s , SSH 连接成功", host)
|
||||
} else {
|
||||
logrus.Errorf("Host: %s , SSH 连接失败", host)
|
||||
}
|
||||
}
|
||||
|
||||
// 按 , 切 config.Host 中的每一个 host
|
||||
for _, host := range strings.Split(envConfig.Host, ",") {
|
||||
serverURL := fmt.Sprintf("%s:%d", host, port)
|
||||
|
||||
client, err := ssh.Dial("tcp", serverURL, config)
|
||||
if err != nil {
|
||||
logrus.Error("Failed to dial: ", err)
|
||||
}
|
||||
|
||||
_ = remoteShellExec(*client, []string{"mkdir -p " + dstDir}, false)
|
||||
|
||||
filepath.Walk(srcDir, func(path string, info os.FileInfo, err error) error {
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if info.IsDir() {
|
||||
remoteShellExec(*client, []string{"mkdir -p " + dstDir + path}, false)
|
||||
} else {
|
||||
scpClient, err := scp.NewClientBySSH(client)
|
||||
scpClient.Connect()
|
||||
file, err := os.Open(path)
|
||||
defer file.Close()
|
||||
if err != nil {
|
||||
logrus.Error("Failed to open file: ", err)
|
||||
}
|
||||
err = scpClient.CopyFile(context.Background(), file, dstDir+path, "0655")
|
||||
if err != nil {
|
||||
logrus.Error("Failed to copy file: ", err)
|
||||
}
|
||||
scpClient.Close()
|
||||
|
||||
}
|
||||
return nil
|
||||
})
|
||||
|
||||
// 执行 远端 脚本
|
||||
logrus.Infof("Host: %s , 正在执行远端脚本中...", host)
|
||||
|
||||
_ = remoteShellExec(*client, []string{
|
||||
"cd " + dstDir + srcDir + " && " + "./run.sh " + hostIPS[0] + " " + "'" + geekbenchLicense + "'",
|
||||
}, true)
|
||||
|
||||
// 删除远端文件
|
||||
remoteShellExec(*client, []string{"rm -rf " + dstDir}, false)
|
||||
}
|
||||
|
||||
// 打印执行时间
|
||||
logrus.Infof("本次总计执行时间: %s", time.Since(start))
|
||||
}
|
||||
Reference in New Issue
Block a user