Initial commit (code only without large binaries)

EdgeNode/internal/compressions/writer_pool.go

@@ -0,0 +1,82 @@
+// Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
+
+package compressions
+
+import (
+	teaconst "github.com/TeaOSLab/EdgeNode/internal/const"
+	"github.com/TeaOSLab/EdgeNode/internal/utils/goman"
+	"io"
+	"time"
+)
+
+const maxWriterHits = 1 << 20
+
+var isBusy = false
+
+func init() {
+	if !teaconst.IsMain {
+		return
+	}
+
+	goman.New(func() {
+		var ticker = time.NewTicker(100 * time.Millisecond)
+		for range ticker.C {
+			if isBusy {
+				isBusy = false
+			}
+		}
+	})
+}
+
+func IsBusy() bool {
+	return isBusy
+}
+
+type WriterPool struct {
+	c       chan Writer // level => chan Writer
+	newFunc func(writer io.Writer, level int) (Writer, error)
+}
+
+func NewWriterPool(maxSize int, newFunc func(writer io.Writer, level int) (Writer, error)) *WriterPool {
+	if maxSize <= 0 {
+		maxSize = 1024
+	}
+
+	return &WriterPool{
+		c:       make(chan Writer, maxSize),
+		newFunc: newFunc,
+	}
+}
+
+func (this *WriterPool) Get(parentWriter io.Writer, level int) (Writer, error) {
+	if isBusy {
+		return nil, ErrIsBusy
+	}
+
+	select {
+	case writer := <-this.c:
+		writer.Reset(parentWriter)
+		writer.ResetFinish()
+		return writer, nil
+	default:
+		writer, err := this.newFunc(parentWriter, level)
+		if err != nil {
+			return nil, err
+		}
+		writer.SetPool(this)
+		return writer, nil
+	}
+}
+
+func (this *WriterPool) Put(writer Writer) {
+	if writer.IncreaseHit() > maxWriterHits {
+		// do nothing to discard it
+		return
+	}
+
+	select {
+	case this.c <- writer:
+	default:
+		isBusy = true
+	}
+}