github.com/NVIDIA/aistore@v1.3.23-0.20240517131212-7df6609be51d/transport/pdu.go (about)

     1  // Package transport provides long-lived http/tcp connections for
     2  // intra-cluster communications (see README for details and usage example).
     3  /*
     4   * Copyright (c) 2018-2023, NVIDIA CORPORATION. All rights reserved.
     5   */
     6  package transport
     7  
     8  import (
     9  	"fmt"
    10  	"io"
    11  
    12  	"github.com/NVIDIA/aistore/cmn/cos"
    13  	"github.com/NVIDIA/aistore/cmn/debug"
    14  	"github.com/NVIDIA/aistore/memsys"
    15  )
    16  
    17  type (
    18  	pdu struct {
    19  		buf  []byte
    20  		roff int
    21  		woff int
    22  		done bool
    23  		last bool
    24  	}
    25  	spdu struct {
    26  		pdu
    27  	}
    28  	rpdu struct {
    29  		body io.Reader
    30  		pdu
    31  		flags uint64
    32  		plen  int
    33  	}
    34  )
    35  
    36  /////////
    37  // pdu //
    38  /////////
    39  
    40  func (pdu *pdu) plength() int { return pdu.woff - sizeProtoHdr } // just the payload
    41  func (pdu *pdu) slength() int { return pdu.roff - sizeProtoHdr } // payload transmitted/received so far
    42  func (pdu *pdu) rlength() int { return pdu.woff - pdu.roff }     // not yet sent/received part of the PDU
    43  
    44  func (pdu *pdu) read(b []byte) (n int) {
    45  	n = copy(b, pdu.buf[pdu.roff:pdu.woff])
    46  	pdu.roff += n
    47  	return
    48  }
    49  
    50  func (pdu *pdu) free(mm *memsys.MMSA) {
    51  	if pdu.buf != nil {
    52  		mm.Free(pdu.buf)
    53  	}
    54  }
    55  
    56  //////////
    57  // spdu //
    58  //////////
    59  
    60  func newSendPDU(buf []byte) (p *spdu) {
    61  	debug.Assert(len(buf) >= cos.KiB && len(buf) <= maxSizePDU)
    62  	p = &spdu{pdu{buf: buf}}
    63  	p.reset()
    64  	return
    65  }
    66  
    67  func (pdu *spdu) readFrom(sendoff *sendoff) (err error) {
    68  	var (
    69  		obj = &sendoff.obj
    70  		b   = pdu.buf[pdu.woff:]
    71  		n   int
    72  	)
    73  	n, err = obj.Reader.Read(b)
    74  	pdu.woff += n
    75  	pdu.done = pdu.woff == len(pdu.buf)
    76  	if err != nil {
    77  		pdu.done, pdu.last = true, true
    78  	} else if !obj.IsUnsized() && sendoff.off+int64(pdu.plength()) >= obj.Hdr.ObjAttrs.Size {
    79  		pdu.done, pdu.last = true, true
    80  	}
    81  	return
    82  }
    83  
    84  func (pdu *spdu) reset() {
    85  	pdu.roff, pdu.woff = 0, sizeProtoHdr
    86  	pdu.done, pdu.last = false, false
    87  }
    88  
    89  //////////
    90  // rpdu //
    91  //////////
    92  
    93  func newRecvPDU(body io.Reader, buf []byte) (p *rpdu) {
    94  	p = &rpdu{body: body, pdu: pdu{buf: buf}}
    95  	p.reset()
    96  	return
    97  }
    98  
    99  func (pdu *rpdu) readHdr(loghdr string) (err error) {
   100  	const fmterr = "sbrk %s: invalid PDU header [plen=%d, flags=%s]"
   101  	var n int
   102  	debug.Assert(pdu.woff == 0)
   103  	n, err = pdu.body.Read(pdu.buf[:sizeProtoHdr])
   104  	if n < sizeProtoHdr {
   105  		if err == nil {
   106  			err = fmt.Errorf("sbrk %s: failed to receive PDU header (n=%d)", loghdr, n)
   107  		}
   108  		return
   109  	}
   110  	pdu.plen, pdu.flags, err = extProtoHdr(pdu.buf, loghdr)
   111  	if err != nil {
   112  		return
   113  	}
   114  	if pdu.flags&pduFl == 0 || pdu.plen > maxSizePDU || pdu.plen < 0 {
   115  		err = fmt.Errorf(fmterr, loghdr, pdu.plen, fl2s(pdu.flags))
   116  		debug.AssertNoErr(err)
   117  		return
   118  	}
   119  	pdu.woff = sizeProtoHdr
   120  	pdu.last = pdu.flags&pduLastFl != 0
   121  	debug.Assertf(pdu.plen > 0 || (pdu.plen == 0 && pdu.last), fmterr, loghdr, pdu.plen, fl2s(pdu.flags))
   122  	return
   123  }
   124  
   125  func (pdu *rpdu) reset() {
   126  	pdu.roff, pdu.woff = sizeProtoHdr, 0
   127  	pdu.done, pdu.last = false, false
   128  }
   129  
   130  func (pdu *rpdu) readFrom() (n int, err error) {
   131  	n, err = pdu.body.Read(pdu.buf[pdu.woff : sizeProtoHdr+pdu.plen]) // NOTE: maxSizePDU
   132  	pdu.woff += n
   133  	pdu.done = pdu.plength() == pdu.plen
   134  	if err != nil {
   135  		pdu.done, pdu.last = true, true
   136  	}
   137  	return
   138  }
   139  
   140  //
   141  // misc
   142  //
   143  
   144  func fl2s(flags uint64) (s string) {
   145  	if flags&msgFl == 0 && flags&pduFl == 0 {
   146  		s += "[obj]"
   147  	} else if flags&msgFl != 0 {
   148  		s += "[msg]"
   149  	} else if flags&pduFl != 0 {
   150  		s += "[pdu]"
   151  	}
   152  	if flags&pduStreamFl != 0 {
   153  		s += "[pdu-stream]"
   154  	}
   155  	if flags&pduLastFl != 0 {
   156  		s += "[lst]"
   157  	}
   158  	return
   159  }