github.com/aaa256/atlantis@v0.0.0-20210707112435-42ee889287a2/swarm/network/stream/intervals_test.go (about)

     1  // Copyright 2018 The go-athereum Authors
     2  // This file is part of the go-athereum library.
     3  //
     4  // The go-athereum library is free software: you can redistribute it and/or modify
     5  // it under the terms of the GNU Lesser General Public License as published by
     6  // the Free Software Foundation, either version 3 of the License, or
     7  // (at your option) any later version.
     8  //
     9  // The go-athereum library is distributed in the hope that it will be useful,
    10  // but WITHOUT ANY WARRANTY; without even the implied warranty of
    11  // MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
    12  // GNU Lesser General Public License for more details.
    13  //
    14  // You should have received a copy of the GNU Lesser General Public License
    15  // along with the go-athereum library. If not, see <http://www.gnu.org/licenses/>.
    16  
    17  package stream
    18  
    19  import (
    20  	"context"
    21  	crand "crypto/rand"
    22  	"encoding/binary"
    23  	"fmt"
    24  	"io"
    25  	"sync"
    26  	"testing"
    27  	"time"
    28  
    29  	"github.com/athereum/go-athereum/node"
    30  	"github.com/athereum/go-athereum/p2p/discover"
    31  	"github.com/athereum/go-athereum/p2p/simulations"
    32  	"github.com/athereum/go-athereum/p2p/simulations/adapters"
    33  	"github.com/athereum/go-athereum/rpc"
    34  	"github.com/athereum/go-athereum/swarm/network"
    35  	streamTesting "github.com/athereum/go-athereum/swarm/network/stream/testing"
    36  	"github.com/athereum/go-athereum/swarm/state"
    37  	"github.com/athereum/go-athereum/swarm/storage"
    38  )
    39  
    40  var (
    41  	externalStreamName             = "externalStream"
    42  	externalStreamSessionAt uint64 = 50
    43  	externalStreamMaxKeys   uint64 = 100
    44  )
    45  
    46  func newIntervalsStreamerService(ctx *adapters.ServiceContext) (node.Service, error) {
    47  	id := ctx.Config.ID
    48  	addr := toAddr(id)
    49  	kad := network.NewKademlia(addr.Over(), network.NewKadParams())
    50  	store := stores[id].(*storage.LocalStore)
    51  	db := storage.NewDBAPI(store)
    52  	delivery := NewDelivery(kad, db)
    53  	deliveries[id] = delivery
    54  	r := NewRegistry(addr, delivery, db, state.NewInmemoryStore(), &RegistryOptions{
    55  		SkipCheck: defaultSkipCheck,
    56  	})
    57  
    58  	r.RegisterClientFunc(externalStreamName, func(p *Peer, t string, live bool) (Client, error) {
    59  		return newTestExternalClient(db), nil
    60  	})
    61  	r.RegisterServerFunc(externalStreamName, func(p *Peer, t string, live bool) (Server, error) {
    62  		return newTestExternalServer(t, externalStreamSessionAt, externalStreamMaxKeys, nil), nil
    63  	})
    64  
    65  	go func() {
    66  		waitPeerErrC <- waitForPeers(r, 1*time.Second, peerCount(id))
    67  	}()
    68  	return &TestExternalRegistry{r}, nil
    69  }
    70  
    71  func TestIntervals(t *testing.T) {
    72  	testIntervals(t, true, nil, false)
    73  	testIntervals(t, false, NewRange(9, 26), false)
    74  	testIntervals(t, true, NewRange(9, 26), false)
    75  
    76  	testIntervals(t, true, nil, true)
    77  	testIntervals(t, false, NewRange(9, 26), true)
    78  	testIntervals(t, true, NewRange(9, 26), true)
    79  }
    80  
    81  func testIntervals(t *testing.T, live bool, history *Range, skipCheck bool) {
    82  	nodes := 2
    83  	chunkCount := dataChunkCount
    84  
    85  	defer setDefaultSkipCheck(defaultSkipCheck)
    86  	defaultSkipCheck = skipCheck
    87  
    88  	toAddr = network.NewAddrFromNodeID
    89  	conf := &streamTesting.RunConfig{
    90  		Adapter:        *adapter,
    91  		NodeCount:      nodes,
    92  		ConnLevel:      1,
    93  		ToAddr:         toAddr,
    94  		Services:       services,
    95  		DefaultService: "intervalsStreamer",
    96  	}
    97  
    98  	sim, teardown, err := streamTesting.NewSimulation(conf)
    99  	var rpcSubscriptionsWg sync.WaitGroup
   100  	defer func() {
   101  		rpcSubscriptionsWg.Wait()
   102  		teardown()
   103  	}()
   104  	if err != nil {
   105  		t.Fatal(err)
   106  	}
   107  
   108  	stores = make(map[discover.NodeID]storage.ChunkStore)
   109  	deliveries = make(map[discover.NodeID]*Delivery)
   110  	for i, id := range sim.IDs {
   111  		stores[id] = sim.Stores[i]
   112  	}
   113  
   114  	peerCount = func(id discover.NodeID) int {
   115  		return 1
   116  	}
   117  
   118  	fileStore := storage.NewFileStore(sim.Stores[0], storage.NewFileStoreParams())
   119  	size := chunkCount * chunkSize
   120  	_, wait, err := fileStore.Store(io.LimitReader(crand.Reader, int64(size)), int64(size), false)
   121  	wait()
   122  	if err != nil {
   123  		t.Fatal(err)
   124  	}
   125  
   126  	errc := make(chan error, 1)
   127  	waitPeerErrC = make(chan error)
   128  	quitC := make(chan struct{})
   129  	defer close(quitC)
   130  
   131  	action := func(ctx context.Context) error {
   132  		i := 0
   133  		for err := range waitPeerErrC {
   134  			if err != nil {
   135  				return fmt.Errorf("error waiting for peers: %s", err)
   136  			}
   137  			i++
   138  			if i == nodes {
   139  				break
   140  			}
   141  		}
   142  
   143  		id := sim.IDs[1]
   144  
   145  		err := sim.CallClient(id, func(client *rpc.Client) error {
   146  
   147  			sid := sim.IDs[0]
   148  
   149  			doneC, err := streamTesting.WatchDisconnections(id, client, errc, quitC)
   150  			if err != nil {
   151  				return err
   152  			}
   153  			rpcSubscriptionsWg.Add(1)
   154  			go func() {
   155  				<-doneC
   156  				rpcSubscriptionsWg.Done()
   157  			}()
   158  			ctx, cancel := context.WithTimeout(ctx, 100*time.Second)
   159  			defer cancel()
   160  
   161  			err = client.CallContext(ctx, nil, "stream_subscribeStream", sid, NewStream(externalStreamName, "", live), history, Top)
   162  			if err != nil {
   163  				return err
   164  			}
   165  
   166  			liveErrC := make(chan error)
   167  			historyErrC := make(chan error)
   168  
   169  			go func() {
   170  				if !live {
   171  					close(liveErrC)
   172  					return
   173  				}
   174  
   175  				var err error
   176  				defer func() {
   177  					liveErrC <- err
   178  				}()
   179  
   180  				// live stream
   181  				liveHashesChan := make(chan []byte)
   182  				liveSubscription, err := client.Subscribe(ctx, "stream", liveHashesChan, "getHashes", sid, NewStream(externalStreamName, "", true))
   183  				if err != nil {
   184  					return
   185  				}
   186  				defer liveSubscription.Unsubscribe()
   187  
   188  				i := externalStreamSessionAt
   189  
   190  				// we have subscribed, enable notifications
   191  				err = client.CallContext(ctx, nil, "stream_enableNotifications", sid, NewStream(externalStreamName, "", true))
   192  				if err != nil {
   193  					return
   194  				}
   195  
   196  				for {
   197  					select {
   198  					case hash := <-liveHashesChan:
   199  						h := binary.BigEndian.Uint64(hash)
   200  						if h != i {
   201  							err = fmt.Errorf("expected live hash %d, got %d", i, h)
   202  							return
   203  						}
   204  						i++
   205  						if i > externalStreamMaxKeys {
   206  							return
   207  						}
   208  					case err = <-liveSubscription.Err():
   209  						return
   210  					case <-ctx.Done():
   211  						return
   212  					}
   213  				}
   214  			}()
   215  
   216  			go func() {
   217  				if live && history == nil {
   218  					close(historyErrC)
   219  					return
   220  				}
   221  
   222  				var err error
   223  				defer func() {
   224  					historyErrC <- err
   225  				}()
   226  
   227  				// history stream
   228  				historyHashesChan := make(chan []byte)
   229  				historySubscription, err := client.Subscribe(ctx, "stream", historyHashesChan, "getHashes", sid, NewStream(externalStreamName, "", false))
   230  				if err != nil {
   231  					return
   232  				}
   233  				defer historySubscription.Unsubscribe()
   234  
   235  				var i uint64
   236  				historyTo := externalStreamMaxKeys
   237  				if history != nil {
   238  					i = history.From
   239  					if history.To != 0 {
   240  						historyTo = history.To
   241  					}
   242  				}
   243  
   244  				// we have subscribed, enable notifications
   245  				err = client.CallContext(ctx, nil, "stream_enableNotifications", sid, NewStream(externalStreamName, "", false))
   246  				if err != nil {
   247  					return
   248  				}
   249  
   250  				for {
   251  					select {
   252  					case hash := <-historyHashesChan:
   253  						h := binary.BigEndian.Uint64(hash)
   254  						if h != i {
   255  							err = fmt.Errorf("expected history hash %d, got %d", i, h)
   256  							return
   257  						}
   258  						i++
   259  						if i > historyTo {
   260  							return
   261  						}
   262  					case err = <-historySubscription.Err():
   263  						return
   264  					case <-ctx.Done():
   265  						return
   266  					}
   267  				}
   268  			}()
   269  
   270  			if err := <-liveErrC; err != nil {
   271  				return err
   272  			}
   273  			if err := <-historyErrC; err != nil {
   274  				return err
   275  			}
   276  
   277  			return nil
   278  		})
   279  		return err
   280  	}
   281  	check := func(ctx context.Context, id discover.NodeID) (bool, error) {
   282  		select {
   283  		case err := <-errc:
   284  			return false, err
   285  		case <-ctx.Done():
   286  			return false, ctx.Err()
   287  		default:
   288  		}
   289  		return true, nil
   290  	}
   291  
   292  	conf.Step = &simulations.Step{
   293  		Action:  action,
   294  		Trigger: streamTesting.Trigger(10*time.Millisecond, quitC, sim.IDs[0]),
   295  		Expect: &simulations.Expectation{
   296  			Nodes: sim.IDs[1:1],
   297  			Check: check,
   298  		},
   299  	}
   300  	startedAt := time.Now()
   301  	timeout := 300 * time.Second
   302  	ctx, cancel := context.WithTimeout(context.Background(), timeout)
   303  	defer cancel()
   304  	result, err := sim.Run(ctx, conf)
   305  	finishedAt := time.Now()
   306  	if err != nil {
   307  		t.Fatalf("Setting up simulation failed: %v", err)
   308  	}
   309  	if result.Error != nil {
   310  		t.Fatalf("Simulation failed: %s", result.Error)
   311  	}
   312  	streamTesting.CheckResult(t, result, startedAt, finishedAt)
   313  }