github.com/sagernet/sing@v0.4.0-beta.19.0.20240518125136-f67a0988a636/common/bufio/deadline/reader_fallback.go (about)

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