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

     1  // Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
     2  
     3  package compressions
     4  
     5  import (
     6  	"io"
     7  )
     8  
     9  const maxReadHits = 1 << 20
    10  
    11  type ReaderPool struct {
    12  	c       chan Reader
    13  	newFunc func(reader io.Reader) (Reader, error)
    14  }
    15  
    16  func NewReaderPool(maxSize int, newFunc func(reader io.Reader) (Reader, error)) *ReaderPool {
    17  	if maxSize <= 0 {
    18  		maxSize = 1024
    19  	}
    20  
    21  	return &ReaderPool{
    22  		c:       make(chan Reader, maxSize),
    23  		newFunc: newFunc,
    24  	}
    25  }
    26  
    27  func (this *ReaderPool) Get(parentReader io.Reader) (Reader, error) {
    28  	select {
    29  	case reader := <-this.c:
    30  		err := reader.Reset(parentReader)
    31  		if err != nil {
    32  			// create new
    33  			reader, err = this.newFunc(parentReader)
    34  			if err != nil {
    35  				return nil, err
    36  			}
    37  			reader.SetPool(this)
    38  			return reader, nil
    39  		}
    40  		reader.ResetFinish()
    41  		return reader, nil
    42  	default:
    43  		// create new
    44  		reader, err := this.newFunc(parentReader)
    45  		if err != nil {
    46  			return nil, err
    47  		}
    48  		reader.SetPool(this)
    49  		return reader, nil
    50  	}
    51  }
    52  
    53  func (this *ReaderPool) Put(reader Reader) {
    54  	if reader.IncreaseHit() > maxReadHits {
    55  		// do nothing to discard it
    56  		return
    57  	}
    58  
    59  	select {
    60  	case this.c <- reader:
    61  	default:
    62  	}
    63  }