github.com/sagernet/sing@v0.2.6/common/bufio/deadline/reader_fallback.go (about)

     1  package deadline
     2  
     3  import (
     4  	"time"
     5  
     6  	"github.com/sagernet/sing/common/atomic"
     7  	"github.com/sagernet/sing/common/buf"
     8  )
     9  
    10  type fallbackReader struct {
    11  	*reader
    12  	disablePipe atomic.Bool
    13  	inRead      atomic.Bool
    14  }
    15  
    16  func NewFallbackReader(timeoutReader TimeoutReader) Reader {
    17  	return &fallbackReader{reader: NewReader(timeoutReader).(*reader)}
    18  }
    19  
    20  func (r *fallbackReader) Read(p []byte) (n int, err error) {
    21  	select {
    22  	case result := <-r.result:
    23  		return r.pipeReturn(result, p)
    24  	default:
    25  	}
    26  	if r.disablePipe.Load() {
    27  		return r.ExtendedReader.Read(p)
    28  	}
    29  	select {
    30  	case <-r.done:
    31  		if r.deadline.Load().IsZero() {
    32  			r.done <- struct{}{}
    33  			r.inRead.Store(true)
    34  			defer r.inRead.Store(false)
    35  			n, err = r.ExtendedReader.Read(p)
    36  			return
    37  		}
    38  		go r.pipeRead(len(p))
    39  	default:
    40  	}
    41  	return r.reader.read(p)
    42  }
    43  
    44  func (r *fallbackReader) ReadBuffer(buffer *buf.Buffer) error {
    45  	select {
    46  	case result := <-r.result:
    47  		return r.pipeReturnBuffer(result, buffer)
    48  	default:
    49  	}
    50  	if r.disablePipe.Load() {
    51  		return r.ExtendedReader.ReadBuffer(buffer)
    52  	}
    53  	select {
    54  	case <-r.done:
    55  		if r.deadline.Load().IsZero() {
    56  			r.done <- struct{}{}
    57  			r.inRead.Store(true)
    58  			defer r.inRead.Store(false)
    59  			return r.ExtendedReader.ReadBuffer(buffer)
    60  		}
    61  		go r.pipeReadBuffer(buffer.FreeLen())
    62  	default:
    63  	}
    64  	return r.readBuffer(buffer)
    65  }
    66  
    67  func (r *fallbackReader) SetReadDeadline(t time.Time) error {
    68  	if r.disablePipe.Load() {
    69  		return r.timeoutReader.SetReadDeadline(t)
    70  	} else if r.inRead.Load() {
    71  		r.disablePipe.Store(true)
    72  		return r.timeoutReader.SetReadDeadline(t)
    73  	}
    74  	return r.reader.SetReadDeadline(t)
    75  }
    76  
    77  func (r *fallbackReader) ReaderReplaceable() bool {
    78  	return r.disablePipe.Load() || r.reader.ReaderReplaceable()
    79  }
    80  
    81  func (r *fallbackReader) UpstreamReader() any {
    82  	return r.reader.UpstreamReader()
    83  }