mirror of
https://github.com/TeaOSLab/EdgeNode.git
synced 2025-11-04 16:00:25 +08:00
483 lines
12 KiB
Go
483 lines
12 KiB
Go
package nodes
|
|
|
|
import (
|
|
"encoding/json"
|
|
"fmt"
|
|
"github.com/TeaOSLab/EdgeCommon/pkg/messageconfigs"
|
|
"github.com/TeaOSLab/EdgeCommon/pkg/rpc/pb"
|
|
"github.com/TeaOSLab/EdgeCommon/pkg/serverconfigs"
|
|
"github.com/TeaOSLab/EdgeNode/internal/caches"
|
|
"github.com/TeaOSLab/EdgeNode/internal/errors"
|
|
"github.com/TeaOSLab/EdgeNode/internal/events"
|
|
"github.com/TeaOSLab/EdgeNode/internal/logs"
|
|
"github.com/TeaOSLab/EdgeNode/internal/rpc"
|
|
"io"
|
|
"net/http"
|
|
"strconv"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
type APIStream struct {
|
|
stream pb.NodeService_NodeStreamClient
|
|
}
|
|
|
|
func NewAPIStream() *APIStream {
|
|
return &APIStream{}
|
|
}
|
|
|
|
func (this *APIStream) Start() {
|
|
isQuiting := false
|
|
events.On(events.EventQuit, func() {
|
|
isQuiting = true
|
|
})
|
|
for {
|
|
if isQuiting {
|
|
return
|
|
}
|
|
err := this.loop()
|
|
if err != nil {
|
|
logs.Error("API_STREAM", err.Error())
|
|
time.Sleep(10 * time.Second)
|
|
continue
|
|
}
|
|
time.Sleep(1 * time.Second)
|
|
}
|
|
}
|
|
|
|
func (this *APIStream) loop() error {
|
|
rpcClient, err := rpc.SharedRPC()
|
|
if err != nil {
|
|
return errors.Wrap(err)
|
|
}
|
|
isQuiting := false
|
|
events.On(events.EventQuit, func() {
|
|
isQuiting = true
|
|
})
|
|
nodeStream, err := rpcClient.NodeRPC().NodeStream(rpcClient.Context())
|
|
if err != nil {
|
|
if isQuiting {
|
|
return nil
|
|
}
|
|
return errors.Wrap(err)
|
|
}
|
|
this.stream = nodeStream
|
|
|
|
for {
|
|
if isQuiting {
|
|
break
|
|
}
|
|
|
|
message, err := nodeStream.Recv()
|
|
if err != nil {
|
|
if isQuiting {
|
|
return nil
|
|
}
|
|
return errors.Wrap(err)
|
|
}
|
|
|
|
// 处理消息
|
|
switch message.Code {
|
|
case messageconfigs.MessageCodeConnectedAPINode: // 连接API节点成功
|
|
err = this.handleConnectedAPINode(message)
|
|
case messageconfigs.MessageCodeWriteCache: // 写入缓存
|
|
err = this.handleWriteCache(message)
|
|
case messageconfigs.MessageCodeReadCache: // 读取缓存
|
|
err = this.handleReadCache(message)
|
|
case messageconfigs.MessageCodeStatCache: // 统计缓存
|
|
err = this.handleStatCache(message)
|
|
case messageconfigs.MessageCodeCleanCache: // 清理缓存
|
|
err = this.handleCleanCache(message)
|
|
case messageconfigs.MessageCodePurgeCache: // 删除缓存
|
|
err = this.handlePurgeCache(message)
|
|
case messageconfigs.MessageCodePreheatCache: // 预热缓存
|
|
err = this.handlePreheatCache(message)
|
|
case messageconfigs.MessageCodeConfigChanged: // 配置变化
|
|
err = this.handleConfigChanged(message)
|
|
default:
|
|
err = this.handleUnknownMessage(message)
|
|
}
|
|
if err != nil {
|
|
logs.Error("API_STREAM", "handle message failed: "+err.Error())
|
|
}
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
// 连接API节点成功
|
|
func (this *APIStream) handleConnectedAPINode(message *pb.NodeStreamMessage) error {
|
|
// 更改连接的APINode信息
|
|
if len(message.DataJSON) == 0 {
|
|
return nil
|
|
}
|
|
msg := &messageconfigs.ConnectedAPINodeMessage{}
|
|
err := json.Unmarshal(message.DataJSON, msg)
|
|
if err != nil {
|
|
return errors.Wrap(err)
|
|
}
|
|
|
|
rpcClient, err := rpc.SharedRPC()
|
|
if err != nil {
|
|
return errors.Wrap(err)
|
|
}
|
|
|
|
_, err = rpcClient.NodeRPC().UpdateNodeConnectedAPINodes(rpcClient.Context(), &pb.UpdateNodeConnectedAPINodesRequest{ApiNodeIds: []int64{msg.APINodeId}})
|
|
if err != nil {
|
|
return errors.Wrap(err)
|
|
}
|
|
logs.Println("API_STREAM", "connected to api node '"+strconv.FormatInt(msg.APINodeId, 10)+"'")
|
|
return nil
|
|
}
|
|
|
|
// 写入缓存
|
|
func (this *APIStream) handleWriteCache(message *pb.NodeStreamMessage) error {
|
|
msg := &messageconfigs.WriteCacheMessage{}
|
|
err := json.Unmarshal(message.DataJSON, msg)
|
|
if err != nil {
|
|
this.replyFail(message.RequestId, "decode message data failed: "+err.Error())
|
|
return err
|
|
}
|
|
|
|
storage, shouldStop, err := this.cacheStorage(message, msg.CachePolicyJSON)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if shouldStop {
|
|
defer func() {
|
|
storage.Stop()
|
|
}()
|
|
}
|
|
|
|
expiredAt := time.Now().Unix() + msg.LifeSeconds
|
|
writer, err := storage.Open(msg.Key, expiredAt)
|
|
if err != nil {
|
|
this.replyFail(message.RequestId, "prepare writing failed: "+err.Error())
|
|
return err
|
|
}
|
|
|
|
_, err = writer.Write(msg.Value)
|
|
if err != nil {
|
|
_ = writer.Discard()
|
|
this.replyFail(message.RequestId, "write failed: "+err.Error())
|
|
return err
|
|
}
|
|
err = writer.Close()
|
|
if err == nil {
|
|
storage.AddToList(&caches.Item{
|
|
Key: msg.Key,
|
|
ExpiredAt: expiredAt,
|
|
})
|
|
}
|
|
|
|
this.replyOk(message.RequestId, "write ok")
|
|
|
|
return nil
|
|
}
|
|
|
|
// 读取缓存
|
|
func (this *APIStream) handleReadCache(message *pb.NodeStreamMessage) error {
|
|
msg := &messageconfigs.ReadCacheMessage{}
|
|
err := json.Unmarshal(message.DataJSON, msg)
|
|
if err != nil {
|
|
this.replyFail(message.RequestId, "decode message data failed: "+err.Error())
|
|
return err
|
|
}
|
|
|
|
storage, shouldStop, err := this.cacheStorage(message, msg.CachePolicyJSON)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if shouldStop {
|
|
defer func() {
|
|
storage.Stop()
|
|
}()
|
|
}
|
|
|
|
buf := make([]byte, 1024)
|
|
size := 0
|
|
err = storage.Read(msg.Key, buf, func(data []byte, valueSize int64, expiredAt int64, isEOF bool) {
|
|
size += len(data)
|
|
})
|
|
if err != nil {
|
|
if err == caches.ErrNotFound {
|
|
this.replyFail(message.RequestId, "key not found")
|
|
return nil
|
|
}
|
|
this.replyFail(message.RequestId, "read key failed: "+err.Error())
|
|
return err
|
|
}
|
|
|
|
this.replyOk(message.RequestId, "value "+strconv.Itoa(size)+" bytes")
|
|
|
|
return nil
|
|
}
|
|
|
|
// 统计缓存
|
|
func (this *APIStream) handleStatCache(message *pb.NodeStreamMessage) error {
|
|
msg := &messageconfigs.ReadCacheMessage{}
|
|
err := json.Unmarshal(message.DataJSON, msg)
|
|
if err != nil {
|
|
this.replyFail(message.RequestId, "decode message data failed: "+err.Error())
|
|
return err
|
|
}
|
|
|
|
storage, shouldStop, err := this.cacheStorage(message, msg.CachePolicyJSON)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if shouldStop {
|
|
defer func() {
|
|
storage.Stop()
|
|
}()
|
|
}
|
|
|
|
stat, err := storage.Stat()
|
|
if err != nil {
|
|
this.replyFail(message.RequestId, "stat failed: "+err.Error())
|
|
return err
|
|
}
|
|
|
|
sizeFormat := ""
|
|
if stat.Size < 1024 {
|
|
sizeFormat = strconv.FormatInt(stat.Size, 10) + " Bytes"
|
|
} else if stat.Size < 1024*1024 {
|
|
sizeFormat = fmt.Sprintf("%.2f KB", float64(stat.Size)/1024)
|
|
} else if stat.Size < 1024*1024*1024 {
|
|
sizeFormat = fmt.Sprintf("%.2f MB", float64(stat.Size)/1024/1024)
|
|
} else {
|
|
sizeFormat = fmt.Sprintf("%.2f GB", float64(stat.Size)/1024/1024/1024)
|
|
}
|
|
this.replyOk(message.RequestId, "size:"+sizeFormat+", count:"+strconv.Itoa(stat.Count))
|
|
|
|
return nil
|
|
}
|
|
|
|
// 清理缓存
|
|
func (this *APIStream) handleCleanCache(message *pb.NodeStreamMessage) error {
|
|
msg := &messageconfigs.ReadCacheMessage{}
|
|
err := json.Unmarshal(message.DataJSON, msg)
|
|
if err != nil {
|
|
this.replyFail(message.RequestId, "decode message data failed: "+err.Error())
|
|
return err
|
|
}
|
|
|
|
storage, shouldStop, err := this.cacheStorage(message, msg.CachePolicyJSON)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if shouldStop {
|
|
defer func() {
|
|
storage.Stop()
|
|
}()
|
|
}
|
|
|
|
err = storage.CleanAll()
|
|
if err != nil {
|
|
this.replyFail(message.RequestId, "clean cache failed: "+err.Error())
|
|
return err
|
|
}
|
|
|
|
this.replyOk(message.RequestId, "ok")
|
|
|
|
return nil
|
|
}
|
|
|
|
// 删除缓存
|
|
func (this *APIStream) handlePurgeCache(message *pb.NodeStreamMessage) error {
|
|
msg := &messageconfigs.PurgeCacheMessage{}
|
|
err := json.Unmarshal(message.DataJSON, msg)
|
|
if err != nil {
|
|
this.replyFail(message.RequestId, "decode message data failed: "+err.Error())
|
|
return err
|
|
}
|
|
|
|
storage, shouldStop, err := this.cacheStorage(message, msg.CachePolicyJSON)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if shouldStop {
|
|
defer func() {
|
|
storage.Stop()
|
|
}()
|
|
}
|
|
|
|
err = storage.Purge(msg.Keys)
|
|
if err != nil {
|
|
this.replyFail(message.RequestId, "purge keys failed: "+err.Error())
|
|
return err
|
|
}
|
|
|
|
this.replyOk(message.RequestId, "ok")
|
|
|
|
return nil
|
|
}
|
|
|
|
// 预热缓存
|
|
func (this *APIStream) handlePreheatCache(message *pb.NodeStreamMessage) error {
|
|
msg := &messageconfigs.PreheatCacheMessage{}
|
|
err := json.Unmarshal(message.DataJSON, msg)
|
|
if err != nil {
|
|
this.replyFail(message.RequestId, "decode message data failed: "+err.Error())
|
|
return err
|
|
}
|
|
|
|
storage, shouldStop, err := this.cacheStorage(message, msg.CachePolicyJSON)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if shouldStop {
|
|
defer func() {
|
|
storage.Stop()
|
|
}()
|
|
}
|
|
|
|
if len(msg.Keys) == 0 {
|
|
this.replyOk(message.RequestId, "ok")
|
|
return nil
|
|
}
|
|
|
|
wg := sync.WaitGroup{}
|
|
wg.Add(len(msg.Keys))
|
|
client := http.Client{} // TODO 可以设置请求超时事件
|
|
errorMessages := []string{}
|
|
locker := sync.Mutex{}
|
|
for _, key := range msg.Keys {
|
|
go func(key string) {
|
|
defer wg.Done()
|
|
|
|
req, err := http.NewRequest(http.MethodGet, key, nil)
|
|
if err != nil {
|
|
locker.Lock()
|
|
errorMessages = append(errorMessages, "invalid url: "+key+": "+err.Error())
|
|
locker.Unlock()
|
|
return
|
|
}
|
|
// TODO 可以在管理界面自定义Header
|
|
req.Header.Set("User-Agent", "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/85.0.4183.121 Safari/537.36")
|
|
req.Header.Set("Accept-Encoding", "gzip, deflate, br") // TODO 这里需要记录下缓存是否为gzip的
|
|
resp, err := client.Do(req)
|
|
if err != nil {
|
|
locker.Lock()
|
|
errorMessages = append(errorMessages, "request failed: "+key+": "+err.Error())
|
|
locker.Unlock()
|
|
return
|
|
}
|
|
|
|
defer func() {
|
|
_ = resp.Body.Close()
|
|
}()
|
|
|
|
expiredAt := time.Now().Unix() + 8600
|
|
writer, err := storage.Open(key, expiredAt) // TODO 可以设置缓存过期事件
|
|
if err != nil {
|
|
locker.Lock()
|
|
errorMessages = append(errorMessages, "open cache writer failed: "+key+": "+err.Error())
|
|
locker.Unlock()
|
|
return
|
|
}
|
|
|
|
buf := make([]byte, 16*1024)
|
|
isClosed := false
|
|
for {
|
|
n, err := resp.Body.Read(buf)
|
|
if n > 0 {
|
|
_, writerErr := writer.Write(buf[:n])
|
|
if writerErr != nil {
|
|
locker.Lock()
|
|
errorMessages = append(errorMessages, "write failed: "+key+": "+writerErr.Error())
|
|
locker.Unlock()
|
|
break
|
|
}
|
|
}
|
|
if err != nil {
|
|
if err == io.EOF {
|
|
err = writer.Close()
|
|
if err == nil {
|
|
storage.AddToList(&caches.Item{
|
|
Key: key,
|
|
ExpiredAt: expiredAt,
|
|
})
|
|
}
|
|
isClosed = true
|
|
} else {
|
|
locker.Lock()
|
|
errorMessages = append(errorMessages, "read url failed: "+key+": "+err.Error())
|
|
locker.Unlock()
|
|
}
|
|
break
|
|
}
|
|
}
|
|
|
|
if !isClosed {
|
|
_ = writer.Close()
|
|
}
|
|
}(key)
|
|
}
|
|
wg.Wait()
|
|
|
|
if len(errorMessages) > 0 {
|
|
this.replyFail(message.RequestId, strings.Join(errorMessages, ", "))
|
|
return nil
|
|
}
|
|
|
|
this.replyOk(message.RequestId, "ok")
|
|
|
|
return nil
|
|
}
|
|
|
|
// 处理配置变化
|
|
func (this *APIStream) handleConfigChanged(message *pb.NodeStreamMessage) error {
|
|
select {
|
|
case changeNotify <- true:
|
|
default:
|
|
|
|
}
|
|
this.replyOk(message.RequestId, "ok")
|
|
return nil
|
|
}
|
|
|
|
// 处理未知消息
|
|
func (this *APIStream) handleUnknownMessage(message *pb.NodeStreamMessage) error {
|
|
this.replyFail(message.RequestId, "unknown message code '"+message.Code+"'")
|
|
return nil
|
|
}
|
|
|
|
// 回复失败
|
|
func (this *APIStream) replyFail(requestId int64, message string) {
|
|
_ = this.stream.Send(&pb.NodeStreamMessage{RequestId: requestId, IsOk: false, Message: message})
|
|
}
|
|
|
|
// 回复成功
|
|
func (this *APIStream) replyOk(requestId int64, message string) {
|
|
_ = this.stream.Send(&pb.NodeStreamMessage{RequestId: requestId, IsOk: true, Message: message})
|
|
}
|
|
|
|
// 获取缓存存取对象
|
|
func (this *APIStream) cacheStorage(message *pb.NodeStreamMessage, cachePolicyJSON []byte) (storage caches.StorageInterface, shouldStop bool, err error) {
|
|
cachePolicy := &serverconfigs.HTTPCachePolicy{}
|
|
err = json.Unmarshal(cachePolicyJSON, cachePolicy)
|
|
if err != nil {
|
|
this.replyFail(message.RequestId, "decode cache policy config failed: "+err.Error())
|
|
return nil, false, err
|
|
}
|
|
|
|
storage = caches.SharedManager.FindStorageWithPolicy(cachePolicy.Id)
|
|
if storage == nil {
|
|
storage = caches.SharedManager.NewStorageWithPolicy(cachePolicy)
|
|
if storage == nil {
|
|
this.replyFail(message.RequestId, "invalid storage type '"+cachePolicy.Type+"'")
|
|
return nil, false, errors.New("invalid storage type '" + cachePolicy.Type + "'")
|
|
}
|
|
err = storage.Init()
|
|
if err != nil {
|
|
this.replyFail(message.RequestId, "storage init failed: "+err.Error())
|
|
return nil, false, err
|
|
}
|
|
shouldStop = true
|
|
}
|
|
|
|
return
|
|
}
|