github.com/zignig/go-ipfs@v0.0.0-20141111235910-c9e5fdf55a52/net/conn/multiconn_test.go (about)

     1  package conn
     2  
     3  import (
     4  	"fmt"
     5  	"sync"
     6  	"testing"
     7  	"time"
     8  
     9  	peer "github.com/jbenet/go-ipfs/peer"
    10  
    11  	context "github.com/jbenet/go-ipfs/Godeps/_workspace/src/code.google.com/p/go.net/context"
    12  	ma "github.com/jbenet/go-ipfs/Godeps/_workspace/src/github.com/jbenet/go-multiaddr"
    13  )
    14  
    15  func tcpAddr(t *testing.T, port int) ma.Multiaddr {
    16  	tcp, err := ma.NewMultiaddr(tcpAddrString(port))
    17  	if err != nil {
    18  		t.Fatal(err)
    19  	}
    20  	return tcp
    21  }
    22  
    23  func tcpAddrString(port int) string {
    24  	return fmt.Sprintf("/ip4/127.0.0.1/tcp/%d", port)
    25  }
    26  
    27  type msg struct {
    28  	sent     bool
    29  	received bool
    30  	payload  string
    31  }
    32  
    33  func (m *msg) Sent(t *testing.T) {
    34  	if m.sent {
    35  		t.Fatal("sent msg at incorrect state:", m)
    36  	}
    37  	m.sent = true
    38  }
    39  
    40  func (m *msg) Received(t *testing.T) {
    41  	if m.received {
    42  		t.Fatal("received msg at incorrect state:", m)
    43  	}
    44  	m.received = true
    45  }
    46  
    47  type msgMap struct {
    48  	sent int
    49  	recv int
    50  	msgs map[string]*msg
    51  }
    52  
    53  func (mm *msgMap) Sent(t *testing.T, payload string) {
    54  	mm.msgs[payload].Sent(t)
    55  	mm.sent++
    56  }
    57  
    58  func (mm *msgMap) Received(t *testing.T, payload string) {
    59  	mm.msgs[payload].Received(t)
    60  	mm.recv++
    61  }
    62  
    63  func (mm *msgMap) CheckDone(t *testing.T) {
    64  	if mm.sent != len(mm.msgs) {
    65  		t.Fatal("failed to send all msgs", mm.sent, len(mm.msgs))
    66  	}
    67  
    68  	if mm.sent != len(mm.msgs) {
    69  		t.Fatal("failed to send all msgs", mm.sent, len(mm.msgs))
    70  	}
    71  }
    72  
    73  func genMessages(num int, tag string) *msgMap {
    74  	msgs := &msgMap{msgs: map[string]*msg{}}
    75  	for i := 0; i < num; i++ {
    76  		s := fmt.Sprintf("Message #%d -- %s", i, tag)
    77  		msgs.msgs[s] = &msg{payload: s}
    78  	}
    79  	return msgs
    80  }
    81  
    82  func setupMultiConns(t *testing.T, ctx context.Context) (a, b *MultiConn) {
    83  
    84  	log.Info("Setting up peers")
    85  	p1, err := setupPeer(tcpAddrString(11000))
    86  	if err != nil {
    87  		t.Fatal("error setting up peer", err)
    88  	}
    89  
    90  	p2, err := setupPeer(tcpAddrString(12000))
    91  	if err != nil {
    92  		t.Fatal("error setting up peer", err)
    93  	}
    94  
    95  	// peerstores
    96  	p1ps := peer.NewPeerstore()
    97  	p2ps := peer.NewPeerstore()
    98  	p1ps.Add(p1)
    99  	p2ps.Add(p2)
   100  
   101  	// listeners
   102  	listen := func(addr ma.Multiaddr, p peer.Peer, ps peer.Peerstore) Listener {
   103  		l, err := Listen(ctx, addr, p, ps)
   104  		if err != nil {
   105  			t.Fatal(err)
   106  		}
   107  		return l
   108  	}
   109  
   110  	log.Info("Setting up listeners")
   111  	p1l := listen(p1.Addresses()[0], p1, p1ps)
   112  	p2l := listen(p2.Addresses()[0], p2, p2ps)
   113  
   114  	// dialers
   115  	p1d := &Dialer{Peerstore: p1ps, LocalPeer: p1}
   116  	p2d := &Dialer{Peerstore: p2ps, LocalPeer: p2}
   117  
   118  	dial := func(d *Dialer, dst peer.Peer) <-chan Conn {
   119  		cc := make(chan Conn)
   120  		go func() {
   121  			c, err := d.Dial(ctx, "tcp", dst)
   122  			if err != nil {
   123  				t.Fatal("error dialing peer", err)
   124  			}
   125  			cc <- c
   126  		}()
   127  		return cc
   128  	}
   129  
   130  	// connect simultaneously
   131  	log.Info("Connecting...")
   132  	p1dc := dial(p1d, p2)
   133  	p2dc := dial(p2d, p1)
   134  
   135  	c12a := <-p1l.Accept()
   136  	c12b := <-p1dc
   137  	c21a := <-p2l.Accept()
   138  	c21b := <-p2dc
   139  
   140  	log.Info("Ok, making multiconns")
   141  	c1, err := NewMultiConn(ctx, p1, p2, []Conn{c12a, c12b})
   142  	if err != nil {
   143  		t.Fatal(err)
   144  	}
   145  
   146  	c2, err := NewMultiConn(ctx, p2, p1, []Conn{c21a, c21b})
   147  	if err != nil {
   148  		t.Fatal(err)
   149  	}
   150  
   151  	p1l.Close()
   152  	p2l.Close()
   153  
   154  	log.Info("did you make multiconns?")
   155  	return c1, c2
   156  }
   157  
   158  func TestMulticonnSend(t *testing.T) {
   159  	// t.Skip("fooo")
   160  
   161  	log.Info("TestMulticonnSend")
   162  	ctx := context.Background()
   163  	ctxC, cancel := context.WithCancel(ctx)
   164  
   165  	c1, c2 := setupMultiConns(t, ctx)
   166  
   167  	log.Info("gen msgs")
   168  	num := 100
   169  	msgsFrom1 := genMessages(num, "from p1 to p2")
   170  	msgsFrom2 := genMessages(num, "from p2 to p1")
   171  
   172  	var wg sync.WaitGroup
   173  
   174  	send := func(c *MultiConn, msgs *msgMap) {
   175  		defer wg.Done()
   176  
   177  		for _, m := range msgs.msgs {
   178  			log.Info("send: %s", m.payload)
   179  			c.Out() <- []byte(m.payload)
   180  			msgs.Sent(t, m.payload)
   181  			<-time.After(time.Microsecond * 10)
   182  		}
   183  	}
   184  
   185  	recv := func(ctx context.Context, c *MultiConn, msgs *msgMap) {
   186  		defer wg.Done()
   187  
   188  		for {
   189  			select {
   190  			case payload := <-c.In():
   191  				msgs.Received(t, string(payload))
   192  				log.Info("recv: %s", payload)
   193  				if msgs.recv == len(msgs.msgs) {
   194  					return
   195  				}
   196  
   197  			case <-ctx.Done():
   198  				return
   199  
   200  			}
   201  		}
   202  
   203  	}
   204  
   205  	log.Info("msg send + recv")
   206  
   207  	wg.Add(4)
   208  	go send(c1, msgsFrom1)
   209  	go send(c2, msgsFrom2)
   210  	go recv(ctxC, c1, msgsFrom2)
   211  	go recv(ctxC, c2, msgsFrom1)
   212  	wg.Wait()
   213  	cancel()
   214  	c1.Close()
   215  	c2.Close()
   216  
   217  	msgsFrom1.CheckDone(t)
   218  	msgsFrom2.CheckDone(t)
   219  	<-time.After(100 * time.Millisecond)
   220  }
   221  
   222  func TestMulticonnSendUnderlying(t *testing.T) {
   223  	// t.Skip("fooo")
   224  
   225  	log.Info("TestMulticonnSendUnderlying")
   226  	ctx := context.Background()
   227  	ctxC, cancel := context.WithCancel(ctx)
   228  
   229  	c1, c2 := setupMultiConns(t, ctx)
   230  
   231  	log.Info("gen msgs")
   232  	num := 100
   233  	msgsFrom1 := genMessages(num, "from p1 to p2")
   234  	msgsFrom2 := genMessages(num, "from p2 to p1")
   235  
   236  	var wg sync.WaitGroup
   237  
   238  	send := func(c *MultiConn, msgs *msgMap) {
   239  		defer wg.Done()
   240  
   241  		conns := make([]Conn, 0, len(c.conns))
   242  		for _, c1 := range c.conns {
   243  			conns = append(conns, c1)
   244  		}
   245  
   246  		i := 0
   247  		for _, m := range msgs.msgs {
   248  			log.Info("send: %s", m.payload)
   249  			switch i % 3 {
   250  			case 0:
   251  				conns[0].Out() <- []byte(m.payload)
   252  			case 1:
   253  				conns[1].Out() <- []byte(m.payload)
   254  			case 2:
   255  				c.Out() <- []byte(m.payload)
   256  			}
   257  			msgs.Sent(t, m.payload)
   258  			<-time.After(time.Microsecond * 10)
   259  			i++
   260  		}
   261  	}
   262  
   263  	recv := func(ctx context.Context, c *MultiConn, msgs *msgMap) {
   264  		defer wg.Done()
   265  
   266  		for {
   267  			select {
   268  			case payload := <-c.In():
   269  				msgs.Received(t, string(payload))
   270  				log.Info("recv: %s", payload)
   271  				if msgs.recv == len(msgs.msgs) {
   272  					return
   273  				}
   274  
   275  			case <-ctx.Done():
   276  				return
   277  
   278  			}
   279  		}
   280  
   281  	}
   282  
   283  	log.Info("msg send + recv")
   284  
   285  	wg.Add(4)
   286  	go send(c1, msgsFrom1)
   287  	go send(c2, msgsFrom2)
   288  	go recv(ctxC, c1, msgsFrom2)
   289  	go recv(ctxC, c2, msgsFrom1)
   290  	wg.Wait()
   291  	cancel()
   292  	c1.Close()
   293  	c2.Close()
   294  
   295  	msgsFrom1.CheckDone(t)
   296  	msgsFrom2.CheckDone(t)
   297  }
   298  
   299  func TestMulticonnClose(t *testing.T) {
   300  	// t.Skip("fooo")
   301  
   302  	log.Info("TestMulticonnSendUnderlying")
   303  	ctx := context.Background()
   304  	c1, c2 := setupMultiConns(t, ctx)
   305  
   306  	for _, c := range c1.conns {
   307  		c.Close()
   308  	}
   309  
   310  	for _, c := range c2.conns {
   311  		c.Close()
   312  	}
   313  
   314  	timeout := time.After(100 * time.Millisecond)
   315  	select {
   316  	case <-c1.Closed():
   317  	case <-timeout:
   318  		t.Fatal("timeout")
   319  	}
   320  
   321  	select {
   322  	case <-c2.Closed():
   323  	case <-timeout:
   324  		t.Fatal("timeout")
   325  	}
   326  }