|
@@ -3,12 +3,14 @@ package service
|
|
|
import (
|
|
|
"context"
|
|
|
"esproxy/internal/consts"
|
|
|
+ "fmt"
|
|
|
"github.com/gogf/gf/v2/frame/g"
|
|
|
"github.com/gogf/gf/v2/os/glog"
|
|
|
"net/http/httputil"
|
|
|
"net/url"
|
|
|
"os"
|
|
|
"sync"
|
|
|
+ "time"
|
|
|
)
|
|
|
|
|
|
var (
|
|
@@ -16,9 +18,10 @@ var (
|
|
|
)
|
|
|
|
|
|
type esProxyManager struct {
|
|
|
- lock sync.Mutex
|
|
|
- esStatus consts.EsStatus
|
|
|
- proxyPool chan *httputil.ReverseProxy
|
|
|
+ lock sync.Mutex
|
|
|
+ esStatus consts.EsStatus
|
|
|
+ proxyPool chan *httputil.ReverseProxy
|
|
|
+ simpleQueryPool, aggsQueryPool, complexQueryPool chan struct{}
|
|
|
}
|
|
|
|
|
|
// initManager 初始化代理Manager
|
|
@@ -39,9 +42,12 @@ func initManager(ctx context.Context) *esProxyManager {
|
|
|
}
|
|
|
|
|
|
emp := &esProxyManager{
|
|
|
- lock: sync.Mutex{},
|
|
|
- esStatus: consts.EsStatus_Free,
|
|
|
- proxyPool: pool,
|
|
|
+ lock: sync.Mutex{},
|
|
|
+ esStatus: consts.EsStatus_Free,
|
|
|
+ proxyPool: pool,
|
|
|
+ simpleQueryPool: makeEmptyChan(g.Cfg().MustGet(ctx, "elasticSearch.pool.simple", 10).Int()),
|
|
|
+ aggsQueryPool: makeEmptyChan(g.Cfg().MustGet(ctx, "elasticSearch.pool.aggs", 10).Int()),
|
|
|
+ complexQueryPool: makeEmptyChan(g.Cfg().MustGet(ctx, "elasticSearch.pool.complex", 10).Int()),
|
|
|
}
|
|
|
|
|
|
go emp.UpdateEsStatus(ctx)
|
|
@@ -49,6 +55,14 @@ func initManager(ctx context.Context) *esProxyManager {
|
|
|
return emp
|
|
|
}
|
|
|
|
|
|
+func makeEmptyChan(total int) chan struct{} {
|
|
|
+ c := make(chan struct{}, total)
|
|
|
+ for i := 0; i < total; i++ {
|
|
|
+ c <- struct{}{}
|
|
|
+ }
|
|
|
+ return c
|
|
|
+}
|
|
|
+
|
|
|
// UpdateEsStatus 更新es状态
|
|
|
func (epm *esProxyManager) UpdateEsStatus(ctx context.Context) {
|
|
|
if cron := g.Cfg().MustGet(ctx, "elasticSearch.queryState.cron").String(); cron != "" {
|
|
@@ -66,10 +80,32 @@ func (epm *esProxyManager) UpdateEsStatus(ctx context.Context) {
|
|
|
|
|
|
// GetProxy 获取代理对象
|
|
|
func (epm *esProxyManager) GetProxy(ctx context.Context, queryLevel int) (*httputil.ReverseProxy, error) {
|
|
|
- //epm.lock.Lock()
|
|
|
- //defer epm.lock.Unlock()
|
|
|
+ if queryLevel+int(epm.esStatus) > 2 {
|
|
|
+ return nil, fmt.Errorf("server is busy")
|
|
|
+ }
|
|
|
+ select {
|
|
|
+ case <-epm.getPool(queryLevel):
|
|
|
+ c := <-epm.proxyPool
|
|
|
+ epm.proxyPool <- c
|
|
|
+ return c, nil
|
|
|
+ case <-time.After(time.Second * time.Duration(g.Cfg().MustGet(ctx, "elasticSearch.queryState.waitTime", 5).Int())):
|
|
|
+ return nil, fmt.Errorf("wait time out")
|
|
|
+ }
|
|
|
+}
|
|
|
+
|
|
|
+// Release 归还链接池
|
|
|
+func (epm *esProxyManager) Release(ctx context.Context, queryLevel int) {
|
|
|
+ epm.getPool(queryLevel) <- struct{}{}
|
|
|
+}
|
|
|
|
|
|
- c := <-epm.proxyPool
|
|
|
- epm.proxyPool <- c
|
|
|
- return c, nil
|
|
|
+// 获取对应查询复杂度的链接池
|
|
|
+func (epm *esProxyManager) getPool(queryLevel int) chan struct{} {
|
|
|
+ switch queryLevel {
|
|
|
+ case int(consts.QueryLevelSimple):
|
|
|
+ return epm.simpleQueryPool
|
|
|
+ case int(consts.QueryLevelAggs):
|
|
|
+ return epm.aggsQueryPool
|
|
|
+ default:
|
|
|
+ return epm.complexQueryPool
|
|
|
+ }
|
|
|
}
|