ceremonyclient/go-libp2p/p2p/test/resource-manager/echo.go
2023-08-20 23:07:43 -05:00

294 lines
4.9 KiB
Go

package itest
import (
"context"
"fmt"
"io"
"sync"
"time"
"github.com/libp2p/go-libp2p/core/host"
"github.com/libp2p/go-libp2p/core/network"
"github.com/libp2p/go-libp2p/core/peer"
logging "github.com/ipfs/go-log/v2"
)
const (
EchoService = "test.echo"
EchoProtoID = "/test/echo"
)
var (
echoLog = logging.Logger("echo")
)
type Echo struct {
Host host.Host
mx sync.Mutex
status EchoStatus
beforeReserve, beforeRead, beforeWrite, beforeDone func() error
done func()
}
type EchoStatus struct {
StreamsIn int
EchosIn, EchosOut int
IOErrors int
ResourceServiceErrors int
ResourceReservationErrors int
}
func NewEcho(h host.Host) *Echo {
e := &Echo{Host: h}
h.SetStreamHandler(EchoProtoID, e.handleStream)
return e
}
func (e *Echo) Status() EchoStatus {
e.mx.Lock()
defer e.mx.Unlock()
return e.status
}
func (e *Echo) BeforeReserve(f func() error) {
e.mx.Lock()
defer e.mx.Unlock()
e.beforeReserve = f
}
func (e *Echo) BeforeRead(f func() error) {
e.mx.Lock()
defer e.mx.Unlock()
e.beforeRead = f
}
func (e *Echo) BeforeWrite(f func() error) {
e.mx.Lock()
defer e.mx.Unlock()
e.beforeWrite = f
}
func (e *Echo) BeforeDone(f func() error) {
e.mx.Lock()
defer e.mx.Unlock()
e.beforeDone = f
}
func (e *Echo) Done(f func()) {
e.mx.Lock()
defer e.mx.Unlock()
e.done = f
}
func (e *Echo) getBeforeReserve() func() error {
e.mx.Lock()
defer e.mx.Unlock()
return e.beforeReserve
}
func (e *Echo) getBeforeRead() func() error {
e.mx.Lock()
defer e.mx.Unlock()
return e.beforeRead
}
func (e *Echo) getBeforeWrite() func() error {
e.mx.Lock()
defer e.mx.Unlock()
return e.beforeWrite
}
func (e *Echo) getBeforeDone() func() error {
e.mx.Lock()
defer e.mx.Unlock()
return e.beforeDone
}
func (e *Echo) getDone() func() {
e.mx.Lock()
defer e.mx.Unlock()
return e.done
}
func (e *Echo) handleStream(s network.Stream) {
defer s.Close()
if done := e.getDone(); done != nil {
defer done()
}
e.mx.Lock()
e.status.StreamsIn++
e.mx.Unlock()
if beforeReserve := e.getBeforeReserve(); beforeReserve != nil {
if err := beforeReserve(); err != nil {
echoLog.Debugf("error syncing before reserve: %s", err)
s.Reset()
return
}
}
if err := s.Scope().SetService(EchoService); err != nil {
echoLog.Debugf("error attaching stream to echo service: %s", err)
e.mx.Lock()
e.status.ResourceServiceErrors++
e.mx.Unlock()
s.Reset()
return
}
if err := s.Scope().ReserveMemory(4096, network.ReservationPriorityAlways); err != nil {
echoLog.Debugf("error reserving memory: %s", err)
e.mx.Lock()
e.status.ResourceReservationErrors++
e.mx.Unlock()
s.Reset()
return
}
if beforeRead := e.getBeforeRead(); beforeRead != nil {
if err := beforeRead(); err != nil {
echoLog.Debugf("error syncing before read: %s", err)
s.Reset()
return
}
}
buf := make([]byte, 4096)
s.SetReadDeadline(time.Now().Add(5 * time.Second))
n, err := s.Read(buf)
switch {
case err == io.EOF:
if n == 0 {
return
}
case err != nil:
echoLog.Debugf("I/O error : %s", err)
e.mx.Lock()
e.status.IOErrors++
e.mx.Unlock()
s.Reset()
return
}
e.mx.Lock()
e.status.EchosIn++
e.mx.Unlock()
if beforeWrite := e.getBeforeWrite(); beforeWrite != nil {
if err := beforeWrite(); err != nil {
echoLog.Debugf("error syncing before write: %s", err)
s.Reset()
return
}
}
s.SetWriteDeadline(time.Now().Add(5 * time.Second))
_, err = s.Write(buf[:n])
if err != nil {
echoLog.Debugf("I/O error: %s", err)
e.mx.Lock()
e.status.IOErrors++
e.mx.Unlock()
s.Reset()
return
}
e.mx.Lock()
e.status.EchosOut++
e.mx.Unlock()
s.CloseWrite()
if beforeDone := e.getBeforeDone(); beforeDone != nil {
if err := beforeDone(); err != nil {
echoLog.Debugf("error syncing before done: %s", err)
s.Reset()
}
}
}
func (e *Echo) Echo(p peer.ID, what string) error {
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
defer cancel()
s, err := e.Host.NewStream(ctx, p, EchoProtoID)
if err != nil {
return err
}
defer s.Close()
if err := s.Scope().SetService(EchoService); err != nil {
echoLog.Debugf("error attaching stream to echo service: %s", err)
s.Reset()
return err
}
if err := s.Scope().ReserveMemory(4096, network.ReservationPriorityAlways); err != nil {
echoLog.Debugf("error reserving memory: %s", err)
s.Reset()
return err
}
s.SetWriteDeadline(time.Now().Add(5 * time.Second))
_, err = s.Write([]byte(what))
if err != nil {
return err
}
s.CloseWrite()
buf := make([]byte, 4096)
s.SetReadDeadline(time.Now().Add(5 * time.Second))
n, err := s.Read(buf)
switch {
case err == io.EOF:
if n == 0 {
return err
}
case err != nil:
echoLog.Debugf("I/O error : %s", err)
s.Reset()
return err
}
if what != string(buf[:n]) {
return fmt.Errorf("echo output doesn't match input")
}
return nil
}