extend p2p/protobuf reader and writer with cancelation methods

pkg/p2p/protobuf/protobuf.go

@@ -5,7 +5,10 @@
 package protobuf
 
 import (
+	"context"
+	"errors"
 	"io"
+	"time"
 
 	"github.com/ethersphere/bee/pkg/p2p"
 	ggio "github.com/gogo/protobuf/io"
@@ -14,20 +17,22 @@ import (
 
 const delimitedReaderMaxSize = 128 * 1024 // max message size
 
+var ErrTimeout = errors.New("timeout")
+
 type Message = proto.Message
 
-func NewWriterAndReader(s p2p.Stream) (w ggio.Writer, r ggio.Reader) {
-	r = ggio.NewDelimitedReader(s, delimitedReaderMaxSize)
-	w = ggio.NewDelimitedWriter(s)
-	return w, r
+func NewWriterAndReader(s p2p.Stream) (Writer, Reader) {
+	w := ggio.NewDelimitedWriter(s)
+	r := ggio.NewDelimitedReader(s, delimitedReaderMaxSize)
+	return newWriter(w), newReader(r)
 }
 
-func NewReader(r io.Reader) ggio.Reader {
-	return ggio.NewDelimitedReader(r, delimitedReaderMaxSize)
+func NewReader(r io.Reader) Reader {
+	return newReader(ggio.NewDelimitedReader(r, delimitedReaderMaxSize))
 }
 
 func NewWriter(w io.Writer) ggio.Writer {
-	return ggio.NewDelimitedWriter(w)
+	return newWriter(ggio.NewDelimitedWriter(w))
 }
 
 func ReadMessages(r io.Reader, newMessage func() Message) (m []Message, err error) {
@@ -44,3 +49,81 @@ func ReadMessages(r io.Reader, newMessage func() Message) (m []Message, err erro
 	}
 	return m, nil
 }
+
+type Reader struct {
+	ggio.Reader
+}
+
+func newReader(r ggio.Reader) Reader {
+	return Reader{Reader: r}
+}
+
+func (r Reader) ReadMsgWithContext(ctx context.Context, msg proto.Message) error {
+	errChan := make(chan error, 1)
+	go func() {
+		errChan <- r.ReadMsg(msg)
+	}()
+
+	select {
+	case err := <-errChan:
+		return err
+	case <-ctx.Done():
+		return ctx.Err()
+	}
+}
+
+func (r Reader) ReadMsgWithTimeout(d time.Duration, msg proto.Message) error {
+	errChan := make(chan error, 1)
+	go func() {
+		errChan <- r.ReadMsg(msg)
+	}()
+
+	timer := time.NewTimer(d)
+	defer timer.Stop()
+
+	select {
+	case err := <-errChan:
+		return err
+	case <-timer.C:
+		return ErrTimeout
+	}
+}
+
+type Writer struct {
+	ggio.Writer
+}
+
+func newWriter(r ggio.Writer) Writer {
+	return Writer{Writer: r}
+}
+
+func (w Writer) WriteMsgWithContext(ctx context.Context, msg proto.Message) error {
+	errChan := make(chan error, 1)
+	go func() {
+		errChan <- w.WriteMsg(msg)
+	}()
+
+	select {
+	case err := <-errChan:
+		return err
+	case <-ctx.Done():
+		return ctx.Err()
+	}
+}
+
+func (w Writer) WriteMsgWithTimeout(d time.Duration, msg proto.Message) error {
+	errChan := make(chan error, 1)
+	go func() {
+		errChan <- w.WriteMsg(msg)
+	}()
+
+	timer := time.NewTimer(d)
+	defer timer.Stop()
+
+	select {
+	case err := <-errChan:
+		return err
+	case <-timer.C:
+		return ErrTimeout
+	}
+}