gitee.com/ks-custle/core-gm@v0.0.0-20230922171213-b83bdd97b62c/grpc/test/bufconn/bufconn.go (about)

     1  /*
     2   *
     3   * Copyright 2017 gRPC authors.
     4   *
     5   * Licensed under the Apache License, Version 2.0 (the "License");
     6   * you may not use this file except in compliance with the License.
     7   * You may obtain a copy of the License at
     8   *
     9   *     http://www.apache.org/licenses/LICENSE-2.0
    10   *
    11   * Unless required by applicable law or agreed to in writing, software
    12   * distributed under the License is distributed on an "AS IS" BASIS,
    13   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    14   * See the License for the specific language governing permissions and
    15   * limitations under the License.
    16   *
    17   */
    18  
    19  // Package bufconn provides a net.Conn implemented by a buffer and related
    20  // dialing and listening functionality.
    21  package bufconn
    22  
    23  import (
    24  	"context"
    25  	"fmt"
    26  	"io"
    27  	"net"
    28  	"sync"
    29  	"time"
    30  )
    31  
    32  // Listener implements a net.Listener that creates local, buffered net.Conns
    33  // via its Accept and Dial method.
    34  type Listener struct {
    35  	mu   sync.Mutex
    36  	sz   int
    37  	ch   chan net.Conn
    38  	done chan struct{}
    39  }
    40  
    41  // Implementation of net.Error providing timeout
    42  type netErrorTimeout struct {
    43  	error
    44  }
    45  
    46  func (e netErrorTimeout) Timeout() bool   { return true }
    47  func (e netErrorTimeout) Temporary() bool { return false }
    48  
    49  var errClosed = fmt.Errorf("closed")
    50  var errTimeout net.Error = netErrorTimeout{error: fmt.Errorf("i/o timeout")}
    51  
    52  // Listen returns a Listener that can only be contacted by its own Dialers and
    53  // creates buffered connections between the two.
    54  func Listen(sz int) *Listener {
    55  	return &Listener{sz: sz, ch: make(chan net.Conn), done: make(chan struct{})}
    56  }
    57  
    58  // Accept blocks until Dial is called, then returns a net.Conn for the server
    59  // half of the connection.
    60  func (l *Listener) Accept() (net.Conn, error) {
    61  	select {
    62  	case <-l.done:
    63  		return nil, errClosed
    64  	case c := <-l.ch:
    65  		return c, nil
    66  	}
    67  }
    68  
    69  // Close stops the listener.
    70  func (l *Listener) Close() error {
    71  	l.mu.Lock()
    72  	defer l.mu.Unlock()
    73  	select {
    74  	case <-l.done:
    75  		// Already closed.
    76  		break
    77  	default:
    78  		close(l.done)
    79  	}
    80  	return nil
    81  }
    82  
    83  // Addr reports the address of the listener.
    84  func (l *Listener) Addr() net.Addr { return addr{} }
    85  
    86  // Dial creates an in-memory full-duplex network connection, unblocks Accept by
    87  // providing it the server half of the connection, and returns the client half
    88  // of the connection.
    89  func (l *Listener) Dial() (net.Conn, error) {
    90  	return l.DialContext(context.Background())
    91  }
    92  
    93  // DialContext creates an in-memory full-duplex network connection, unblocks Accept by
    94  // providing it the server half of the connection, and returns the client half
    95  // of the connection.  If ctx is Done, returns ctx.Err()
    96  func (l *Listener) DialContext(ctx context.Context) (net.Conn, error) {
    97  	p1, p2 := newPipe(l.sz), newPipe(l.sz)
    98  	select {
    99  	case <-ctx.Done():
   100  		return nil, ctx.Err()
   101  	case <-l.done:
   102  		return nil, errClosed
   103  	case l.ch <- &conn{p1, p2}:
   104  		return &conn{p2, p1}, nil
   105  	}
   106  }
   107  
   108  type pipe struct {
   109  	mu sync.Mutex
   110  
   111  	// buf contains the data in the pipe.  It is a ring buffer of fixed capacity,
   112  	// with r and w pointing to the offset to read and write, respsectively.
   113  	//
   114  	// Data is read between [r, w) and written to [w, r), wrapping around the end
   115  	// of the slice if necessary.
   116  	//
   117  	// The buffer is empty if r == len(buf), otherwise if r == w, it is full.
   118  	//
   119  	// w and r are always in the range [0, cap(buf)) and [0, len(buf)].
   120  	buf  []byte
   121  	w, r int
   122  
   123  	wwait sync.Cond
   124  	rwait sync.Cond
   125  
   126  	// Indicate that a write/read timeout has occurred
   127  	wtimedout bool
   128  	rtimedout bool
   129  
   130  	wtimer *time.Timer
   131  	rtimer *time.Timer
   132  
   133  	closed      bool
   134  	writeClosed bool
   135  }
   136  
   137  func newPipe(sz int) *pipe {
   138  	p := &pipe{buf: make([]byte, 0, sz)}
   139  	p.wwait.L = &p.mu
   140  	p.rwait.L = &p.mu
   141  
   142  	p.wtimer = time.AfterFunc(0, func() {})
   143  	p.rtimer = time.AfterFunc(0, func() {})
   144  	return p
   145  }
   146  
   147  func (p *pipe) empty() bool {
   148  	return p.r == len(p.buf)
   149  }
   150  
   151  func (p *pipe) full() bool {
   152  	return p.r < len(p.buf) && p.r == p.w
   153  }
   154  
   155  func (p *pipe) Read(b []byte) (n int, err error) {
   156  	p.mu.Lock()
   157  	defer p.mu.Unlock()
   158  	// Block until p has data.
   159  	for {
   160  		if p.closed {
   161  			return 0, io.ErrClosedPipe
   162  		}
   163  		if !p.empty() {
   164  			break
   165  		}
   166  		if p.writeClosed {
   167  			return 0, io.EOF
   168  		}
   169  		if p.rtimedout {
   170  			return 0, errTimeout
   171  		}
   172  
   173  		p.rwait.Wait()
   174  	}
   175  	wasFull := p.full()
   176  
   177  	n = copy(b, p.buf[p.r:len(p.buf)])
   178  	p.r += n
   179  	if p.r == cap(p.buf) {
   180  		p.r = 0
   181  		p.buf = p.buf[:p.w]
   182  	}
   183  
   184  	// Signal a blocked writer, if any
   185  	if wasFull {
   186  		p.wwait.Signal()
   187  	}
   188  
   189  	return n, nil
   190  }
   191  
   192  func (p *pipe) Write(b []byte) (n int, err error) {
   193  	p.mu.Lock()
   194  	defer p.mu.Unlock()
   195  	if p.closed {
   196  		return 0, io.ErrClosedPipe
   197  	}
   198  	for len(b) > 0 {
   199  		// Block until p is not full.
   200  		for {
   201  			if p.closed || p.writeClosed {
   202  				return 0, io.ErrClosedPipe
   203  			}
   204  			if !p.full() {
   205  				break
   206  			}
   207  			if p.wtimedout {
   208  				return 0, errTimeout
   209  			}
   210  
   211  			p.wwait.Wait()
   212  		}
   213  		wasEmpty := p.empty()
   214  
   215  		end := cap(p.buf)
   216  		if p.w < p.r {
   217  			end = p.r
   218  		}
   219  		x := copy(p.buf[p.w:end], b)
   220  		b = b[x:]
   221  		n += x
   222  		p.w += x
   223  		if p.w > len(p.buf) {
   224  			p.buf = p.buf[:p.w]
   225  		}
   226  		if p.w == cap(p.buf) {
   227  			p.w = 0
   228  		}
   229  
   230  		// Signal a blocked reader, if any.
   231  		if wasEmpty {
   232  			p.rwait.Signal()
   233  		}
   234  	}
   235  	return n, nil
   236  }
   237  
   238  func (p *pipe) Close() error {
   239  	p.mu.Lock()
   240  	defer p.mu.Unlock()
   241  	p.closed = true
   242  	// Signal all blocked readers and writers to return an error.
   243  	p.rwait.Broadcast()
   244  	p.wwait.Broadcast()
   245  	return nil
   246  }
   247  
   248  func (p *pipe) closeWrite() error {
   249  	p.mu.Lock()
   250  	defer p.mu.Unlock()
   251  	p.writeClosed = true
   252  	// Signal all blocked readers and writers to return an error.
   253  	p.rwait.Broadcast()
   254  	p.wwait.Broadcast()
   255  	return nil
   256  }
   257  
   258  type conn struct {
   259  	io.Reader
   260  	io.Writer
   261  }
   262  
   263  func (c *conn) Close() error {
   264  	err1 := c.Reader.(*pipe).Close()
   265  	err2 := c.Writer.(*pipe).closeWrite()
   266  	if err1 != nil {
   267  		return err1
   268  	}
   269  	return err2
   270  }
   271  
   272  func (c *conn) SetDeadline(t time.Time) error {
   273  	c.SetReadDeadline(t)
   274  	c.SetWriteDeadline(t)
   275  	return nil
   276  }
   277  
   278  func (c *conn) SetReadDeadline(t time.Time) error {
   279  	p := c.Reader.(*pipe)
   280  	p.mu.Lock()
   281  	defer p.mu.Unlock()
   282  	p.rtimer.Stop()
   283  	p.rtimedout = false
   284  	if !t.IsZero() {
   285  		p.rtimer = time.AfterFunc(time.Until(t), func() {
   286  			p.mu.Lock()
   287  			defer p.mu.Unlock()
   288  			p.rtimedout = true
   289  			p.rwait.Broadcast()
   290  		})
   291  	}
   292  	return nil
   293  }
   294  
   295  func (c *conn) SetWriteDeadline(t time.Time) error {
   296  	p := c.Writer.(*pipe)
   297  	p.mu.Lock()
   298  	defer p.mu.Unlock()
   299  	p.wtimer.Stop()
   300  	p.wtimedout = false
   301  	if !t.IsZero() {
   302  		p.wtimer = time.AfterFunc(time.Until(t), func() {
   303  			p.mu.Lock()
   304  			defer p.mu.Unlock()
   305  			p.wtimedout = true
   306  			p.wwait.Broadcast()
   307  		})
   308  	}
   309  	return nil
   310  }
   311  
   312  func (*conn) LocalAddr() net.Addr  { return addr{} }
   313  func (*conn) RemoteAddr() net.Addr { return addr{} }
   314  
   315  type addr struct{}
   316  
   317  func (addr) Network() string { return "bufconn" }
   318  func (addr) String() string  { return "bufconn" }