// main.go

// 代码仅供参考,无法直接运行.

package main

import (

"bytes"

"encoding/csv"

"fmt"

MQTT "github.com/eclipse/paho.mqtt.golang"

"github.com/myzhan/boomer"

"io"

"io/ioutil"

"log"

"os"

"strconv"

"strings"

"sync"

"time"

)

var rows [][]string // 读取csv文件保存到这里

var clientTopic []map[string]MQTT.Client

var conn = 0 // 调试用

var failCount = 0 // 初始化失败数量

var i = 0 // 控制并发

var j = 1 // 记录消息发送成功

var f = 1 // 记录消息发送失败

var nowStr = strconv.Itoa(int(time.Now().Unix())) // 当前时间戳,用来做后续查询的消息的标识符

func newConn(c MQTT.Client, clientId string, group *sync.WaitGroup) {

defer func() {

group.Add(-1)

err := recover()

if err != nil {

failCount++

fmt.Println("login fail clientId: ", clientId)

}

}()

token := c.Connect()

if token.Wait() && token.Error() != nil {

panic(token.Error())

}

// 组装topic

topic := fmt.Sprintf("msg/%s/supply", clientId)

temp := make(map[string]MQTT.Client)

temp[topic] = c

clientTopic = append(clientTopic, temp)

conn++ // 调试用

}

func initClients() {

var wg sync.WaitGroup

server := "server_ip:1883"

for i := 0; i < len(rows); i++ {

wg.Add(1)

clientId, userName, passWord := rows[i][0], rows[i][1], rows[i][2]

opts := MQTT.NewClientOptions().AddBroker(server)

opts.SetUsername(userName)

opts.SetPassword(passWord)

opts.SetClientID(clientId)

opts.SetKeepAlive(300 * time.Second)

c := MQTT.NewClient(opts)

go newConn(c, clientId, &wg)

}

wg.Wait() // 等到所有协程执行完成

fmt.Printf("init finish, clients len is %d \n", len(clientTopic))

fmt.Printf("conn: %d \n", conn)

fmt.Printf("failCount: %d \n", failCount)

}

func initCsvData() {

pwd, _ := os.Getwd()

b, err := ioutil.ReadFile(pwd + "/clients.csv")

fs := bytes.NewBuffer(b)

if err != nil {

log.Fatalf("can not open the file, err is %+v", err)

}

r := csv.NewReader(fs)

//针对大文件，一行一行的读取文件

for {

row, err := r.Read()

if err != nil && err != io.EOF {

log.Fatalf("can not read, err is %+v", err)

}

if err == io.EOF {

break

}

rows = append(rows, row)

}

}

func login() {

server := "server_ip:port"

clientId, userName, passWord := rows[i][0], rows[i][1], rows[i][2]

start := time.Now()

opts := MQTT.NewClientOptions().AddBroker(server)

opts.SetUsername(userName)

opts.SetPassword(passWord)

opts.SetClientID(clientId)

c := MQTT.NewClient(opts)

token := c.Connect()

elapsed := time.Since(start)

if token.Error() == nil {

log.Println("success" + strconv.Itoa(j))

boomer.RecordSuccess("tcp", "login", elapsed.Nanoseconds()/int64(time.Millisecond), int64(10))

} else {

log.Println(token.Error())

boomer.RecordFailure("tcp", "login", elapsed.Nanoseconds()/int64(time.Millisecond), clientId)

}

c.Disconnect(5)

// avoid out of array

if i < len(clientTopic)-1 {

i++

} else {

i = 0

}

j++

}

func sendMsg() {

start := time.Now()

msgId := "msg" + strconv.Itoa(i)

var clientId string

var topic string

var c MQTT.Client

for k, v := range clientTopic[i] {

clientId = k[6:19]

topic = k

c = v // v就是一个connected的client

}

deviceTime := nowStr

str := []string{msgId, clientId, deviceTime}

msgPayload := strings.Join(str, "|")

if c.IsConnected() == true {

token := c.Publish(topic, 1, false, msgPayload)

token.Wait() 等待消息发送完成,虽然会拉低并发,但必须要这么做,确保消息发送成功

elapsed := time.Since(start)

if token.Error() == nil {

fmt.Printf("this topic name is: %s \n", topic)

fmt.Printf("this topic payload is: %s \n", msgPayload)

fmt.Printf("success msg index: %v elapsed: %v \n", j, elapsed)

j++ // 消息发送成功, 记录一条,并且也给locust记录一条,方便后续校对数据量

boomer.RecordSuccess("tcp", "task", elapsed.Nanoseconds()/int64(time.Millisecond), int64(j))

// 避免数组越界

if i < len(clientTopic)-1 {

i++

} else {

i = 0

}

} else {

boomer.RecordFailure("tcp", "task", elapsed.Nanoseconds()/int64(time.Millisecond), msgPayload)

fmt.Printf("发送失败, fail msg index: %v \n", f)

}

} else {

if token := c.Connect(); token.Wait() && token.Error() != nil {

elapsed := time.Since(start)

fmt.Printf("fail msg index: %v \n", f)

f++

boomer.RecordFailure("tcp", "task", elapsed.Nanoseconds()/int64(time.Millisecond), msgPayload)

}

}

}

func main() {

initCsvData()

initClients()

task1 := &boomer.Task{

Name: "myTask",

Weight: 1,

Fn: sendMsg,

}

//task2 := &boomer.Task{

// Name: "login",

// Weight: 1,

// Fn: login,

//}

boomer.Run(task1)

}