diff --git a/client.go b/client.go index dcc3b6a..7d08bd4 100644 --- a/client.go +++ b/client.go @@ -1,8 +1,12 @@ package yar import ( + "context" "errors" + "fmt" + "github.com/flyhope/go-yar/logger" "github.com/flyhope/go-yar/pack" + "github.com/flyhope/go-yar/request" "github.com/sirupsen/logrus" "io/ioutil" "net/http" @@ -10,15 +14,22 @@ import ( ) type Client struct { - Request *pack.Request - Response *pack.Response - Http *http.Request - HttpClient *http.Client - LogTrace LogTrace + ctx context.Context + Request *pack.Request + Response *pack.Response + Http *http.Request + RequestClient request.Handler + logTrace logger.LogTrace } // 初始化一个客户端 -func NewClient(addr string, method string, params interface{}) (*Client, error) { +func NewClient(ctx context.Context, addr string, method string, params interface{}) (*Client, error) { + handler := &request.Http{Client: &http.Client{Timeout: time.Second}} + return NewWithHandler(ctx, handler, addr, method, params) +} + +// 初始化一个客户端 +func NewWithHandler(ctx context.Context, requestClient request.Handler, addr string, method string, params interface{}) (*Client, error) { httpRequest, err := http.NewRequest(http.MethodPost, addr, nil) if err != nil { return nil, err @@ -26,15 +37,22 @@ func NewClient(addr string, method string, params interface{}) (*Client, error) httpRequest.Header.Set("User-Agent", "Go Yar Rpc-0.1") c := &Client{ - Request: pack.NewRequest(addr, method, params), - Response: new(pack.Response), - Http: httpRequest, - HttpClient: &http.Client{Timeout: time.Second}, + ctx: ctx, + Request: pack.NewRequest(addr, method, params), + Response: new(pack.Response), + Http: httpRequest, + RequestClient: requestClient, } return c, nil } +// 设置日志追踪方法 +func (c *Client) SetLogTrace(logTrace logger.LogTrace) { + c.logTrace = logTrace + c.RequestClient.SetLog(logTrace) +} + // 设置返回值结构体 func (c *Client) SetResponseRetStruct(retVal interface{}) *Client { c.Response.Retval = retVal @@ -54,33 +72,19 @@ func (c *Client) Send() error { buffer := header.Bytes() buffer.Write(data) + c.Http.Body = ioutil.NopCloser(buffer) c.Http.Header.Set("Content-Type", packHandler.ContentType()) + c.Http.Header.Add("Content-Length", fmt.Sprintf("%d", buffer.Len())) - Log.WithFields(logrus.Fields{"YAR": "Request"}).Debug(string(data)) + logger.Log.WithFields(logrus.Fields{"YAR": "Request"}).Debug(string(data)) // 发送请求 - timeStart := time.Now() - resp, err := c.HttpClient.Do(c.Http) - - // 通过接口记录跟踪日志 - if c.LogTrace != nil { - timeEnd := time.Now() - traceData := &LogTraceData{ - TimeStart: timeStart, - TimeEnd: timeEnd, - Request: c.Http, - Err: err, - } - c.LogTrace.Trace(traceData) - } - + body, err := c.RequestClient.Do(c.ctx, c.Http) if err != nil { return err } - body, err := ioutil.ReadAll(resp.Body) - // 解析处理 headerData := pack.NewHeaderWithBody(body, c.Request.Protocol) packHandler = pack.GetPackHandler(headerData.Packager) @@ -91,10 +95,10 @@ func (c *Client) Send() error { err = packHandler.Decode(bodyContent, c.Response) if c.Response.Except != nil { - Log.WithFields(logrus.Fields{"YAR": "Except"}).Debug(c.Response.Except) + logger.Log.WithFields(logrus.Fields{"YAR": "Except"}).Debug(c.Response.Except) } - Log.WithFields(logrus.Fields{"YAR": "BodyContent"}).Debug(string(bodyContent)) + logger.Log.WithFields(logrus.Fields{"YAR": "BodyContent"}).Debug(string(bodyContent)) if err != nil { return err diff --git a/go.mod b/go.mod index 11a8b10..3f2dfda 100644 --- a/go.mod +++ b/go.mod @@ -1,6 +1,6 @@ module github.com/flyhope/go-yar -go 1.13 +go 1.15 require ( github.com/niemeyer/pretty v0.0.0-20200227124842-a10e7caefd8e // indirect diff --git a/log.go b/logger/log.go similarity index 97% rename from log.go rename to logger/log.go index 7e1f02c..b9be45c 100644 --- a/log.go +++ b/logger/log.go @@ -1,4 +1,4 @@ -package yar +package logger import ( "github.com/sirupsen/logrus" diff --git a/set.go b/logger/set.go similarity index 96% rename from set.go rename to logger/set.go index 4b7963c..a0c3479 100644 --- a/set.go +++ b/logger/set.go @@ -1,4 +1,4 @@ -package yar +package logger import ( "github.com/sirupsen/logrus" diff --git a/request/define.go b/request/define.go new file mode 100644 index 0000000..424310c --- /dev/null +++ b/request/define.go @@ -0,0 +1,20 @@ +package request + +import ( + "context" + "github.com/flyhope/go-yar/logger" + "net/http" +) + +type Handler interface { + SetLog(logger.LogTrace) + Do(ctx context.Context, req *http.Request) ([]byte, error) +} + +type Abstract struct { + logger.LogTrace +} + +func (a *Abstract) SetLog(l logger.LogTrace) { + a.LogTrace = l +} diff --git a/request/http.go b/request/http.go new file mode 100644 index 0000000..7fd0c8e --- /dev/null +++ b/request/http.go @@ -0,0 +1,38 @@ +package request + +import ( + "context" + "github.com/flyhope/go-yar/logger" + "io/ioutil" + "net/http" + "time" +) + +type Http struct { + Abstract + Client *http.Client +} + +func (h *Http) Do(ctx context.Context, req *http.Request) ([]byte, error) { + // 发送请求 + timeStart := time.Now() + resp, err := h.Client.Do(req.WithContext(ctx)) + + // 通过接口记录跟踪日志 + if h.LogTrace != nil { + timeEnd := time.Now() + traceData := &logger.LogTraceData{ + TimeStart: timeStart, + TimeEnd: timeEnd, + Request: req, + Err: err, + } + h.LogTrace.Trace(traceData) + } + + if err != nil { + return nil, err + } + + return ioutil.ReadAll(resp.Body) +}