support custom handler

client.go

@@ -1,40 +1,58 @@
 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"
 	"time"
 )
 
 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
 	}
 
 	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

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

log.go → logger/log.go

@@ -1,4 +1,4 @@
-package yar
+package logger
 
 import (
 	"github.com/sirupsen/logrus"

set.go → logger/set.go

@@ -1,4 +1,4 @@
-package yar
+package logger
 
 import (
 	"github.com/sirupsen/logrus"

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
+}

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)
+}