github.com/keltia/go-ipfs@v0.3.8-0.20150909044612-210793031c63/p2p/net/mock/mock_stream.go (about)

     1  package mocknet
     2  
     3  import (
     4  	"bytes"
     5  	"io"
     6  	"time"
     7  
     8  	process "github.com/ipfs/go-ipfs/Godeps/_workspace/src/github.com/jbenet/goprocess"
     9  
    10  	inet "github.com/ipfs/go-ipfs/p2p/net"
    11  )
    12  
    13  // stream implements inet.Stream
    14  type stream struct {
    15  	io.Reader
    16  	io.Writer
    17  	conn      *conn
    18  	toDeliver chan *transportObject
    19  	proc      process.Process
    20  }
    21  
    22  type transportObject struct {
    23  	msg         []byte
    24  	arrivalTime time.Time
    25  }
    26  
    27  func NewStream(w io.Writer, r io.Reader) *stream {
    28  	s := &stream{
    29  		Reader:    r,
    30  		Writer:    w,
    31  		toDeliver: make(chan *transportObject),
    32  	}
    33  
    34  	s.proc = process.WithTeardown(s.teardown)
    35  	s.proc.Go(s.transport)
    36  	return s
    37  }
    38  
    39  //  How to handle errors with writes?
    40  func (s *stream) Write(p []byte) (n int, err error) {
    41  	l := s.conn.link
    42  	delay := l.GetLatency() + l.RateLimit(len(p))
    43  	t := time.Now().Add(delay)
    44  	select {
    45  	case <-s.proc.Closing(): // bail out if we're closing.
    46  		return 0, io.ErrClosedPipe
    47  	case s.toDeliver <- &transportObject{msg: p, arrivalTime: t}:
    48  	}
    49  	return len(p), nil
    50  }
    51  
    52  func (s *stream) Close() error {
    53  	return s.proc.Close()
    54  }
    55  
    56  // teardown shuts down the stream. it is called by s.proc.Close()
    57  // after all the children of this s.proc (i.e. transport's proc)
    58  // are done.
    59  func (s *stream) teardown() error {
    60  	// at this point, no streams are writing.
    61  
    62  	s.conn.removeStream(s)
    63  	if r, ok := (s.Reader).(io.Closer); ok {
    64  		r.Close()
    65  	}
    66  	if w, ok := (s.Writer).(io.Closer); ok {
    67  		w.Close()
    68  	}
    69  	s.conn.net.notifyAll(func(n inet.Notifiee) {
    70  		n.ClosedStream(s.conn.net, s)
    71  	})
    72  	return nil
    73  }
    74  
    75  func (s *stream) Conn() inet.Conn {
    76  	return s.conn
    77  }
    78  
    79  // transport will grab message arrival times, wait until that time, and
    80  // then write the message out when it is scheduled to arrive
    81  func (s *stream) transport(proc process.Process) {
    82  	bufsize := 256
    83  	buf := new(bytes.Buffer)
    84  	ticker := time.NewTicker(time.Millisecond * 4)
    85  
    86  	// writeBuf writes the contents of buf through to the s.Writer.
    87  	// done only when arrival time makes sense.
    88  	drainBuf := func() {
    89  		if buf.Len() > 0 {
    90  			_, err := s.Writer.Write(buf.Bytes())
    91  			if err != nil {
    92  				return
    93  			}
    94  			buf.Reset()
    95  		}
    96  	}
    97  
    98  	// deliverOrWait is a helper func that processes
    99  	// an incoming packet. it waits until the arrival time,
   100  	// and then writes things out.
   101  	deliverOrWait := func(o *transportObject) {
   102  		buffered := len(o.msg) + buf.Len()
   103  
   104  		now := time.Now()
   105  		if now.Before(o.arrivalTime) {
   106  			if buffered < bufsize {
   107  				buf.Write(o.msg)
   108  				return
   109  			}
   110  
   111  			// we do not buffer + return here, instead hanging the
   112  			// call (i.e. not accepting any more transportObjects)
   113  			// so that we apply back-pressure to the sender.
   114  			// this sleep should wake up same time as ticker.
   115  			time.Sleep(o.arrivalTime.Sub(now))
   116  		}
   117  
   118  		// ok, we waited our due time. now rite the buf + msg.
   119  
   120  		// drainBuf first, before we write this message.
   121  		drainBuf()
   122  
   123  		// write this message.
   124  		_, err := s.Writer.Write(o.msg)
   125  		if err != nil {
   126  			log.Error("mock_stream", err)
   127  		}
   128  	}
   129  
   130  	for {
   131  		select {
   132  		case <-proc.Closing():
   133  			return // bail out of here.
   134  
   135  		case o, ok := <-s.toDeliver:
   136  			if !ok {
   137  				return
   138  			}
   139  			deliverOrWait(o)
   140  
   141  		case <-ticker.C: // ok, due to write it out.
   142  			drainBuf()
   143  		}
   144  	}
   145  }