github.com/hdt3213/godis@v1.2.9/redis/connection/fake.go (about)

     1  package connection
     2  
     3  import (
     4  	"fmt"
     5  	"github.com/hdt3213/godis/lib/logger"
     6  	"io"
     7  	"sync"
     8  )
     9  
    10  // FakeConn implements redis.Connection for test
    11  type FakeConn struct {
    12  	Connection
    13  	buf    []byte
    14  	offset int
    15  	waitOn chan struct{}
    16  	closed bool
    17  	mu     sync.Mutex
    18  }
    19  
    20  func NewFakeConn() *FakeConn {
    21  	c := &FakeConn{}
    22  	return c
    23  }
    24  
    25  // Write writes data to buffer
    26  func (c *FakeConn) Write(b []byte) (int, error) {
    27  	if c.closed {
    28  		return 0, io.EOF
    29  	}
    30  	c.mu.Lock()
    31  	c.buf = append(c.buf, b...)
    32  	c.mu.Unlock()
    33  	c.notify()
    34  	return len(b), nil
    35  }
    36  
    37  func (c *FakeConn) notify() {
    38  	if c.waitOn != nil {
    39  		c.mu.Lock()
    40  		if c.waitOn != nil {
    41  			logger.Debug(fmt.Sprintf("notify %p", c.waitOn))
    42  			close(c.waitOn)
    43  			c.waitOn = nil
    44  		}
    45  		c.mu.Unlock()
    46  	}
    47  }
    48  
    49  func (c *FakeConn) wait(offset int) {
    50  	c.mu.Lock()
    51  	if c.offset != offset { // new data during waiting lock
    52  		return
    53  	}
    54  	if c.waitOn == nil {
    55  		c.waitOn = make(chan struct{})
    56  	}
    57  	waitOn := c.waitOn
    58  	logger.Debug(fmt.Sprintf("wait on %p", waitOn))
    59  	c.mu.Unlock()
    60  	<-waitOn
    61  	logger.Debug(fmt.Sprintf("wait on %p finish", waitOn))
    62  }
    63  
    64  // Read reads data from buffer
    65  func (c *FakeConn) Read(p []byte) (int, error) {
    66  	c.mu.Lock()
    67  	n := copy(p, c.buf[c.offset:])
    68  	c.offset += n
    69  	offset := c.offset
    70  	c.mu.Unlock()
    71  	if n == 0 {
    72  		if c.closed {
    73  			return n, io.EOF
    74  		}
    75  		c.wait(offset)
    76  		// after notify
    77  		if c.closed {
    78  			return n, io.EOF
    79  		}
    80  		n = copy(p, c.buf[c.offset:])
    81  		c.offset += n
    82  		return n, nil
    83  	}
    84  	if c.closed {
    85  		return n, io.EOF
    86  	}
    87  	return n, nil
    88  }
    89  
    90  // Clean resets the buffer
    91  func (c *FakeConn) Clean() {
    92  	c.waitOn = make(chan struct{})
    93  	c.buf = nil
    94  	c.offset = 0
    95  }
    96  
    97  // Bytes returns written data
    98  func (c *FakeConn) Bytes() []byte {
    99  	return c.buf
   100  }
   101  
   102  func (c *FakeConn) Close() error {
   103  	c.closed = true
   104  	c.notify()
   105  	return nil
   106  }