github.com/anacrolix/torrent@v1.61.0/peer-conn-msg-writer.go (about)

     1  package torrent
     2  
     3  import (
     4  	"bytes"
     5  	"io"
     6  	"time"
     7  
     8  	"github.com/anacrolix/chansync"
     9  	"github.com/anacrolix/log"
    10  	"github.com/anacrolix/sync"
    11  
    12  	pp "github.com/anacrolix/torrent/peer_protocol"
    13  )
    14  
    15  func (pc *PeerConn) initMessageWriter() {
    16  	w := &pc.messageWriter
    17  	*w = peerConnMsgWriter{
    18  		fillWriteBuffer: func() {
    19  			pc.locker().Lock()
    20  			defer pc.locker().Unlock()
    21  			if pc.closed.IsSet() {
    22  				return
    23  			}
    24  			pc.fillWriteBuffer()
    25  		},
    26  		closed: &pc.closed,
    27  		logger: pc.logger,
    28  		w:      pc.w,
    29  		keepAlive: func() bool {
    30  			pc.locker().RLock()
    31  			defer pc.locker().RUnlock()
    32  			return pc.useful()
    33  		},
    34  		writeBuffer: new(peerConnMsgWriterBuffer),
    35  	}
    36  }
    37  
    38  func (pc *PeerConn) startMessageWriter() {
    39  	pc.initMessageWriter()
    40  	go pc.messageWriterRunner()
    41  }
    42  
    43  func (pc *PeerConn) messageWriterRunner() {
    44  	defer pc.locker().Unlock()
    45  	defer pc.close()
    46  	defer pc.locker().Lock()
    47  	pc.messageWriter.run(pc.t.cl.config.KeepAliveTimeout)
    48  }
    49  
    50  type peerConnMsgWriterBuffer struct {
    51  	// The number of bytes in the buffer that are part of a piece message. When
    52  	// the whole buffer is written, we can count this many bytes.
    53  	pieceDataBytes int
    54  	bytes.Buffer
    55  }
    56  
    57  type peerConnMsgWriter struct {
    58  	// Must not be called with the local mutex held, as it will call back into the write method.
    59  	fillWriteBuffer func()
    60  	closed          *chansync.SetOnce
    61  	logger          log.Logger
    62  	w               io.Writer
    63  	keepAlive       func() bool
    64  
    65  	mu        sync.Mutex
    66  	writeCond chansync.BroadcastCond
    67  	// Pointer so we can swap with the "front buffer".
    68  	writeBuffer *peerConnMsgWriterBuffer
    69  
    70  	totalWriteDuration    time.Duration
    71  	totalBytesWritten     int64
    72  	totalDataBytesWritten int64
    73  	dataUploadRate        float64
    74  }
    75  
    76  // Routine that writes to the peer. Some of what to write is buffered by
    77  // activity elsewhere in the Client, and some is determined locally when the
    78  // connection is writable.
    79  func (cn *peerConnMsgWriter) run(keepAliveTimeout time.Duration) {
    80  	lastWrite := time.Now()
    81  	keepAliveTimer := time.NewTimer(keepAliveTimeout)
    82  	frontBuf := new(peerConnMsgWriterBuffer)
    83  	for {
    84  		if cn.closed.IsSet() {
    85  			return
    86  		}
    87  		cn.fillWriteBuffer()
    88  		keepAlive := cn.keepAlive()
    89  		cn.mu.Lock()
    90  		if cn.writeBuffer.Len() == 0 && time.Since(lastWrite) >= keepAliveTimeout && keepAlive {
    91  			cn.writeBuffer.Write(pp.Message{Keepalive: true}.MustMarshalBinary())
    92  			torrent.Add("written keepalives", 1)
    93  		}
    94  		if cn.writeBuffer.Len() == 0 {
    95  			writeCond := cn.writeCond.Signaled()
    96  			cn.mu.Unlock()
    97  			select {
    98  			case <-cn.closed.Done():
    99  			case <-writeCond:
   100  			case <-keepAliveTimer.C:
   101  			}
   102  			continue
   103  		}
   104  		// Flip the buffers.
   105  		frontBuf, cn.writeBuffer = cn.writeBuffer, frontBuf
   106  		cn.mu.Unlock()
   107  		if frontBuf.Len() == 0 {
   108  			panic("expected non-empty front buffer")
   109  		}
   110  		var err error
   111  		startedWriting := time.Now()
   112  		startingBufLen := frontBuf.Len()
   113  		for frontBuf.Len() != 0 {
   114  			next := frontBuf.Bytes()
   115  			var n int
   116  			n, err = cn.w.Write(next)
   117  			frontBuf.Next(n)
   118  			if err == nil && n != len(next) {
   119  				panic("expected full write")
   120  			}
   121  			if err != nil {
   122  				break
   123  			}
   124  		}
   125  		if err != nil {
   126  			cn.logger.WithDefaultLevel(log.Debug).Printf("error writing: %v", err)
   127  			return
   128  		}
   129  		// Track what was sent and how long it took.
   130  		writeDuration := time.Since(startedWriting)
   131  		cn.mu.Lock()
   132  		cn.dataUploadRate = float64(frontBuf.pieceDataBytes) / writeDuration.Seconds()
   133  		cn.totalWriteDuration += writeDuration
   134  		cn.totalBytesWritten += int64(startingBufLen)
   135  		cn.totalDataBytesWritten += int64(frontBuf.pieceDataBytes)
   136  		cn.mu.Unlock()
   137  		frontBuf.pieceDataBytes = 0
   138  		lastWrite = time.Now()
   139  		keepAliveTimer.Reset(keepAliveTimeout)
   140  	}
   141  }
   142  
   143  func (cn *peerConnMsgWriter) writeToBuffer(msg pp.Message) (err error) {
   144  	originalLen := cn.writeBuffer.Len()
   145  	defer func() {
   146  		if err != nil {
   147  			// Since an error occurred during buffer write, revert buffer to its original state before the write.
   148  			cn.writeBuffer.Truncate(originalLen)
   149  		}
   150  	}()
   151  	err = msg.WriteTo(cn.writeBuffer)
   152  	if err == nil {
   153  		cn.writeBuffer.pieceDataBytes += len(msg.Piece)
   154  	}
   155  	return
   156  }
   157  
   158  func (cn *peerConnMsgWriter) write(msg pp.Message) bool {
   159  	cn.mu.Lock()
   160  	defer cn.mu.Unlock()
   161  	cn.writeToBuffer(msg)
   162  	cn.writeCond.Broadcast()
   163  	return !cn.writeBufferFull()
   164  }
   165  
   166  func (cn *peerConnMsgWriter) writeBufferFull() bool {
   167  	return cn.writeBuffer.Len() >= writeBufferHighWaterLen
   168  }