mirror of
https://github.com/httprunner/httprunner.git
synced 2026-09-08 17:09:33 +08:00
move ghdc to pkg
This commit is contained in:
lilong.129
@@ -0,0 +1,345 @@
|
||||
package ghdc
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"container/list"
|
||||
"encoding/binary"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net"
|
||||
"strconv"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
type uitestKitTransport struct {
|
||||
connectionPool *ConnectionPool
|
||||
socketMap map[string]*SocketContext
|
||||
serial string
|
||||
mu sync.Mutex
|
||||
}
|
||||
|
||||
type SocketContext struct {
|
||||
conn net.Conn
|
||||
socketId string
|
||||
writeLock sync.Mutex
|
||||
|
||||
callbackMap map[string]UitestKitCallback
|
||||
queue *responseList
|
||||
}
|
||||
|
||||
type response struct {
|
||||
sessionId uint32
|
||||
payload []byte
|
||||
}
|
||||
|
||||
type responseList struct {
|
||||
list *list.List
|
||||
lock sync.Mutex
|
||||
}
|
||||
|
||||
func (r *responseList) Clear() {
|
||||
r.lock.Lock()
|
||||
defer r.lock.Unlock()
|
||||
r.list.Init()
|
||||
}
|
||||
|
||||
func (r *responseList) push(v interface{}) {
|
||||
defer r.lock.Unlock()
|
||||
r.lock.Lock()
|
||||
r.list.PushBack(v)
|
||||
}
|
||||
|
||||
func (r *responseList) traverse(action func(response response)) {
|
||||
r.lock.Lock()
|
||||
defer r.lock.Unlock()
|
||||
for e := r.list.Front(); e != nil; e = e.Next() {
|
||||
action(e.Value.(response))
|
||||
}
|
||||
}
|
||||
|
||||
func (r *responseList) remove(sessionId uint32) *response {
|
||||
r.lock.Lock()
|
||||
defer r.lock.Unlock()
|
||||
for e := r.list.Front(); e != nil; e = e.Next() {
|
||||
response := e.Value.(response)
|
||||
if response.sessionId == sessionId {
|
||||
r.list.Remove(e)
|
||||
return &response
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
type UitestKitCallback interface {
|
||||
OnData([]byte)
|
||||
OnError(error)
|
||||
}
|
||||
|
||||
type ReqTypeEnum int
|
||||
|
||||
const (
|
||||
DEFAULT ReqTypeEnum = iota
|
||||
SCREEN_CAPTURE
|
||||
UI_ACTION_CAPTURE
|
||||
)
|
||||
|
||||
const (
|
||||
HEADER_BYTES = "_uitestkit_rpc_message_head_"
|
||||
TAILER_BYTES = "_uitestkit_rpc_message_tail_"
|
||||
)
|
||||
|
||||
func newUitestKitTransport(serial string, host string, port string) (uKtp uitestKitTransport, err error) {
|
||||
pool, err := newConnectionPool(host, port, 3)
|
||||
if err != nil {
|
||||
err = fmt.Errorf("[uitest] failed to init connection pool \n%v", err)
|
||||
return
|
||||
}
|
||||
uKtp.connectionPool = pool
|
||||
uKtp.socketMap = make(map[string]*SocketContext)
|
||||
uKtp.serial = serial
|
||||
return
|
||||
}
|
||||
|
||||
func (uKtp *uitestKitTransport) initializeSocket(reqType ReqTypeEnum) error {
|
||||
socketId := fmt.Sprintf("socket_%d_%s", reqType, uKtp.serial)
|
||||
uKtp.mu.Lock()
|
||||
defer uKtp.mu.Unlock()
|
||||
if uKtp.socketMap[socketId] != nil {
|
||||
return nil
|
||||
}
|
||||
connection, err := uKtp.connectionPool.getConnection()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
uKtp.socketMap[socketId] = &SocketContext{socketId: socketId, conn: connection, callbackMap: make(map[string]UitestKitCallback), queue: newResponseList()}
|
||||
go uKtp.receive(reqType)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (uKtp *uitestKitTransport) disconnect(reqType ReqTypeEnum) {
|
||||
socketId := fmt.Sprintf("socket_%d_%s", reqType, uKtp.serial)
|
||||
uKtp.mu.Lock()
|
||||
defer uKtp.mu.Unlock()
|
||||
socketContext := uKtp.socketMap[socketId]
|
||||
if socketContext == nil {
|
||||
return
|
||||
}
|
||||
socketContext.Close()
|
||||
delete(uKtp.socketMap, socketId)
|
||||
}
|
||||
|
||||
func (uKtp *uitestKitTransport) registerCallback(reqType ReqTypeEnum, sessionId uint32, callback UitestKitCallback) error {
|
||||
socketId := fmt.Sprintf("socket_%d_%s", reqType, uKtp.serial)
|
||||
socketContext := uKtp.socketMap[socketId]
|
||||
if socketContext == nil {
|
||||
if err := uKtp.initializeSocket(reqType); err != nil {
|
||||
return err
|
||||
}
|
||||
socketContext = uKtp.socketMap[socketId]
|
||||
}
|
||||
for {
|
||||
socketContext.writeLock.Lock()
|
||||
res := socketContext.queue.remove(sessionId)
|
||||
socketContext.writeLock.Unlock()
|
||||
if res != nil && callback != nil {
|
||||
callback.OnData(res.payload)
|
||||
}
|
||||
if res == nil {
|
||||
break
|
||||
}
|
||||
}
|
||||
socketContext.writeLock.Lock()
|
||||
socketContext.callbackMap[strconv.Itoa(int(sessionId))] = callback
|
||||
socketContext.writeLock.Unlock()
|
||||
return nil
|
||||
}
|
||||
|
||||
func (uKtp *uitestKitTransport) receive(reqType ReqTypeEnum) {
|
||||
socketId := fmt.Sprintf("socket_%d_%s", reqType, uKtp.serial)
|
||||
defer uKtp.disconnect(reqType)
|
||||
socketContext := uKtp.socketMap[socketId]
|
||||
var receiveError error
|
||||
defer func() {
|
||||
if receiveError != nil {
|
||||
socketContext.onException(receiveError)
|
||||
}
|
||||
}()
|
||||
if socketContext == nil {
|
||||
if receiveError = uKtp.initializeSocket(reqType); receiveError != nil {
|
||||
return
|
||||
}
|
||||
socketContext = uKtp.socketMap[socketId]
|
||||
}
|
||||
|
||||
for {
|
||||
headerSize := len(HEADER_BYTES)
|
||||
raw, err := _readN(socketContext.conn, headerSize+8)
|
||||
if err != nil {
|
||||
receiveError = err
|
||||
break
|
||||
}
|
||||
if len(raw) == 0 {
|
||||
continue
|
||||
}
|
||||
|
||||
header := raw[:len(HEADER_BYTES)]
|
||||
if !bytes.Equal(header, []byte(HEADER_BYTES)) {
|
||||
receiveError = fmt.Errorf("verify message head failed on channel: %s", socketId)
|
||||
break
|
||||
}
|
||||
var sessionId, length uint32
|
||||
if receiveError = binary.Read(bytes.NewReader(raw[headerSize:headerSize+4]), binary.BigEndian, &sessionId); receiveError != nil {
|
||||
break
|
||||
}
|
||||
if receiveError = binary.Read(bytes.NewReader(raw[headerSize+4:headerSize+8]), binary.BigEndian, &length); receiveError != nil {
|
||||
break
|
||||
}
|
||||
payload, err := _readN(socketContext.conn, int(length))
|
||||
if err != nil {
|
||||
receiveError = err
|
||||
break
|
||||
}
|
||||
tail, err := _readN(socketContext.conn, len(TAILER_BYTES))
|
||||
if err != nil {
|
||||
receiveError = err
|
||||
break
|
||||
}
|
||||
if !bytes.Equal(tail, []byte(TAILER_BYTES)) {
|
||||
receiveError = fmt.Errorf("verify message tail failed on channel: %s", socketId)
|
||||
break
|
||||
}
|
||||
socketContext.writeLock.Lock()
|
||||
callback := socketContext.callbackMap[strconv.Itoa(int(sessionId))]
|
||||
socketContext.writeLock.Unlock()
|
||||
if callback != nil {
|
||||
callback.OnData(payload)
|
||||
} else {
|
||||
socketContext.queue.push(response{sessionId: sessionId, payload: payload})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (uKtp *uitestKitTransport) sendMessage(reqType ReqTypeEnum, sessionId uint32, message string) (response UitestKitResponse, err error) {
|
||||
defer func() {
|
||||
if err != nil {
|
||||
uKtp.disconnect(reqType)
|
||||
}
|
||||
}()
|
||||
if err = uKtp._sendMessage(reqType, sessionId, message); err != nil {
|
||||
return
|
||||
}
|
||||
return uKtp.receiveMessage(reqType, sessionId)
|
||||
}
|
||||
|
||||
func (uKtp *uitestKitTransport) _sendMessage(reqType ReqTypeEnum, sessionId uint32, message string) (err error) {
|
||||
socketId := fmt.Sprintf("socket_%d_%s", reqType, uKtp.serial)
|
||||
socketContext := uKtp.socketMap[socketId]
|
||||
if socketContext == nil {
|
||||
if err = uKtp.initializeSocket(reqType); err != nil {
|
||||
return
|
||||
}
|
||||
socketContext = uKtp.socketMap[socketId]
|
||||
}
|
||||
buffer := new(bytes.Buffer)
|
||||
if err = binary.Write(buffer, binary.BigEndian, []byte(HEADER_BYTES)); err != nil {
|
||||
return
|
||||
}
|
||||
if err = binary.Write(buffer, binary.BigEndian, sessionId); err != nil {
|
||||
return
|
||||
}
|
||||
if err = binary.Write(buffer, binary.BigEndian, uint32(len(message))); err != nil {
|
||||
return
|
||||
}
|
||||
if err = binary.Write(buffer, binary.BigEndian, []byte(message)); err != nil {
|
||||
return
|
||||
}
|
||||
if err = binary.Write(buffer, binary.BigEndian, []byte(TAILER_BYTES)); err != nil {
|
||||
return
|
||||
}
|
||||
socketContext.writeLock.Lock()
|
||||
defer socketContext.writeLock.Unlock()
|
||||
return _send(socketContext.conn, buffer.Bytes())
|
||||
}
|
||||
|
||||
func (sc *SocketContext) Close() {
|
||||
sc.writeLock.Lock()
|
||||
defer sc.writeLock.Unlock()
|
||||
|
||||
if sc.conn != nil {
|
||||
_ = sc.conn.Close()
|
||||
}
|
||||
|
||||
// 清空callbackMap
|
||||
for key := range sc.callbackMap {
|
||||
delete(sc.callbackMap, key)
|
||||
}
|
||||
sc.callbackMap = nil
|
||||
|
||||
// 清理队列
|
||||
if sc.queue != nil {
|
||||
sc.queue.Clear()
|
||||
}
|
||||
}
|
||||
|
||||
func (uKtp *uitestKitTransport) Close() {
|
||||
uKtp.mu.Lock()
|
||||
defer uKtp.mu.Unlock()
|
||||
|
||||
// 关闭所有的SocketContext
|
||||
if uKtp.socketMap != nil {
|
||||
for _, socketContext := range uKtp.socketMap {
|
||||
socketContext.Close()
|
||||
}
|
||||
}
|
||||
|
||||
// 关闭连接池
|
||||
if uKtp.connectionPool != nil {
|
||||
uKtp.connectionPool.close()
|
||||
}
|
||||
|
||||
uKtp.socketMap = nil
|
||||
uKtp.connectionPool = nil
|
||||
}
|
||||
|
||||
func (uKtp *uitestKitTransport) receiveMessage(reqType ReqTypeEnum, sessionId uint32) (response UitestKitResponse, err error) {
|
||||
socketId := fmt.Sprintf("socket_%d_%s", reqType, uKtp.serial)
|
||||
socketContext := uKtp.socketMap[socketId]
|
||||
if socketContext == nil {
|
||||
err = fmt.Errorf("failed to read message. not found target connection")
|
||||
return
|
||||
}
|
||||
// 创建一个计时器，设置超时时间为 10 秒
|
||||
timeout := time.After(2 * time.Second)
|
||||
|
||||
// 创建一个 ticker，每秒触发一次
|
||||
ticker := time.NewTicker(10 * time.Millisecond)
|
||||
defer ticker.Stop()
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-timeout:
|
||||
err = fmt.Errorf("failed to read message in 2 second")
|
||||
return
|
||||
case <-ticker.C:
|
||||
res := socketContext.queue.remove(sessionId)
|
||||
if res != nil {
|
||||
err = json.Unmarshal(res.payload, &response)
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (sc *SocketContext) onException(err error) {
|
||||
for _, callback := range sc.callbackMap {
|
||||
if callback != nil {
|
||||
callback.OnError(err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func newResponseList() *responseList {
|
||||
return &responseList{list: list.New()}
|
||||
}
|
||||
Reference in New Issue
Block a user