github.com/gospider007/requests@v0.0.0-20240506025355-c73d46169a23/pip.go (about)

     1  package requests
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  )
     7  
     8  type pipCon struct {
     9  	reader  <-chan []byte
    10  	writer  chan<- []byte
    11  	readerI <-chan int
    12  	writerI chan<- int
    13  	lock    sync.Mutex
    14  	ctx     context.Context
    15  	cnl     context.CancelCauseFunc
    16  }
    17  
    18  func (obj *pipCon) Read(b []byte) (n int, err error) {
    19  	select {
    20  	case con := <-obj.reader:
    21  		n = copy(b, con)
    22  		select {
    23  		case obj.writerI <- n:
    24  			return
    25  		case <-obj.ctx.Done():
    26  			return n, context.Cause(obj.ctx)
    27  		}
    28  	case <-obj.ctx.Done():
    29  		return n, context.Cause(obj.ctx)
    30  	}
    31  }
    32  func (obj *pipCon) Write(b []byte) (n int, err error) {
    33  	obj.lock.Lock()
    34  	defer obj.lock.Unlock()
    35  	for once := true; once || len(b) > 0; once = false {
    36  		select {
    37  		case obj.writer <- b:
    38  			select {
    39  			case i := <-obj.readerI:
    40  				b = b[i:]
    41  				n += i
    42  			case <-obj.ctx.Done():
    43  				return n, context.Cause(obj.ctx)
    44  			}
    45  		case <-obj.ctx.Done():
    46  			return n, context.Cause(obj.ctx)
    47  		}
    48  	}
    49  	return
    50  }
    51  func (obj *pipCon) Close(err error) error {
    52  	obj.cnl(err)
    53  	return nil
    54  }
    55  
    56  func pipe(preCtx context.Context) (*pipCon, *pipCon) {
    57  	ctx, cnl := context.WithCancelCause(preCtx)
    58  	readerCha := make(chan []byte)
    59  	writerCha := make(chan []byte)
    60  
    61  	readerI := make(chan int)
    62  	writerI := make(chan int)
    63  	localpipCon := &pipCon{
    64  		reader:  readerCha,
    65  		readerI: readerI,
    66  		writer:  writerCha,
    67  		writerI: writerI,
    68  		ctx:     ctx,
    69  		cnl:     cnl,
    70  	}
    71  	remotepipCon := &pipCon{
    72  		reader:  writerCha,
    73  		readerI: writerI,
    74  		writer:  readerCha,
    75  		writerI: readerI,
    76  		ctx:     ctx,
    77  		cnl:     cnl,
    78  	}
    79  	return localpipCon, remotepipCon
    80  }