github.com/NVIDIA/aistore@v1.3.23-0.20240517131212-7df6609be51d/transport/obj_test.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-2024, NVIDIA CORPORATION. All rights reserved.
     5   */
     6  package transport_test
     7  
     8  // How to run:
     9  //
    10  // 1) run all unit tests
    11  // go test -v
    12  //
    13  // 2) run tests matching "Multi" with debug enabled:
    14  // go test -v -run=Multi -tags=debug
    15  
    16  import (
    17  	"encoding/binary"
    18  	"flag"
    19  	"fmt"
    20  	"io"
    21  	"math"
    22  	"math/rand"
    23  	"net/http"
    24  	"net/http/httptest"
    25  	"os"
    26  	"path"
    27  	"reflect"
    28  	"strconv"
    29  	"sync"
    30  	"testing"
    31  	"time"
    32  
    33  	"github.com/NVIDIA/aistore/3rdparty/golang/mux"
    34  	"github.com/NVIDIA/aistore/api/apc"
    35  	"github.com/NVIDIA/aistore/cmn"
    36  	"github.com/NVIDIA/aistore/cmn/atomic"
    37  	"github.com/NVIDIA/aistore/cmn/cos"
    38  	"github.com/NVIDIA/aistore/cmn/mono"
    39  	"github.com/NVIDIA/aistore/memsys"
    40  	"github.com/NVIDIA/aistore/tools"
    41  	"github.com/NVIDIA/aistore/tools/tassert"
    42  	"github.com/NVIDIA/aistore/tools/tlog"
    43  	"github.com/NVIDIA/aistore/transport"
    44  )
    45  
    46  const (
    47  	lorem = `Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut
    48  labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat.`
    49  	duis = `Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur.
    50  Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum.`
    51  	et = `Et harum quidem rerum facilis est et expedita distinctio. Nam libero tempore, cum soluta nobis est
    52  eligendi optio, cumque nihil impedit, quo minus id, quod maxime placeat, facere possimus, omnis voluptas assumenda est, omnis dolor repellendus.`
    53  	temporibus = `Temporibus autem quibusdam et aut officiis debitis aut rerum necessitatibus saepe eveniet,
    54  ut et voluptates repudiandae sint et molestiae non-recusandae.`
    55  	text = lorem + duis + et + temporibus
    56  )
    57  
    58  type dummyStatsTracker struct{}
    59  
    60  // interface guard
    61  var _ cos.StatsUpdater = (*dummyStatsTracker)(nil)
    62  
    63  func (*dummyStatsTracker) Add(string, int64)         {}
    64  func (*dummyStatsTracker) Inc(string)                {}
    65  func (*dummyStatsTracker) Get(string) int64          { return 0 }
    66  func (*dummyStatsTracker) AddMany(...cos.NamedVal64) {}
    67  
    68  var (
    69  	objmux   *mux.ServeMux
    70  	duration time.Duration // test duration
    71  )
    72  
    73  func TestMain(t *testing.M) {
    74  	var (
    75  		d   string
    76  		err error
    77  	)
    78  	flag.StringVar(&d, "duration", "30s", "test duration")
    79  	flag.Parse()
    80  	if duration, err = time.ParseDuration(d); err != nil {
    81  		cos.Exitf("Invalid duration %q", d)
    82  	}
    83  
    84  	config := cmn.GCO.BeginUpdate()
    85  	config.Transport.MaxHeaderSize = memsys.PageSize
    86  	config.Transport.IdleTeardown = cos.Duration(time.Second)
    87  	config.Transport.QuiesceTime = cos.Duration(10 * time.Second)
    88  	config.Log.Level = "3"
    89  	cmn.GCO.CommitUpdate(config)
    90  	sc := transport.Init(&dummyStatsTracker{}, config)
    91  	go sc.Run()
    92  
    93  	objmux = mux.NewServeMux()
    94  	path := transport.ObjURLPath("")
    95  	objmux.HandleFunc(path, transport.RxAnyStream)
    96  	objmux.HandleFunc(path+"/", transport.RxAnyStream)
    97  
    98  	os.Exit(t.Run())
    99  }
   100  
   101  func Example_headers() {
   102  	f := func(_ http.ResponseWriter, r *http.Request) {
   103  		body, err := io.ReadAll(r.Body)
   104  		if err != nil {
   105  			panic(err)
   106  		}
   107  		if len(body) == 0 {
   108  			return
   109  		}
   110  		var (
   111  			hdr       transport.ObjHdr
   112  			hlen, off int
   113  		)
   114  		for {
   115  			hlen = int(binary.BigEndian.Uint64(body[off:]))
   116  			off += 16 // hlen and hlen-checksum
   117  			hdr = transport.ExtObjHeader(body[off:], hlen)
   118  
   119  			if transport.ReservedOpcode(hdr.Opcode) {
   120  				break
   121  			}
   122  
   123  			fmt.Printf("%+v (%d)\n", hdr, hlen)
   124  			off += hlen + int(hdr.ObjAttrs.Size)
   125  		}
   126  	}
   127  
   128  	ts := httptest.NewServer(http.HandlerFunc(f))
   129  	defer ts.Close()
   130  
   131  	httpclient := transport.NewIntraDataClient()
   132  	stream := transport.NewObjStream(httpclient, ts.URL, cos.GenTie(), nil)
   133  
   134  	sendText(stream, lorem, duis)
   135  	stream.Fin()
   136  
   137  	// Output:
   138  	// {Bck:s3://@uuid#namespace/abc ObjName:X SID: Opaque:[] ObjAttrs:{Cksum:xxhash[h1] CustomMD:map[] Ver:1 Atime:663346294 Size:231} Opcode:0} (69)
   139  	// {Bck:ais://abracadabra ObjName:p/q/s SID: Opaque:[49 50 51] ObjAttrs:{Cksum:xxhash[h2] CustomMD:map[xx:11 yy:22] Ver:222222222222222222222222 Atime:663346294 Size:213} Opcode:0} (110)
   140  }
   141  
   142  func sendText(stream *transport.Stream, txt1, txt2 string) {
   143  	var wg sync.WaitGroup
   144  	cb := func(*transport.ObjHdr, io.ReadCloser, any, error) {
   145  		wg.Done()
   146  	}
   147  	sgl1 := memsys.PageMM().NewSGL(0)
   148  	sgl1.Write([]byte(txt1))
   149  	hdr := transport.ObjHdr{
   150  		Bck: cmn.Bck{
   151  			Name:     "abc",
   152  			Provider: apc.AWS,
   153  			Ns:       cmn.Ns{UUID: "uuid", Name: "namespace"},
   154  		},
   155  		ObjName: "X",
   156  		ObjAttrs: cmn.ObjAttrs{
   157  			Size:  sgl1.Size(),
   158  			Atime: 663346294,
   159  			Cksum: cos.NewCksum(cos.ChecksumXXHash, "h1"),
   160  			Ver:   "1",
   161  		},
   162  		Opaque: nil,
   163  	}
   164  	wg.Add(1)
   165  	stream.Send(&transport.Obj{Hdr: hdr, Reader: sgl1, Callback: cb})
   166  	wg.Wait()
   167  
   168  	sgl2 := memsys.PageMM().NewSGL(0)
   169  	sgl2.Write([]byte(txt2))
   170  	hdr = transport.ObjHdr{
   171  		Bck: cmn.Bck{
   172  			Name:     "abracadabra",
   173  			Provider: apc.AIS,
   174  			Ns:       cmn.NsGlobal,
   175  		},
   176  		ObjName: "p/q/s",
   177  		ObjAttrs: cmn.ObjAttrs{
   178  			Size:  sgl2.Size(),
   179  			Atime: 663346294,
   180  			Cksum: cos.NewCksum(cos.ChecksumXXHash, "h2"),
   181  			Ver:   "222222222222222222222222",
   182  		},
   183  		Opaque: []byte{'1', '2', '3'},
   184  	}
   185  	hdr.ObjAttrs.SetCustomMD(cos.StrKVs{"xx": "11", "yy": "22"})
   186  	wg.Add(1)
   187  	stream.Send(&transport.Obj{Hdr: hdr, Reader: sgl2, Callback: cb})
   188  	wg.Wait()
   189  }
   190  
   191  func Example_obj() {
   192  	receive := func(hdr *transport.ObjHdr, objReader io.Reader, err error) error {
   193  		cos.Assert(err == nil)
   194  		object, err := io.ReadAll(objReader)
   195  		if err != nil {
   196  			panic(err)
   197  		}
   198  		if int64(len(object)) != hdr.ObjAttrs.Size {
   199  			panic(fmt.Sprintf("size %d != %d", len(object), hdr.ObjAttrs.Size))
   200  		}
   201  		fmt.Printf("%s...\n", string(object[:16]))
   202  		return nil
   203  	}
   204  	ts := httptest.NewServer(objmux)
   205  	defer ts.Close()
   206  	trname := "dummy-obj"
   207  	err := transport.Handle(trname, receive)
   208  	if err != nil {
   209  		fmt.Println(err)
   210  		return
   211  	}
   212  	httpclient := transport.NewIntraDataClient()
   213  	stream := transport.NewObjStream(httpclient, ts.URL+transport.ObjURLPath(trname), cos.GenTie(), nil)
   214  	sendText(stream, lorem, duis)
   215  	sendText(stream, et, temporibus)
   216  	stream.Fin()
   217  
   218  	// Output:
   219  	// Lorem ipsum dolo...
   220  	// Duis aute irure ...
   221  	// Et harum quidem ...
   222  	// Temporibus autem...
   223  }
   224  
   225  // test random streaming
   226  func TestOneStream(t *testing.T) {
   227  	tools.CheckSkip(t, &tools.SkipTestArgs{Long: true})
   228  	ts := httptest.NewServer(objmux)
   229  	defer ts.Close()
   230  
   231  	streamWriteUntil(t, 55, nil, ts, nil, nil, false /*compress*/, true /*PDU*/)
   232  	printNetworkStats()
   233  }
   234  
   235  func TestMultiStream(t *testing.T) {
   236  	tools.CheckSkip(t, &tools.SkipTestArgs{Long: true})
   237  
   238  	tlog.Logf("Duration %v\n", duration)
   239  	ts := httptest.NewServer(objmux)
   240  	defer ts.Close()
   241  
   242  	wg := &sync.WaitGroup{}
   243  	netstats := make(map[string]transport.RxStats)
   244  	lock := &sync.Mutex{}
   245  	for i := range 16 {
   246  		wg.Add(1)
   247  		go streamWriteUntil(t, i, wg, ts, netstats, lock, false /*compress*/, false /*PDU*/)
   248  	}
   249  	wg.Wait()
   250  	compareNetworkStats(netstats)
   251  }
   252  
   253  func printNetworkStats() {
   254  	netstats := transport.GetRxStats()
   255  	for trname, eps := range netstats {
   256  		for uid, stats := range eps { // RxStats by session ID
   257  			xx, sessID := transport.UID2SessID(uid)
   258  			fmt.Printf("recv$ %s[%d:%d]: offset=%d, num=%d\n",
   259  				trname, xx, sessID, stats.Offset.Load(), stats.Num.Load())
   260  		}
   261  	}
   262  }
   263  
   264  func compareNetworkStats(netstats1 map[string]transport.RxStats) {
   265  	netstats2 := transport.GetRxStats()
   266  	for trname, eps2 := range netstats2 {
   267  		eps1, ok := netstats1[trname]
   268  		for uid, stats2 := range eps2 { // RxStats by session ID
   269  			xx, sessID := transport.UID2SessID(uid)
   270  			fmt.Printf("recv$ %s[%d:%d]: offset=%d, num=%d\n", trname, xx, sessID,
   271  				stats2.Offset.Load(), stats2.Num.Load())
   272  			if ok {
   273  				stats1, ok := eps1[sessID]
   274  				if ok {
   275  					fmt.Printf("send$ %s[%d]: offset=%d, num=%d\n",
   276  						trname, sessID, stats1.Offset.Load(), stats1.Num.Load())
   277  				} else {
   278  					fmt.Printf("send$ %s[%d]: -- not present --\n", trname, sessID)
   279  				}
   280  			} else {
   281  				fmt.Printf("send$ %s[%d]: -- not present --\n", trname, sessID)
   282  			}
   283  		}
   284  	}
   285  }
   286  
   287  func TestMultipleNetworks(t *testing.T) {
   288  	totalRecv, recvFunc := makeRecvFunc(t)
   289  
   290  	streams := make([]*transport.Stream, 0, 10)
   291  	for idx := range 10 {
   292  		ts := httptest.NewServer(objmux)
   293  		defer ts.Close()
   294  		trname := "endpoint" + strconv.Itoa(idx)
   295  		err := transport.Handle(trname, recvFunc)
   296  		tassert.CheckFatal(t, err)
   297  		defer transport.Unhandle(trname)
   298  
   299  		httpclient := transport.NewIntraDataClient()
   300  		url := ts.URL + transport.ObjURLPath(trname)
   301  		streams = append(streams, transport.NewObjStream(httpclient, url, cos.GenTie(), nil))
   302  	}
   303  
   304  	totalSend := int64(0)
   305  	random := newRand(mono.NanoTime())
   306  	for _, stream := range streams {
   307  		hdr, reader := makeRandReader(random, false)
   308  		totalSend += hdr.ObjAttrs.Size
   309  		stream.Send(&transport.Obj{Hdr: hdr, Reader: reader})
   310  	}
   311  
   312  	for _, stream := range streams {
   313  		stream.Fin()
   314  	}
   315  	time.Sleep(3 * time.Second) // FIN has been sent but not necessarily received
   316  
   317  	if *totalRecv != totalSend {
   318  		t.Fatalf("total received bytes %d is different from expected: %d", *totalRecv, totalSend)
   319  	}
   320  }
   321  
   322  func TestSendCallback(t *testing.T) {
   323  	objectCnt := 10000
   324  	if testing.Short() {
   325  		objectCnt = 1000
   326  	}
   327  
   328  	ts := httptest.NewServer(objmux)
   329  	defer ts.Close()
   330  
   331  	totalRecv, recvFunc := makeRecvFunc(t)
   332  	trname := "callback"
   333  	err := transport.Handle(trname, recvFunc)
   334  	tassert.CheckFatal(t, err)
   335  	defer transport.Unhandle(trname)
   336  	httpclient := transport.NewIntraDataClient()
   337  	url := ts.URL + transport.ObjURLPath(trname)
   338  	stream := transport.NewObjStream(httpclient, url, cos.GenTie(), nil)
   339  
   340  	var (
   341  		totalSend int64
   342  		mu        sync.Mutex
   343  		posted    = make([]*randReader, objectCnt)
   344  	)
   345  	random := newRand(mono.NanoTime())
   346  	for idx := range len(posted) {
   347  		hdr, rr := makeRandReader(random, false)
   348  		mu.Lock()
   349  		posted[idx] = rr
   350  		mu.Unlock()
   351  		rrc := &randReaderCtx{t, rr, posted, &mu, idx}
   352  		totalSend += hdr.ObjAttrs.Size
   353  		stream.Send(&transport.Obj{Hdr: hdr, Reader: rr, Callback: rrc.sentCallback})
   354  	}
   355  	stream.Fin()
   356  
   357  	for idx := range posted {
   358  		if posted[idx] != nil {
   359  			t.Errorf("sent-callback %d never fired", idx)
   360  		}
   361  	}
   362  	if *totalRecv != totalSend {
   363  		t.Fatalf("total received bytes %d is different from expected: %d", *totalRecv, totalSend)
   364  	}
   365  }
   366  
   367  func TestObjAttrs(t *testing.T) {
   368  	testAttrs := []cmn.ObjAttrs{
   369  		{
   370  			Size:  1024,
   371  			Atime: 1024,
   372  			Cksum: cos.NewCksum("", ""),
   373  			Ver:   "102.44",
   374  		},
   375  		{
   376  			Size:  1024,
   377  			Atime: math.MaxInt64,
   378  			Cksum: cos.NewCksum(cos.ChecksumXXHash, "120421"),
   379  			Ver:   "102.44",
   380  		},
   381  		{
   382  			Size:  0,
   383  			Atime: 0,
   384  			Cksum: cos.NewCksum(cos.ChecksumNone, "120421"),
   385  			Ver:   "",
   386  		},
   387  	}
   388  
   389  	ts := httptest.NewServer(objmux)
   390  	defer ts.Close()
   391  
   392  	var receivedCount atomic.Int64
   393  	recvFunc := func(hdr *transport.ObjHdr, objReader io.Reader, err error) error {
   394  		cos.Assert(err == nil)
   395  
   396  		idx := hdr.Opaque[0]
   397  		cos.AssertMsg(hdr.Bck.IsAIS(), "expecting ais bucket")
   398  		cos.Assertf(reflect.DeepEqual(testAttrs[idx], hdr.ObjAttrs),
   399  			"attrs are not equal: %v; %v;", testAttrs[idx], hdr.ObjAttrs)
   400  
   401  		written, err := io.Copy(io.Discard, objReader)
   402  		cos.Assert(err == nil)
   403  		cos.Assertf(written == hdr.ObjAttrs.Size, "written: %d, expected: %d", written, hdr.ObjAttrs.Size)
   404  
   405  		receivedCount.Inc()
   406  		return nil
   407  	}
   408  	trname := "objattrs"
   409  	err := transport.Handle(trname, recvFunc)
   410  	tassert.CheckFatal(t, err)
   411  	defer transport.Unhandle(trname)
   412  	httpclient := transport.NewIntraDataClient()
   413  	url := ts.URL + transport.ObjURLPath(trname)
   414  	stream := transport.NewObjStream(httpclient, url, cos.GenTie(), nil)
   415  
   416  	random := newRand(mono.NanoTime())
   417  	for idx, attrs := range testAttrs {
   418  		var (
   419  			reader io.ReadCloser
   420  			hdr    = transport.ObjHdr{
   421  				Bck: cmn.Bck{
   422  					Provider: apc.AIS,
   423  				},
   424  				ObjAttrs: attrs,
   425  				Opaque:   []byte{byte(idx)},
   426  			}
   427  		)
   428  		slab, err := memsys.PageMM().GetSlab(memsys.PageSize)
   429  		if err != nil {
   430  			t.Fatal(err)
   431  		}
   432  		if hdr.ObjAttrs.Size > 0 {
   433  			reader = newRandReader(random, hdr, slab)
   434  		}
   435  		if err := stream.Send(&transport.Obj{Hdr: hdr, Reader: reader}); err != nil {
   436  			t.Fatal(err)
   437  		}
   438  	}
   439  	stream.Fin()
   440  	if receivedCount.Load() != int64(len(testAttrs)) {
   441  		t.Fatalf("invalid received count: %d, expected: %d", receivedCount.Load(), len(testAttrs))
   442  	}
   443  }
   444  
   445  func receive10G(hdr *transport.ObjHdr, objReader io.Reader, err error) error {
   446  	cos.Assert(err == nil || cos.IsEOF(err))
   447  	written, _ := io.Copy(io.Discard, objReader)
   448  	cos.Assert(written == hdr.ObjAttrs.Size)
   449  	return nil
   450  }
   451  
   452  func TestCompressedOne(t *testing.T) {
   453  	trname := "cmpr-one"
   454  	config := cmn.GCO.BeginUpdate()
   455  	config.Transport.LZ4BlockMaxSize = 256 * cos.KiB
   456  	config.Transport.IdleTeardown = cos.Duration(time.Second)
   457  	config.Transport.QuiesceTime = cos.Duration(8 * time.Second)
   458  	cmn.GCO.CommitUpdate(config)
   459  	if err := config.Transport.Validate(); err != nil {
   460  		tassert.CheckFatal(t, err)
   461  	}
   462  
   463  	ts := httptest.NewServer(objmux)
   464  	defer ts.Close()
   465  
   466  	err := transport.Handle(trname, receive10G, true /*with Rx stats*/)
   467  	tassert.CheckFatal(t, err)
   468  	defer transport.Unhandle(trname)
   469  
   470  	httpclient := transport.NewIntraDataClient()
   471  	url := ts.URL + transport.ObjURLPath(trname)
   472  	t.Setenv("AIS_STREAM_BURST_NUM", "2")
   473  	stream := transport.NewObjStream(httpclient, url, cos.GenTie(), &transport.Extra{Compression: apc.CompressAlways})
   474  
   475  	slab, _ := memsys.PageMM().GetSlab(memsys.MaxPageSlabSize)
   476  	random := newRand(mono.NanoTime())
   477  	buf := slab.Alloc()
   478  	_, _ = random.Read(buf)
   479  	hdr := genStaticHeader(random)
   480  	size, prevsize, num, numhdr, numGs := int64(0), int64(0), 0, 0, int64(16)
   481  	if testing.Short() {
   482  		numGs = 2
   483  	}
   484  	for size < cos.GiB*numGs {
   485  		if num%7 == 0 { // header-only
   486  			hdr.ObjAttrs.Size = 0
   487  			stream.Send(&transport.Obj{Hdr: hdr})
   488  			numhdr++
   489  		} else {
   490  			var reader io.ReadCloser
   491  			if num%3 == 0 {
   492  				hdr.ObjAttrs.Size = int64(random.Intn(100) + 1)
   493  				// fully random to prevent compression
   494  				reader = io.NopCloser(&io.LimitedReader{R: random, N: hdr.ObjAttrs.Size})
   495  			} else {
   496  				hdr.ObjAttrs.Size = int64(random.Intn(cos.GiB) + 1)
   497  				reader = &randReader{buf: buf, hdr: hdr, clone: true}
   498  			}
   499  			stream.Send(&transport.Obj{Hdr: hdr, Reader: reader})
   500  		}
   501  		num++
   502  		size += hdr.ObjAttrs.Size
   503  		if size-prevsize >= cos.GiB*4 {
   504  			stats := stream.GetStats()
   505  			tlog.Logf("%s: %d GiB compression-ratio=%.2f\n", stream, size/cos.GiB, stats.CompressionRatio())
   506  			prevsize = size
   507  		}
   508  	}
   509  	stream.Fin()
   510  	stats := stream.GetStats()
   511  
   512  	slab.Free(buf)
   513  
   514  	fmt.Printf("send$ %s: offset=%d, num=%d(%d/%d), compression-ratio=%.2f\n",
   515  		stream, stats.Offset.Load(), stats.Num.Load(), num, numhdr, stats.CompressionRatio())
   516  
   517  	printNetworkStats()
   518  }
   519  
   520  func TestDryRun(t *testing.T) {
   521  	tools.CheckSkip(t, &tools.SkipTestArgs{Long: true})
   522  
   523  	t.Setenv("AIS_STREAM_DRY_RUN", "true")
   524  
   525  	stream := transport.NewObjStream(nil, "dummy/null", cos.GenTie(), nil)
   526  
   527  	random := newRand(mono.NanoTime())
   528  	sgl := memsys.PageMM().NewSGL(cos.MiB)
   529  	defer sgl.Free()
   530  	buf, slab := memsys.PageMM().AllocSize(cos.KiB * 128)
   531  	defer slab.Free(buf)
   532  	for sgl.Len() < cos.MiB {
   533  		random.Read(buf)
   534  		sgl.Write(buf)
   535  	}
   536  
   537  	size, num, prevsize := int64(0), 0, int64(0)
   538  	hdr := genStaticHeader(random)
   539  	total := int64(cos.TiB)
   540  	if testing.Short() {
   541  		total = cos.TiB / 4
   542  	}
   543  
   544  	for size < total {
   545  		hdr.ObjAttrs.Size = cos.KiB * 128
   546  		for i := range cos.MiB / hdr.ObjAttrs.Size {
   547  			reader := memsys.NewReader(sgl)
   548  			reader.Seek(i*hdr.ObjAttrs.Size, io.SeekStart)
   549  
   550  			stream.Send(&transport.Obj{Hdr: hdr, Reader: reader})
   551  			num++
   552  			size += hdr.ObjAttrs.Size
   553  			if size-prevsize >= cos.GiB*100 {
   554  				prevsize = size
   555  				tlog.Logf("[dry]: %d GiB\n", size/cos.GiB)
   556  			}
   557  		}
   558  	}
   559  	stream.Fin()
   560  	stats := stream.GetStats()
   561  
   562  	fmt.Printf("[dry]: offset=%d, num=%d(%d)\n", stats.Offset.Load(), stats.Num.Load(), num)
   563  }
   564  
   565  func TestCompletionCount(t *testing.T) {
   566  	tools.CheckSkip(t, &tools.SkipTestArgs{Long: true})
   567  	var (
   568  		numSent                   int64
   569  		numCompleted, numReceived atomic.Int64
   570  	)
   571  
   572  	receive := func(hdr *transport.ObjHdr, objReader io.Reader, err error) error {
   573  		cos.Assert(err == nil)
   574  		written, _ := io.Copy(io.Discard, objReader)
   575  		cos.Assert(written == hdr.ObjAttrs.Size)
   576  		numReceived.Inc()
   577  		return nil
   578  	}
   579  	callback := func(_ *transport.ObjHdr, _ io.ReadCloser, _ any, _ error) {
   580  		numCompleted.Inc()
   581  	}
   582  
   583  	ts := httptest.NewServer(objmux)
   584  	defer ts.Close()
   585  
   586  	trname := "cmpl-cnt"
   587  	err := transport.Handle(trname, receive)
   588  	tassert.CheckFatal(t, err)
   589  	defer transport.Unhandle(trname)
   590  	httpclient := transport.NewIntraDataClient()
   591  	url := ts.URL + transport.ObjURLPath(trname)
   592  	t.Setenv("AIS_STREAM_BURST_NUM", "256")
   593  	stream := transport.NewObjStream(httpclient, url, cos.GenTie(), nil) // provide for sizeable queue at any point
   594  	random := newRand(mono.NanoTime())
   595  	rem := int64(0)
   596  	for idx := range 10000 {
   597  		if idx%7 == 0 {
   598  			hdr := genStaticHeader(random)
   599  			hdr.ObjAttrs.Size = 0
   600  			hdr.Opaque = []byte(strconv.FormatInt(104729*int64(idx), 10))
   601  			stream.Send(&transport.Obj{Hdr: hdr, Callback: callback})
   602  			rem = random.Int63() % 13
   603  		} else {
   604  			hdr, rr := makeRandReader(random, false)
   605  			stream.Send(&transport.Obj{Hdr: hdr, Reader: rr, Callback: callback})
   606  		}
   607  		numSent++
   608  		if numSent > 5000 && rem == 3 {
   609  			stream.Stop()
   610  			break
   611  		}
   612  	}
   613  	// collect all pending completions until timeout
   614  	started := time.Now()
   615  	for numCompleted.Load() < numSent {
   616  		time.Sleep(time.Millisecond * 10)
   617  		if time.Since(started) > time.Second*10 {
   618  			break
   619  		}
   620  	}
   621  	if numSent == numCompleted.Load() {
   622  		tlog.Logf("sent %d = %d completed, %d received\n", numSent, numCompleted.Load(), numReceived.Load())
   623  	} else {
   624  		t.Fatalf("sent %d != %d completed\n", numSent, numCompleted.Load())
   625  	}
   626  }
   627  
   628  //
   629  // test helpers
   630  //
   631  
   632  func streamWriteUntil(t *testing.T, ii int, wg *sync.WaitGroup, ts *httptest.Server,
   633  	netstats map[string]transport.RxStats, lock sync.Locker, compress, usePDU bool) {
   634  	if wg != nil {
   635  		defer wg.Done()
   636  	}
   637  	totalRecv, recvFunc := makeRecvFunc(t)
   638  	trname := fmt.Sprintf("rand-rx-%d", ii)
   639  	err := transport.Handle(trname, recvFunc, true /* with Rx stats */)
   640  	tassert.CheckFatal(t, err)
   641  	defer transport.Unhandle(trname)
   642  
   643  	if compress {
   644  		config := cmn.GCO.BeginUpdate()
   645  		config.Transport.LZ4BlockMaxSize = cos.KiB * 256
   646  		cmn.GCO.CommitUpdate(config)
   647  		if err := config.Transport.Validate(); err != nil {
   648  			tassert.CheckFatal(t, err)
   649  		}
   650  	}
   651  
   652  	httpclient := transport.NewIntraDataClient()
   653  	url := ts.URL + transport.ObjURLPath(trname)
   654  	var extra *transport.Extra
   655  	if compress || usePDU {
   656  		extra = &transport.Extra{}
   657  		if compress {
   658  			extra.Compression = apc.CompressAlways
   659  		}
   660  		if usePDU {
   661  			extra.SizePDU = memsys.DefaultBufSize
   662  		}
   663  	}
   664  	stream := transport.NewObjStream(httpclient, url, cos.GenTie(), extra)
   665  	trname, sessID := stream.ID()
   666  	now := time.Now()
   667  
   668  	random := newRand(mono.NanoTime())
   669  	size, num, prevsize := int64(0), 0, int64(0)
   670  	runFor := duration
   671  	if testing.Short() {
   672  		runFor = 10 * time.Second
   673  	}
   674  	var randReader *randReader
   675  	for time.Since(now) < runFor {
   676  		obj := transport.AllocSend()
   677  		obj.Hdr, randReader = makeRandReader(random, usePDU)
   678  		obj.Reader = randReader
   679  		stream.Send(obj)
   680  		num++
   681  		if obj.IsUnsized() {
   682  			size += randReader.offEOF
   683  		} else {
   684  			size += obj.Hdr.ObjAttrs.Size
   685  		}
   686  		if size-prevsize >= cos.GiB*4 {
   687  			tlog.Logf("%s: %d GiB\n", stream, size/cos.GiB)
   688  			prevsize = size
   689  			if random.Int63()%7 == 0 {
   690  				time.Sleep(time.Second * 2) // simulate occasional timeout
   691  			}
   692  		}
   693  	}
   694  	stream.Fin()
   695  	stats := stream.GetStats()
   696  	if netstats == nil {
   697  		reason, termErr := stream.TermInfo()
   698  		tassert.Errorf(t, reason != "", "expecting reason for termination")
   699  		fmt.Printf("send$ %s[%d]: offset=%d, num=%d(%d), term(%q, %v)\n",
   700  			trname, sessID, stats.Offset.Load(), stats.Num.Load(), num, reason, termErr)
   701  	} else {
   702  		lock.Lock()
   703  		eps := make(transport.RxStats)
   704  		eps[uint64(sessID)] = &stats
   705  		netstats[trname] = eps
   706  		lock.Unlock()
   707  	}
   708  
   709  	if *totalRecv != size {
   710  		t.Errorf("total received bytes %d is different from expected: %d", *totalRecv, size)
   711  		return
   712  	}
   713  }
   714  
   715  func makeRecvFunc(t *testing.T) (*int64, transport.RecvObj) {
   716  	totalReceived := new(int64)
   717  	return totalReceived, func(hdr *transport.ObjHdr, objReader io.Reader, err error) error {
   718  		cos.Assert(err == nil || cos.IsEOF(err))
   719  		written, err := io.Copy(io.Discard, objReader)
   720  		if err != nil && !cos.IsEOF(err) {
   721  			tassert.CheckFatal(t, err)
   722  		}
   723  		if written != hdr.ObjAttrs.Size && !hdr.IsUnsized() {
   724  			t.Fatalf("size %d != %d", written, hdr.ObjAttrs.Size)
   725  		}
   726  		*totalReceived += written
   727  		return nil
   728  	}
   729  }
   730  
   731  func newRand(seed int64) *rand.Rand {
   732  	src := rand.NewSource(seed)
   733  	random := rand.New(src)
   734  	return random
   735  }
   736  
   737  func genStaticHeader(random *rand.Rand) (hdr transport.ObjHdr) {
   738  	hdr.Bck = cmn.Bck{
   739  		Name:     "a",
   740  		Provider: apc.AIS,
   741  	}
   742  	hdr.ObjName = strconv.FormatInt(random.Int63(), 10)
   743  	hdr.Opaque = []byte("123456789abcdef")
   744  	hdr.ObjAttrs.Size = cos.GiB
   745  	hdr.ObjAttrs.SetCustomKey(strconv.FormatInt(random.Int63(), 10), "d")
   746  	hdr.ObjAttrs.SetCustomKey("e", "")
   747  	hdr.ObjAttrs.SetCksum(cos.ChecksumXXHash, "xxhash")
   748  	return
   749  }
   750  
   751  func genRandomHeader(random *rand.Rand, usePDU bool) (hdr transport.ObjHdr) {
   752  	x := random.Int63()
   753  	hdr.Bck.Name = strconv.FormatInt(x, 10)
   754  	hdr.Bck.Provider = apc.AIS
   755  	hdr.ObjName = path.Join(hdr.Bck.Name, strconv.FormatInt(math.MaxInt64-x, 10))
   756  
   757  	pos := x % int64(len(text))
   758  	hdr.Opaque = []byte(text[int(pos):])
   759  
   760  	// test a variety of payload sizes
   761  	y := x & 3
   762  	s := strconv.FormatInt(x, 10)
   763  	switch y {
   764  	case 0:
   765  		hdr.ObjAttrs.Size = (x & 0xffffff) + 1
   766  		hdr.ObjAttrs.Ver = s
   767  		hdr.ObjAttrs.SetCksum(cos.ChecksumNone, "")
   768  	case 1:
   769  		if usePDU {
   770  			hdr.ObjAttrs.Size = transport.SizeUnknown
   771  		} else {
   772  			hdr.ObjAttrs.Size = (x & 0xfffff) + 1
   773  		}
   774  		hdr.ObjAttrs.SetCustomKey(strconv.FormatInt(random.Int63(), 10), s)
   775  		hdr.ObjAttrs.SetCustomKey(s, "")
   776  		hdr.ObjAttrs.SetCksum(cos.ChecksumMD5, "md5")
   777  	case 2:
   778  		hdr.ObjAttrs.Size = (x & 0xffff) + 1
   779  		hdr.ObjAttrs.SetCksum(cos.ChecksumXXHash, "xxhash")
   780  		for range int(x & 0x1f) {
   781  			hdr.ObjAttrs.SetCustomKey(strconv.FormatInt(random.Int63(), 10), s)
   782  		}
   783  	default:
   784  		hdr.ObjAttrs.Size = 0
   785  		hdr.ObjAttrs.Ver = s
   786  		hdr.ObjAttrs.SetCustomKey(s, "")
   787  		hdr.ObjAttrs.SetCksum(cos.ChecksumNone, "")
   788  	}
   789  	return
   790  }
   791  
   792  ////////////////
   793  // randReader //
   794  ////////////////
   795  
   796  type randReader struct {
   797  	buf    []byte
   798  	hdr    transport.ObjHdr
   799  	slab   *memsys.Slab
   800  	off    int64
   801  	random *rand.Rand
   802  	offEOF int64 // when size is unknown
   803  	clone  bool
   804  }
   805  
   806  //nolint:gocritic // can do (hdr) hugeParam
   807  func newRandReader(random *rand.Rand, hdr transport.ObjHdr, slab *memsys.Slab) *randReader {
   808  	buf := slab.Alloc()
   809  	_, err := random.Read(buf)
   810  	if err != nil {
   811  		panic("Failed read rand: " + err.Error())
   812  	}
   813  	r := &randReader{buf: buf, hdr: hdr, slab: slab, random: random}
   814  	if hdr.IsUnsized() {
   815  		r.offEOF = int64(random.Int31()>>1) + 1
   816  	}
   817  	return r
   818  }
   819  
   820  func makeRandReader(random *rand.Rand, usePDU bool) (transport.ObjHdr, *randReader) {
   821  	hdr := genRandomHeader(random, usePDU)
   822  	if hdr.ObjSize() == 0 {
   823  		return hdr, nil
   824  	}
   825  	slab, err := memsys.PageMM().GetSlab(memsys.DefaultBufSize)
   826  	if err != nil {
   827  		panic("Failed getting slab: " + err.Error())
   828  	}
   829  	return hdr, newRandReader(random, hdr, slab)
   830  }
   831  
   832  func (r *randReader) Read(p []byte) (n int, err error) {
   833  	var objSize int64
   834  	if r.hdr.IsUnsized() {
   835  		objSize = r.offEOF
   836  	} else {
   837  		objSize = r.hdr.ObjAttrs.Size
   838  	}
   839  	for loff := 0; ; {
   840  		rem := objSize - r.off
   841  		if rem == 0 {
   842  			return n, io.EOF
   843  		}
   844  		l64 := min(rem, int64(len(p)-n))
   845  		if l64 == 0 {
   846  			return
   847  		}
   848  		nr := copy(p[n:n+int(l64)], r.buf[loff:])
   849  		n += nr
   850  		loff += 16
   851  		if loff >= len(r.buf)-16 {
   852  			loff = 0
   853  		}
   854  		r.off += int64(nr)
   855  	}
   856  }
   857  
   858  func (r *randReader) Open() (cos.ReadOpenCloser, error) {
   859  	buf := r.slab.Alloc()
   860  	copy(buf, r.buf)
   861  	r2 := randReader{buf: buf, hdr: r.hdr, slab: r.slab, offEOF: r.offEOF}
   862  	return &r2, nil
   863  }
   864  
   865  func (r *randReader) Close() error {
   866  	if r != nil && !r.clone {
   867  		r.slab.Free(r.buf)
   868  	}
   869  	return nil
   870  }
   871  
   872  type randReaderCtx struct {
   873  	t      *testing.T
   874  	rr     *randReader
   875  	posted []*randReader // => stream
   876  	mu     *sync.Mutex
   877  	idx    int
   878  }
   879  
   880  func (rrc *randReaderCtx) sentCallback(hdr *transport.ObjHdr, _ io.ReadCloser, _ any, err error) {
   881  	if err != nil {
   882  		rrc.t.Errorf("sent-callback %d(%s) returned an error: %v", rrc.idx, hdr.Cname(), err)
   883  	}
   884  	rr := rrc.rr
   885  	if rr != nil {
   886  		rr.slab.Free(rr.buf)
   887  	}
   888  	rrc.mu.Lock()
   889  	rrc.posted[rrc.idx] = nil
   890  	if rrc.idx > 0 && rrc.posted[rrc.idx-1] != nil {
   891  		rrc.t.Errorf("sent-callback %d(%s) fired out of order", rrc.idx, hdr.Cname())
   892  	}
   893  	rrc.posted[rrc.idx] = nil
   894  	rrc.mu.Unlock()
   895  }