示例#1
0
文件: throttle.go 项目: na--/nedomi
// New creates and returns a ready to used ServerStatusHandler.
func New(cfg *config.Handler, l *types.Location, next types.RequestHandler) (types.RequestHandler, error) {
	var s struct {
		Speed types.BytesSize `json:"speed"`
	}
	if err := json.Unmarshal(cfg.Settings, &s); err != nil {
		return nil, fmt.Errorf("handler.throttle got error while parsing settings - %s", err)
	}
	if s.Speed == 0 {
		return nil, fmt.Errorf("handler.throttle needs to have speed settings > 0")
	}
	return types.RequestHandlerFunc(
		func(ctx context.Context, w http.ResponseWriter, r *http.Request) {
			next.RequestHandle(ctx, &throttledResponseWriter{ResponseWriter: w,
				ThrottlerWriter: iocontrol.ThrottledWriter(w, int(s.Speed.Bytes()), time.Millisecond*10),
			}, r)
		}), nil
}
示例#2
0
func Example_ThrottledWriter() {

	totalSize := 10 * iocontrol.KiB
	readPerSec := 100 * iocontrol.KiB
	maxBurst := 10 * time.Millisecond

	input := randBytes(totalSize)
	src := bytes.NewReader(input)
	dst := bytes.NewBuffer(nil)
	measured := iocontrol.NewMeasuredWriter(dst)
	throttled := iocontrol.ThrottledWriter(measured, readPerSec, maxBurst)

	done := make(chan []byte)
	go func() {
		start := time.Now()
		_, err := io.Copy(throttled, src)
		fmt.Printf("done in %.1fs", time.Since(start).Seconds())
		if err != nil {
			log.Fatalf("error writing: %v", err)
		}
		done <- dst.Bytes()
	}()

	for {
		select {
		case <-time.Tick(time.Millisecond * 10):
			log.Printf("writing at %s/s", humanize.IBytes(measured.BytesPerSec()))

		case output := <-done:
			if !bytes.Equal(input, output) {
				log.Print("==== input ====\n", hex.Dump(input))
				log.Print("==== output ====\n", hex.Dump(output))
				log.Fatalf("mismatch between input and output")
			}
			return
		}
	}

	// Output:
	// done in 0.1s
}
示例#3
0
func TestClientTimeoutSlowBody(t *testing.T) {
	// server that flushes the headers ASAP, but sends the body slowly
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(200)
		if f, ok := w.(http.Flusher); ok {
			f.Flush()
		}
		tw := iocontrol.ThrottledWriter(w, 2, time.Second)
		fmt.Fprint(tw, r.URL.Path)
	}))
	defer srv.Close()

	runWithClient := func(c *http.Client) {
		res, err := c.Get(srv.URL + "/testing")

		// should receive a response
		require.Nil(t, err)
		require.NotNil(t, res)

		// should fail with timeout while reading body
		_, err = io.Copy(ioutil.Discard, res.Body)
		res.Body.Close()
		assertNetTimeoutErr(t, err)
	}

	// test with retry transport
	tr := NewTransport(nil, RetryAll(RetryMaxRetries(2), RetryTemporaryErr()), ConstDelay(time.Second))

	c := &http.Client{
		Transport: tr,
		Timeout:   time.Second,
	}
	runWithClient(c)

	// test with default transport, make sure it behaves the same way
	c = &http.Client{Timeout: time.Second}
	runWithClient(c)
}