github.com/TeaOSLab/EdgeNode@v1.3.8/internal/compressions/writer_pool.go (about)

     1  // Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
     2  
     3  package compressions
     4  
     5  import (
     6  	teaconst "github.com/TeaOSLab/EdgeNode/internal/const"
     7  	"github.com/TeaOSLab/EdgeNode/internal/goman"
     8  	"io"
     9  	"time"
    10  )
    11  
    12  const maxWriterHits = 1 << 20
    13  
    14  var isBusy = false
    15  
    16  func init() {
    17  	if !teaconst.IsMain {
    18  		return
    19  	}
    20  
    21  	goman.New(func() {
    22  		var ticker = time.NewTicker(100 * time.Millisecond)
    23  		for range ticker.C {
    24  			if isBusy {
    25  				isBusy = false
    26  			}
    27  		}
    28  	})
    29  }
    30  
    31  func IsBusy() bool {
    32  	return isBusy
    33  }
    34  
    35  type WriterPool struct {
    36  	c       chan Writer // level => chan Writer
    37  	newFunc func(writer io.Writer, level int) (Writer, error)
    38  }
    39  
    40  func NewWriterPool(maxSize int, newFunc func(writer io.Writer, level int) (Writer, error)) *WriterPool {
    41  	if maxSize <= 0 {
    42  		maxSize = 1024
    43  	}
    44  
    45  	return &WriterPool{
    46  		c:       make(chan Writer, maxSize),
    47  		newFunc: newFunc,
    48  	}
    49  }
    50  
    51  func (this *WriterPool) Get(parentWriter io.Writer, level int) (Writer, error) {
    52  	if isBusy {
    53  		return nil, ErrIsBusy
    54  	}
    55  
    56  	select {
    57  	case writer := <-this.c:
    58  		writer.Reset(parentWriter)
    59  		writer.ResetFinish()
    60  		return writer, nil
    61  	default:
    62  		writer, err := this.newFunc(parentWriter, level)
    63  		if err != nil {
    64  			return nil, err
    65  		}
    66  		writer.SetPool(this)
    67  		return writer, nil
    68  	}
    69  }
    70  
    71  func (this *WriterPool) Put(writer Writer) {
    72  	if writer.IncreaseHit() > maxWriterHits {
    73  		// do nothing to discard it
    74  		return
    75  	}
    76  
    77  	select {
    78  	case this.c <- writer:
    79  	default:
    80  		isBusy = true
    81  	}
    82  }