Skip to content

Commit

Permalink
update
Browse files Browse the repository at this point in the history
  • Loading branch information
luopengift committed Mar 25, 2024
1 parent 511892e commit 310da48
Show file tree
Hide file tree
Showing 11 changed files with 104 additions and 50 deletions.
3 changes: 2 additions & 1 deletion .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -20,4 +20,5 @@
# Go workspace file
go.work

vendor/
vendor/
.idea/
Empty file added go.sum
Empty file.
9 changes: 9 additions & 0 deletions options.go
Original file line number Diff line number Diff line change
Expand Up @@ -260,6 +260,7 @@ func Proxy(addr string) Option {
}
}

// Setup use middleware
func Setup(httpFn ...func(HttpRoundTripFunc) HttpRoundTripFunc) Option {
return func(o *Options) {
for _, fn := range httpFn {
Expand All @@ -268,6 +269,14 @@ func Setup(httpFn ...func(HttpRoundTripFunc) HttpRoundTripFunc) Option {
}
}

// RoundTripFunc set default `*http.Transport` by customer define.
func RoundTripFunc(fn HttpRoundTripFunc) Option {
return func(o *Options) {
o.Transport = fn
}
}

// Logf print log
func Logf(f func(ctx context.Context, stat *Stat)) Option {
return func(o *Options) {
o.RoundTripFunc = append(o.RoundTripFunc, fprintf(f))
Expand Down
27 changes: 27 additions & 0 deletions requests_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ import (
"net/http"
"net/http/httptest"
"net/url"
"os"
"sync/atomic"
"testing"
"time"
Expand Down Expand Up @@ -186,3 +187,29 @@ func Test_ForEach(t *testing.T) {
t.Logf("%v, %v", resp.Stat(), err)

}

func TestResponse_Download(t *testing.T) {
s := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
text := "abc\ndef\nghij\n\n123"
fmt.Fprint(w, text)
}))
defer s.Close()

u := "https://go.dev/dl/go1.22.1.darwin-amd64.tar.gz" // a35015fca6f631f3501a36b3bccba9c5
sess := New(URL(u))
f, err := os.OpenFile("tmp/xx.tar.gz", os.O_CREATE|os.O_TRUNC|os.O_RDWR, 0644)
defer f.Close()
sum := 0
resp, err := sess.DoRequest(context.Background(),
TraceLv(3),
Stream(func(i int64, row []byte) error {
cnt, err := f.Write(row)
sum += cnt
return err
}))
if err != nil {
t.Logf("resp=%d, err=%s", resp.Content, err)
return
}
t.Logf("resp=%d, err=%s", resp.Content, err)
}
24 changes: 15 additions & 9 deletions response.go
Original file line number Diff line number Diff line change
Expand Up @@ -58,20 +58,26 @@ func (resp *Response) Stat() *Stat {
return StatLoad(resp)
}

// streamRead
// streamRead xx
func streamRead(reader io.Reader, f func(int64, []byte) error) (int64, error) {
i, cnt, r := int64(0), int64(0), bufio.NewReaderSize(reader, 1024*1024)
for {
raw, err := r.ReadBytes(10) // ascii('\n') = 10
if err != nil {
if err != io.EOF {
return cnt, err
}
return cnt, nil
raw, err1 := r.ReadBytes(10) // ascii('\n') = 10
if err1 != nil && err1 != io.EOF {
return cnt, err1
}
// 保证最后一行能被处理,并且可以正常返回
i, cnt = i+1, cnt+int64(len(raw))
if err = f(i, raw); err != nil {
return cnt, err
if err2 := f(i, raw); err1 == io.EOF || err2 != nil {
return cnt, err2
}
}
}

// CopyResponseBody xx
func CopyResponseBody(resp *http.Response) (b *bytes.Buffer, err error) {
if b, resp.Body, err = copyBody(resp.Body); err != nil {
return nil, err
}
return b, err
}
8 changes: 2 additions & 6 deletions session.go
Original file line number Diff line number Diff line change
Expand Up @@ -123,16 +123,12 @@ func (s *Session) DoRequest(ctx context.Context, opts ...Option) (*Response, err

defer resp.Response.Body.Close()

var err error
if options.Stream != nil {
_, err = streamRead(resp.Response.Body, options.Stream)
_, resp.Err = streamRead(resp.Response.Body, options.Stream)
resp.Content = bytes.NewBufferString("[consumed]")
} else {
_, err = resp.Content.ReadFrom(resp.Response.Body)
_, resp.Err = resp.Content.ReadFrom(resp.Response.Body)
resp.Response.Body = io.NopCloser(bytes.NewReader(resp.Content.Bytes()))
}
if err != nil {
resp.Err = fmt.Errorf("err1=%w, err2=%w", resp.Err, err)
}
return resp, resp.Err
}
25 changes: 15 additions & 10 deletions stat.go
Original file line number Diff line number Diff line change
@@ -1,9 +1,8 @@
package requests

import (
"bytes"
"encoding/json"
"io"
"fmt"
)

const RequestId = "Request-Id"
Expand Down Expand Up @@ -44,11 +43,10 @@ func StatLoad(resp *Response) *Stat {
}
if resp.Response != nil {
var err error
if resp.Content == nil || resp.Response.ContentLength != 0 {
resp.Content = &bytes.Buffer{}
if resp.Response.ContentLength, err = resp.Content.ReadFrom(resp.Response.Body); err == nil {
defer resp.Response.Body.Close()
resp.Response.Body = io.NopCloser(bytes.NewReader(resp.Content.Bytes()))
if resp.Content == nil || resp.Content.Len() == 0 {
if resp.Content, err = CopyResponseBody(resp.Response); err != nil {
stat.Err += fmt.Sprintf("read response: %s", err)
return stat
}
}
stat.Response.Body = make(map[string]any)
Expand All @@ -68,10 +66,17 @@ func StatLoad(resp *Response) *Stat {
stat.Request.Method = resp.Request.Method
stat.Request.URL = resp.Request.URL.String()
if resp.Request.GetBody != nil {
body, _ := resp.Request.GetBody()
body, err := resp.Request.GetBody()
if err != nil {
stat.Err += fmt.Sprintf("read request1: %s", err)
return stat
}

var buf bytes.Buffer
_, _ = buf.ReadFrom(body)
buf, _, err := copyBody(body)
if err != nil {
stat.Err += fmt.Sprintf("read request2: %s", err)
return stat
}

m := make(map[string]any)

Expand Down
23 changes: 15 additions & 8 deletions transport.go
Original file line number Diff line number Diff line change
Expand Up @@ -58,16 +58,23 @@ func verbose(v int, mLimit ...int) func(fn HttpRoundTripFunc) HttpRoundTripFunc
return nil, err
}

respLog, err := httputil.DumpResponse(resp, v > 3)
if err != nil {
return nil, err
if v >= 3 {
// 答应响应头和响应体长度
Log("< %s %s", resp.Proto, resp.Status)
for k, vs := range resp.Header {
for _, v := range vs {
Log("< %s: %s", k, v)
}
}
}
if v > 3 {
Log(show("< ", respLog, maxLimit))
} else {
Log("* resp.body is skipped")
if v >= 4 {
buf, err := CopyResponseBody(resp)
if err != nil {
Log("! response error: %w", err)
return nil, err
}
Log(show("*", buf.Bytes(), maxLimit))
}
Log("* ")
return resp, nil
}
}
Expand Down
13 changes: 8 additions & 5 deletions transport_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -22,18 +22,20 @@ func Test_Middleware(t *testing.T) {
t.Logf("session.ResponseEach end")
return nil
}),
requests.Setup(func(tripFunc requests.HttpRoundTripFunc) requests.HttpRoundTripFunc {
requests.Setup(func(fn requests.HttpRoundTripFunc) requests.HttpRoundTripFunc {
return func(req *http.Request) (*http.Response, error) {
t.Logf("session.Setup start")
defer t.Logf("session.Setup defer end")
resp, err := fn(req)
t.Logf("session.Setup end")
return tripFunc(req)
return resp, err
}
}),
)

resp, err := sess.DoRequest(
context.Background(), requests.URL(ss.URL), requests.Body(`{"Hello":"World"}`), requests.Logf(requests.LogS), requests.TraceLv(3),
context.Background(), requests.URL(ss.URL), requests.Body(`{"Hello":"World"}`),
//requests.Logf(requests.LogS), requests.TraceLv(4),
requests.RequestEach(func(ctx context.Context, r *http.Request) error {
t.Logf("doRequest.RequestEach start")
defer t.Logf("doRequest.RequestEach defer end")
Expand All @@ -46,12 +48,13 @@ func Test_Middleware(t *testing.T) {
t.Logf("doRequest.ResponseEach end")
return nil
}),
requests.Setup(func(tripFunc requests.HttpRoundTripFunc) requests.HttpRoundTripFunc {
requests.Setup(func(fn requests.HttpRoundTripFunc) requests.HttpRoundTripFunc {
return func(req *http.Request) (*http.Response, error) {
t.Logf("doRequest.Setup start")
defer t.Logf("doRequest.Setup defer end")
resp, err := fn(req)
t.Logf("doRequest.Setup end")
return tripFunc(req)
return resp, err
}
}),
)
Expand Down
4 changes: 2 additions & 2 deletions uid.go
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,6 @@ func GenId(id ...string) string {
if len(id) != 0 && id[0] != "" {
return id[0]
}
s := uint64(time.Now().UnixMicro()*1000 + source.Int63n(1000)) // % 4738381338321616895
return strings.ToUpper(strconv.FormatUint(s, 36))
i := time.Now().UnixMicro()*1000 + source.Int63n(1000) // % 4738381338321616895
return strings.ToUpper(strconv.FormatUint(uint64(i), 36))
}
18 changes: 9 additions & 9 deletions util.go
Original file line number Diff line number Diff line change
Expand Up @@ -23,24 +23,24 @@ func show(prompt string, b []byte, mLimit int) string {
return str
}

// drainBody reads all of b to memory and then returns two equivalent
// copyBody reads all of b to memory and then returns two equivalent
// ReadClosers yielding the same bytes.
//
// It returns an error if the initial slurp of all bytes fails. It does not attempt
// to make the returned ReadClosers have identical error-matching behavior.
func drainBody(b io.ReadCloser) (r1, r2 io.ReadCloser, err error) {
func copyBody(b io.ReadCloser) (*bytes.Buffer, io.ReadCloser, error) {
var buf bytes.Buffer
if b == nil || b == http.NoBody {
// No copying needed. Preserve the magic sentinel meaning of NoBody.
return http.NoBody, http.NoBody, nil
return &buf, http.NoBody, nil
}
var buf bytes.Buffer
if _, err = buf.ReadFrom(b); err != nil {
return nil, b, err
if _, err := buf.ReadFrom(b); err != nil {
return &buf, b, err
}
if err = b.Close(); err != nil {
return nil, b, err
if err := b.Close(); err != nil {
return &buf, b, err
}
return io.NopCloser(&buf), io.NopCloser(bytes.NewReader(buf.Bytes())), nil
return &buf, io.NopCloser(bytes.NewReader(buf.Bytes())), nil
}

// LogS supply default handle Stat, print to stdout.
Expand Down

0 comments on commit 310da48

Please sign in to comment.