lotus/lotus-soup/deals_stress.go

package main

import (
	"context"
	"io/ioutil"
	"math/rand"
	"os"
	"sync"
	"time"

	"github.com/filecoin-project/lotus/api"
	"github.com/ipfs/go-cid"

	"github.com/filecoin-project/oni/lotus-soup/testkit"
)

func dealStressTest(t *testkit.TestEnvironment) error {
	// Dispatch/forward non-client roles to defaults.
	if t.Role != "client" {
		return testkit.HandleDefaultRole(t)
	}

	t.RecordMessage("running client")

	cl, err := testkit.PrepareClient(t)
	if err != nil {
		return err
	}

	ctx := context.Background()
	addrs, err := testkit.CollectMinerAddrs(t, ctx, t.IntParam("miners"))
	if err != nil {
		return err
	}
	t.RecordMessage("got %v miner addrs", len(addrs))

	client := cl.FullApi

	// select a random miner
	minerAddr := addrs[rand.Intn(len(addrs))]
	if err := client.NetConnect(ctx, minerAddr.PeerAddr); err != nil {
		return err
	}

	t.RecordMessage("selected %s as the miner", minerAddr.ActorAddr)

	time.Sleep(2 * time.Second)

	// prepare a number of concurrent data points
	deals := t.IntParam("deals")
	data := make([][]byte, 0, deals)
	files := make([]*os.File, 0, deals)
	cids := make([]cid.Cid, 0, deals)
	rng := rand.NewSource(time.Now().UnixNano())

	for i := 0; i < deals; i++ {
		dealData := make([]byte, 1600)
		rand.New(rng).Read(dealData)

		dealFile, err := ioutil.TempFile("/tmp", "data")
		if err != nil {
			return err
		}
		defer os.Remove(dealFile.Name())

		_, err = dealFile.Write(dealData)
		if err != nil {
			return err
		}

		dealCid, err := client.ClientImport(ctx, api.FileRef{Path: dealFile.Name(), IsCAR: false})
		if err != nil {
			return err
		}

		t.RecordMessage("deal %d file cid: %s", i, dealCid)

		data = append(data, dealData)
		files = append(files, dealFile)
		cids = append(cids, dealCid)
	}

	concurrentDeals := true
	if t.StringParam("deal_mode") == "serial" {
		concurrentDeals = false
	}

	t.RecordMessage("starting storage deals")
	if concurrentDeals {

		var wg1 sync.WaitGroup
		for i := 0; i < deals; i++ {
			wg1.Add(1)
			go func(i int) {
				defer wg1.Done()
				deal := testkit.StartDeal(ctx, minerAddr.ActorAddr, client, cids[i])
				t.RecordMessage("started storage deal %d -> %s", i, deal)
				time.Sleep(2 * time.Second)
				t.RecordMessage("waiting for deal %d to be sealed", i)
				testkit.WaitDealSealed(t, ctx, client, deal)
			}(i)
		}
		t.RecordMessage("waiting for all deals to be sealed")
		wg1.Wait()
		t.RecordMessage("all deals sealed; starting retrieval")

		var wg2 sync.WaitGroup
		for i := 0; i < deals; i++ {
			wg2.Add(1)
			go func(i int) {
				defer wg2.Done()
				t.RecordMessage("retrieving data for deal %d", i)
				testkit.RetrieveData(t, ctx, client, cids[i], true, data[i])
				t.RecordMessage("retrieved data for deal %d", i)
			}(i)
		}
		t.RecordMessage("waiting for all retrieval deals to complete")
		wg2.Wait()
		t.RecordMessage("all retrieval deals successful")

	} else {

		for i := 0; i < deals; i++ {
			deal := testkit.StartDeal(ctx, minerAddr.ActorAddr, client, cids[i])
			t.RecordMessage("started storage deal %d -> %s", i, deal)
			time.Sleep(2 * time.Second)
			t.RecordMessage("waiting for deal %d to be sealed", i)
			testkit.WaitDealSealed(t, ctx, client, deal)
		}

		for i := 0; i < deals; i++ {
			t.RecordMessage("retrieving data for deal %d", i)
			testkit.RetrieveData(t, ctx, client, cids[i], true, data[i])
			t.RecordMessage("retrieved data for deal %d", i)
		}
	}

	t.SyncClient.MustSignalEntry(ctx, testkit.StateStopMining)
	t.SyncClient.MustSignalAndWait(ctx, testkit.StateDone, t.TestInstanceCount)

	return nil
}