fix: more reliable pss and gsoc (#462)

* fix: pss and gsoc

* chore: fix lint

* chore: add timeout

* fix: review comments

* fix: graceful shutdown

* fix: review comments

* fix: review comments

* fix: use 10m request timeout

Author: acha-bill

config/config.yaml

@@ -247,8 +247,8 @@ checks:
       postage-ttl: 24h
       postage-depth: 21
       postage-label: test-label
-      request-timeout: 5m
-    timeout: 5m
+      request-timeout: 10m
+    timeout: 30m
     type: pss
   pullsync:
     options:

config/local.yaml

@@ -223,8 +223,8 @@ checks:
       postage-ttl: 24h
       postage-depth: 21
       postage-label: test-label
-      request-timeout: 5m
-    timeout: 5m
+      request-timeout: 10m
+    timeout: 30m
     type: pss
   ci-pushsync-chunks:
     options:

config/public-testnet.yaml

@@ -93,7 +93,7 @@ checks:
       postage-ttl: 336h
       postage-depth: 21
       postage-label: test-label
-      request-timeout: 5m
+      request-timeout: 10m
     timeout: 30m
     type: pss
   pt-soc:

pkg/check/gsoc/gsoc.go

@@ -6,7 +6,6 @@ import (
 	"encoding/binary"
 	"encoding/hex"
 	"fmt"
-	"net/http"
 	"strconv"
 	"sync"
 	"time"
@@ -19,7 +18,7 @@ import (
 	"github.com/ethersphere/beekeeper/pkg/beekeeper"
 	"github.com/ethersphere/beekeeper/pkg/logging"
 	"github.com/ethersphere/beekeeper/pkg/orchestration"
-	"github.com/gorilla/websocket"
+	"github.com/ethersphere/beekeeper/pkg/wslistener"
 	"golang.org/x/sync/errgroup"
 )
 
@@ -122,22 +121,36 @@ func run(ctx context.Context, uploadClient *bee.Client, listenClient *bee.Client
 	if err != nil {
 		return err
 	}
-	logger.Infof("gsoc: socAddress=%s, listner node address=%s", socAddress, addresses.Overlay)
+	logger.Infof("gsoc: socAddress=%s, listener node address=%s, node=%s", socAddress, addresses.Overlay, listenClient.Name())
 
-	listener := &socListener{}
-	ch, err := listener.Listen(ctx, listenClient, socAddress, logger)
+	ctx, cancel := context.WithCancel(ctx)
+	ch, closer, err := wslistener.ListenWebSocket(ctx, listenClient, fmt.Sprintf("/gsoc/subscribe/%s", socAddress), logger)
 	if err != nil {
+		cancel()
 		return fmt.Errorf("listen websocket: %w", err)
 	}
-	defer listener.Close()
+
+	defer func() {
+		cancel()
+		closer()
+	}()
 
 	received := make(map[string]bool, numChunks)
 	receivedMtx := new(sync.Mutex)
+	done := make(chan struct{})
 
 	go func() {
 		for p := range ch {
 			receivedMtx.Lock()
+			if !received[p] {
+				logger.Infof("gsoc: received message %s on node %s", p, listenClient.Name())
+			}
 			received[p] = true
+			if len(received) == numChunks {
+				close(done)
+				receivedMtx.Unlock()
+				return
+			}
 			receivedMtx.Unlock()
 		}
 	}()
@@ -151,16 +164,14 @@ func run(ctx context.Context, uploadClient *bee.Client, listenClient *bee.Client
 		return err
 	}
 
-	// wait for listener to receive all messages
-	time.Sleep(5 * time.Second)
-	listener.Close()
+	select {
+	case <-time.After(1 * time.Minute):
+		return fmt.Errorf("timeout: not all messages received")
+	case <-done:
+	}
 
 	receivedMtx.Lock()
 	defer receivedMtx.Unlock()
-	if len(received) != numChunks {
-		return fmt.Errorf("expected %d messages, got %d", numChunks, len(received))
-	}
-
 	for i := 0; i < numChunks; i++ {
 		want := fmt.Sprintf("data %d", i)
 		if !received[want] {
@@ -289,58 +300,3 @@ func makeSoc(msg string, id []byte, privKey *ecdsa.PrivateKey) (*socData, error)
 		Data:  ch.Data(),
 	}, nil
 }
-
-type socListener struct {
-	listening bool
-	ws        *websocket.Conn
-	ch        chan string
-}
-
-func (s *socListener) Close() {
-	if !s.listening {
-		return
-	}
-
-	s.listening = false
-	s.ws.Close()
-	close(s.ch)
-}
-
-func (s *socListener) Listen(ctx context.Context, client *bee.Client, addr swarm.Address, logger logging.Logger) (<-chan string, error) {
-	dialer := &websocket.Dialer{
-		Proxy:            http.ProxyFromEnvironment,
-		HandshakeTimeout: 45 * time.Second,
-	}
-
-	ws, _, err := dialer.DialContext(ctx, fmt.Sprintf("ws://%s/gsoc/subscribe/%s", client.Host(), addr), http.Header{})
-	if err != nil {
-		return nil, err
-	}
-	s.ws = ws
-	s.ch = make(chan string)
-	s.listening = true
-
-	go func(nodeName string) {
-		for {
-			select {
-			case <-ctx.Done():
-				return
-			default:
-				msgType, data, err := ws.ReadMessage()
-				if err != nil {
-					logger.WithField("node", nodeName).Infof("gsoc: websocket error: %v", err)
-					return
-				}
-				if msgType != websocket.BinaryMessage {
-					logger.WithField("node", nodeName).Info("gsoc: websocket received non-binary message")
-					continue
-				}
-
-				logger.WithField("node", nodeName).Infof("gsoc: websocket received message: %s", string(data))
-				s.ch <- string(data)
-			}
-		}
-	}(client.Name())
-
-	return s.ch, nil
-}

pkg/check/pss/pss.go

@@ -2,18 +2,16 @@ package pss
 
 import (
 	"context"
-	"errors"
 	"fmt"
 	"math/rand"
-	"net/http"
 	"time"
 
 	"github.com/ethersphere/beekeeper/pkg/bee"
 	"github.com/ethersphere/beekeeper/pkg/beekeeper"
 	"github.com/ethersphere/beekeeper/pkg/logging"
 	"github.com/ethersphere/beekeeper/pkg/orchestration"
 	"github.com/ethersphere/beekeeper/pkg/random"
-	"github.com/gorilla/websocket"
+	"github.com/ethersphere/beekeeper/pkg/wslistener"
 )
 
 // Options represents check options
@@ -100,96 +98,60 @@ func pickAtRandom(r *rand.Rand, names []string, skip string) string {
 	}
 }
 
-var (
-	errDataMismatch        = errors.New("pss: data sent and received are not equal")
-	errWebsocketConnection = errors.New("pss: websocket connection terminated with an error")
-)
-
 var (
 	testData  = []byte("Hello Swarm :)")
 	testTopic = "test"
 )
 
 func (c *Check) testPss(nodeAName, nodeBName string, clients map[string]*bee.Client, o Options) error {
 	ctx, cancel := context.WithTimeout(context.Background(), o.RequestTimeout)
+	defer cancel()
 
 	nodeA := clients[nodeAName]
 	nodeB := clients[nodeBName]
 
 	addrB, err := nodeB.Addresses(ctx)
 	if err != nil {
-		cancel()
 		return err
 	}
 
 	batchID, err := nodeA.GetOrCreateMutableBatch(ctx, o.PostageTTL, o.PostageDepth, o.PostageLabel)
 	if err != nil {
-		cancel()
 		return fmt.Errorf("node %s: batched id %w", nodeAName, err)
 	}
 	c.logger.Infof("node %s: batched id %s", nodeAName, batchID)
 
-	ch, close, err := listenWebsocket(ctx, nodeB.Host(), testTopic, c.logger)
+	ch, closer, err := wslistener.ListenWebSocket(ctx, nodeB, fmt.Sprintf("/pss/subscribe/%s", testTopic), c.logger)
 	if err != nil {
-		cancel()
 		return err
 	}
 
 	c.logger.Infof("pss: sending test data to node %s and listening on node %s", nodeAName, nodeBName)
 
+	defer closer()
+
 	tStart := time.Now()
 	err = nodeA.SendPSSMessage(ctx, addrB.Overlay, addrB.PSSPublicKey, testTopic, o.AddressPrefix, testData, batchID)
 	if err != nil {
-		close()
-		cancel()
 		return err
 	}
+	c.logger.Infof("pss: test data sent successfully to node %s. Waiting for response from node %s", nodeAName, nodeBName)
 
-	msg, ok := <-ch
-	if ok {
-		if msg == string(testData) {
-			c.logger.Info("pss: websocket connection received correct message")
-			c.metrics.SendAndReceiveGauge.WithLabelValues(nodeAName, nodeBName).Set(time.Since(tStart).Seconds())
-		} else {
-			err = errDataMismatch
+	for {
+		select {
+		case <-time.After(1 * time.Minute):
+			return fmt.Errorf("correct message not received after %s", 1*time.Minute)
+		case msg, ok := <-ch:
+			if !ok {
+				return fmt.Errorf("ws closed before receiving correct message")
+			}
+
+			if msg == string(testData) {
+				c.logger.Info("pss: websocket connection received correct message")
+				c.metrics.SendAndReceiveGauge.WithLabelValues(nodeAName, nodeBName).Set(time.Since(tStart).Seconds())
+				return nil
+			}
+			c.logger.Infof("pss: received incorrect message. trying again. want %s, got %s", string(testData), msg)
 		}
-	} else {
-		err = errWebsocketConnection
-	}
-
-	cancel()
-	close()
-
-	if err != nil {
-		return err
 	}
-
-	return nil
-}
-
-func listenWebsocket(ctx context.Context, host string, topic string, logger logging.Logger) (<-chan string, func(), error) {
-	dialer := &websocket.Dialer{
-		Proxy:            http.ProxyFromEnvironment,
-		HandshakeTimeout: 45 * time.Second,
-	}
-
-	ws, _, err := dialer.DialContext(ctx, fmt.Sprintf("ws://%s/pss/subscribe/%s", host, topic), http.Header{})
-	if err != nil {
-		return nil, nil, err
-	}
-
-	ch := make(chan string)
-
-	go func() {
-		_, data, err := ws.ReadMessage()
-		if err != nil {
-			logger.Infof("pss: websocket error %v", err)
-			close(ch)
-			return
-		}
-
-		ch <- string(data)
-	}()
-
-	return ch, func() { ws.Close() }, nil
 }