1.4.5.2

EdgeNode/internal/utils/writers/writer_bytes_counter.go

@@ -0,0 +1,32 @@
+// Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
+
+package writers
+
+import "io"
+
+type BytesCounterWriter struct {
+	writer io.Writer
+	count  int64
+}
+
+func NewBytesCounterWriter(rawWriter io.Writer) *BytesCounterWriter {
+	return &BytesCounterWriter{writer: rawWriter}
+}
+
+func (this *BytesCounterWriter) RawWriter() io.Writer {
+	return this.writer
+}
+
+func (this *BytesCounterWriter) Write(p []byte) (n int, err error) {
+	n, err = this.writer.Write(p)
+	this.count += int64(n)
+	return
+}
+
+func (this *BytesCounterWriter) Close() error {
+	return nil
+}
+
+func (this *BytesCounterWriter) TotalBytes() int64 {
+	return this.count
+}

EdgeNode/internal/utils/writers/writer_closer_tee.go

@@ -0,0 +1,53 @@
+// Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
+
+package writers
+
+import (
+	"io"
+)
+
+type TeeWriterCloser struct {
+	primaryW   io.WriteCloser
+	secondaryW io.WriteCloser
+
+	onFail func(err error)
+}
+
+func NewTeeWriterCloser(primaryW io.WriteCloser, secondaryW io.WriteCloser) *TeeWriterCloser {
+	return &TeeWriterCloser{
+		primaryW:   primaryW,
+		secondaryW: secondaryW,
+	}
+}
+
+func (this *TeeWriterCloser) Write(p []byte) (n int, err error) {
+	{
+		n, err = this.primaryW.Write(p)
+
+		if err != nil {
+			if this.onFail != nil {
+				this.onFail(err)
+			}
+		}
+	}
+
+	{
+		_, err2 := this.secondaryW.Write(p)
+		if err2 != nil {
+			if this.onFail != nil {
+				this.onFail(err2)
+			}
+		}
+	}
+
+	return
+}
+
+func (this *TeeWriterCloser) Close() error {
+	// 这里不关闭secondary
+	return this.primaryW.Close()
+}
+
+func (this *TeeWriterCloser) OnFail(onFail func(err error)) {
+	this.onFail = onFail
+}

EdgeNode/internal/utils/writers/writer_print.go

@@ -0,0 +1,28 @@
+// Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved. Official site: https://goedge.cn .
+
+package writers
+
+import (
+	"io"
+	"log"
+)
+
+type PrintWriter struct {
+	rawWriter io.Writer
+	tag       string
+}
+
+func NewPrintWriter(rawWriter io.Writer, tag string) io.Writer {
+	return &PrintWriter{
+		rawWriter: rawWriter,
+		tag:       tag,
+	}
+}
+
+func (this *PrintWriter) Write(p []byte) (n int, err error) {
+	n, err = this.rawWriter.Write(p)
+	if n > 0 {
+		log.Println("[" + this.tag + "]" + string(p[:n]))
+	}
+	return
+}

EdgeNode/internal/utils/writers/writer_rate_limit.go

@@ -0,0 +1,97 @@
+// Copyright 2021 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
+
+package writers
+
+import (
+	"context"
+	"github.com/iwind/TeaGo/types"
+	"io"
+	"time"
+)
+
+// RateLimitWriter 限速写入
+type RateLimitWriter struct {
+	rawWriter io.WriteCloser
+	ctx       context.Context
+
+	rateBytes int
+
+	written int
+	before  time.Time
+}
+
+func NewRateLimitWriter(ctx context.Context, rawWriter io.WriteCloser, rateBytes int64) io.WriteCloser {
+	return &RateLimitWriter{
+		rawWriter: rawWriter,
+		ctx:       ctx,
+		rateBytes: types.Int(rateBytes),
+		before:    time.Now(),
+	}
+}
+
+func (this *RateLimitWriter) Write(p []byte) (n int, err error) {
+	if this.rateBytes <= 0 {
+		return this.write(p)
+	}
+
+	var size = len(p)
+	if size == 0 {
+		return 0, nil
+	}
+
+	if size <= this.rateBytes {
+		return this.write(p)
+	}
+
+	for {
+		size = len(p)
+
+		var limit = this.rateBytes
+		if limit > size {
+			limit = size
+		}
+		n1, wErr := this.write(p[:limit])
+		n += n1
+		if wErr != nil {
+			return n, wErr
+		}
+
+		if size > limit {
+			p = p[limit:]
+		} else {
+			break
+		}
+	}
+
+	return
+}
+
+func (this *RateLimitWriter) Close() error {
+	return this.rawWriter.Close()
+}
+
+func (this *RateLimitWriter) write(p []byte) (n int, err error) {
+	n, err = this.rawWriter.Write(p)
+
+	if err == nil {
+		select {
+		case <-this.ctx.Done():
+			err = io.EOF
+			return
+		default:
+		}
+
+		this.written += n
+
+		if this.written >= this.rateBytes {
+			var duration = 1*time.Second - time.Since(this.before)
+			if duration > 0 {
+				time.Sleep(duration)
+			}
+			this.before = time.Now()
+			this.written = 0
+		}
+	}
+
+	return
+}

EdgeNode/internal/utils/writers/writer_rate_limit_test.go

@@ -0,0 +1,41 @@
+// Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
+
+package writers
+
+import (
+	"sync"
+	"testing"
+	"time"
+)
+
+func TestSleep(t *testing.T) {
+	var count = 2000
+	var wg = sync.WaitGroup{}
+	wg.Add(count)
+	var before = time.Now()
+	for i := 0; i < count; i++ {
+		go func() {
+			defer wg.Done()
+			time.Sleep(1 * time.Second)
+		}()
+	}
+	wg.Wait()
+	t.Log(time.Since(before).Seconds()*1000, "ms")
+}
+
+func TestTimeout(t *testing.T) {
+	var count = 2000
+	var wg = sync.WaitGroup{}
+	wg.Add(count)
+	var before = time.Now()
+	for i := 0; i < count; i++ {
+		go func() {
+			defer wg.Done()
+
+			var timeout = time.NewTimer(1 * time.Second)
+			<-timeout.C
+		}()
+	}
+	wg.Wait()
+	t.Log(time.Since(before).Seconds()*1000, "ms")
+}