p2p/protocols: optimized accounting

p2p/protocols/accounting.go

@@ -17,8 +17,6 @@
 package protocols
 
 import (
-	"sync"
-
 	"github.com/ethereum/go-ethereum/metrics"
 )
 
@@ -36,87 +34,87 @@ var (
 	mChequesReceived = metrics.NewRegisteredCounterForced("account.cheques.received", nil)
 )
 
-type PriceOracle interface {
+type Prices interface {
-	Price(uint32, interface{}) (EntryDirection, uint64)
+	Price(interface{}) *Price
 }
 
-type BalanceManager interface {
+type Price struct {
-	//Credit is crediting the peer, charging local node
+	Value   int64 //Positive if sender pays, negative if receiver pays
-	Credit(peer *Peer, amount uint64) error
+	PerByte bool  //True if the price is per byte or for unit
-	//Debit is crediting the local node, charging the remote peer
-	Debit(peer *Peer, amount uint64) error
 }
 
-type EntryDirection uint
+func (p *Price) ForSender(size uint32) int64 {
-
+	price := p.Value
-const (
+	if p.PerByte {
-	ChargeSender   EntryDirection = 1
+		price *= int64(size)
-	ChargeReceiver EntryDirection = 2
+	}
-	ChargeNone     EntryDirection = 3
+	return price
-)
-
-type AccountingHook struct {
-	BalanceManager
-	PriceOracle
-	lock sync.RWMutex //lock the balances
 }
 
-func NewAccountingHook(mgr BalanceManager, po PriceOracle) *AccountingHook {
+func (p *Price) ForReceiver(size uint32) int64 {
-	ah := &AccountingHook{
+	price := p.Value
-		PriceOracle:    po,
+	if p.PerByte {
-		BalanceManager: mgr,
+		price *= int64(size)
+	}
+	return 0 - price
+}
+
+type Balance interface {
+	//Adds amount to the local balance with remote node `peer`;
+	//positive amount = credit local node
+	//negative amount = debit local node
+	Add(amount int64, peer *Peer) error
+}
+
+type Accounting struct {
+	Balance
+	Prices
+}
+
+func NewAccounting(mgr Balance, po Prices) *Accounting {
+	ah := &Accounting{
+		Prices:  po,
+		Balance: mgr,
 	}
 	return ah
 }
 
-func (ah *AccountingHook) Send(peer *Peer, size uint32, msg interface{}) error {
+func (ah *Accounting) Send(peer *Peer, size uint32, msg interface{}) error {
-	ah.lock.Lock()
+	price := ah.Price(msg)
-	defer ah.lock.Unlock()
+	if price == nil {
-	var err error
-	direction, price := ah.PriceOracle.Price(size, msg)
-	if direction == ChargeSender {
-		err = ah.BalanceManager.Debit(peer, price)
-		ah.debitMetrics(price, size, err)
-	} else if direction == ChargeReceiver {
-		err = ah.BalanceManager.Credit(peer, price)
-		ah.creditMetrics(price, size, err)
-	} else if direction == ChargeNone {
 		return nil
 	}
+	finalPrice := price.ForSender(size)
+	err := ah.Add(finalPrice, peer)
+	ah.doMetrics(finalPrice, size, err)
 	return err
 }
 
-func (ah *AccountingHook) Receive(peer *Peer, size uint32, msg interface{}) error {
+func (ah *Accounting) Receive(peer *Peer, size uint32, msg interface{}) error {
-	ah.lock.Lock()
+	price := ah.Price(msg)
-	defer ah.lock.Unlock()
+	if price == nil {
-	var err error
-	direction, price := ah.PriceOracle.Price(size, msg)
-	if direction == ChargeReceiver {
-		err = ah.BalanceManager.Debit(peer, price)
-		ah.debitMetrics(price, size, err)
-	} else if direction == ChargeSender {
-		err = ah.BalanceManager.Credit(peer, price)
-		ah.creditMetrics(price, size, err)
-	} else if direction == ChargeNone {
 		return nil
 	}
+	finalPrice := price.ForReceiver(size)
+	err := ah.Add(finalPrice, peer)
+	ah.doMetrics(finalPrice, size, err)
 	return err
 }
 
-func (ah *AccountingHook) debitMetrics(price uint64, size uint32, err error) {
+func (ah *Accounting) doMetrics(price int64, size uint32, err error) {
-	mBalanceDebit.Inc(int64(price))
+	if price > 0 {
-	mBytesDebit.Inc(int64(size))
-	mMsgDebit.Inc(1)
-	if err != nil {
-		mSelfDrops.Inc(1)
-	}
-}
-
-func (ah *AccountingHook) creditMetrics(price uint64, size uint32, err error) {
 		mBalanceCredit.Inc(int64(price))
 		mBytesCredit.Inc(int64(size))
 		mMsgCredit.Inc(1)
 		if err != nil {
 			mPeerDrops.Inc(1)
 		}
+	} else {
+		mBalanceDebit.Inc(int64(price))
+		mBytesDebit.Inc(int64(size))
+		mMsgDebit.Inc(1)
+		if err != nil {
+			mSelfDrops.Inc(1)
+		}
+	}
 }

p2p/protocols/accounting_test.go

@@ -19,7 +19,6 @@ package protocols
 import (
 	"bytes"
 	"context"
-	"errors"
 	"testing"
 	"time"
 
@@ -28,44 +27,176 @@ import (
 	"github.com/ethereum/go-ethereum/rlp"
 )
 
-var (
+type dummyBalance struct {
-	wouldHaveAccounted = errors.New("ignore this error")
+	amount int64
-)
+	peer   *Peer
+}
 
-type dummy struct {
+type dummyPrices struct{}
-	content string
+
+type perBytesMsg struct {
+	Content string
+}
+
+type perUnitMsg struct{}
+type zeroMsg struct{}
+
+func (d *dummyPrices) Price(msg interface{}) *Price {
+	switch msg.(type) {
+	case *perBytesMsg:
+		return &Price{
+			PerByte: true,
+			Value:   int64(100),
+		}
+	case *perUnitMsg:
+		return &Price{
+			PerByte: false,
+			Value:   int64(-99),
+		}
+	}
+	return nil
+}
+
+func (d *dummyBalance) Add(amount int64, peer *Peer) error {
+	d.amount = amount
+	d.peer = peer
+	return nil
+}
+
+func TestNoHook(t *testing.T) {
+	spec := createTestSpec()
+	id := adapters.RandomNodeConfig().ID
+	p := p2p.NewPeer(id, "testPeer", nil)
+	peer := NewPeer(p, &dummyRW{}, spec)
+	ctx := context.TODO()
+	msg := &perBytesMsg{Content: "testBalance"}
+	peer.Send(ctx, msg)
+	peer.handleIncoming(func(ctx context.Context, msg interface{}) error {
+		return nil
+	})
+}
+
+func TestSendBalance(t *testing.T) {
+	balance := &dummyBalance{}
+	prices := &dummyPrices{}
+
+	spec := createTestSpec()
+	spec.Hook = NewAccounting(balance, prices)
+
+	id := adapters.RandomNodeConfig().ID
+	p := p2p.NewPeer(id, "testPeer", nil)
+	peer := NewPeer(p, &dummyRW{}, spec)
+	ctx := context.TODO()
+	msg := &perBytesMsg{Content: "testBalance"}
+	size, _ := rlp.EncodeToBytes(msg)
+	peer.Send(ctx, msg)
+	if balance.amount != int64((len(size) * 100)) {
+		t.Fatalf("Expected price to be %d but is %d", (len(size) * 100), balance.amount)
+	}
+
+	msg2 := &perUnitMsg{}
+	peer.Send(ctx, msg2)
+	if balance.amount != int64(-99) {
+		t.Fatalf("Expected price to be %d but is %d", -99, balance.amount)
+	}
+
+	balance.amount = 77
+	msg3 := &zeroMsg{}
+	peer.Send(ctx, msg3)
+	if balance.amount != int64(77) {
+		t.Fatalf("Expected price to be %d but is %d", 77, balance.amount)
+	}
+}
+
+func TestReceiveBalance(t *testing.T) {
+	balance := &dummyBalance{}
+	prices := &dummyPrices{}
+
+	spec := createTestSpec()
+	spec.Hook = NewAccounting(balance, prices)
+
+	id := adapters.RandomNodeConfig().ID
+	p := p2p.NewPeer(id, "testPeer", nil)
+	rw := &dummyRW{}
+	peer := NewPeer(p, rw, spec)
+	msg := &perBytesMsg{Content: "testBalance"}
+	size, _ := rlp.EncodeToBytes(msg)
+
+	rw.msg = msg
+	rw.code, _ = spec.GetCode(msg)
+	err := peer.handleIncoming(func(ctx context.Context, msg interface{}) error {
+		return nil
+	})
+	if err != nil {
+		t.Fatalf("Expected no error, but got error: %v", err)
+	}
+	if balance.amount != int64((len(size) * (-100))) {
+		t.Fatalf("Expected price to be %d but is %d", (len(size) * (-100)), balance.amount)
+	}
+
+	msg2 := &perUnitMsg{}
+	rw.msg = msg2
+	rw.code, _ = spec.GetCode(msg2)
+	err = peer.handleIncoming(func(ctx context.Context, msg interface{}) error {
+		return nil
+	})
+	if err != nil {
+		t.Fatalf("Expected no error, but got error: %v", err)
+	}
+	if balance.amount != int64(99) {
+		t.Fatalf("Expected price to be %d but is %d", 99, balance.amount)
+	}
+
+	msg3 := &zeroMsg{}
+	rw.msg = msg3
+	rw.code, _ = spec.GetCode(msg3)
+	//need to reset cause no accounting won't overwrite
+	balance.amount = -888
+	err = peer.handleIncoming(func(ctx context.Context, msg interface{}) error {
+		return nil
+	})
+	if err != nil {
+		t.Fatalf("Expected no error, but got error: %v", err)
+	}
+
+	if balance.amount != int64(-888) {
+		t.Fatalf("Expected price to be %d but is %d", -888, balance.amount)
+	}
 }
 
 //dummy implementation of a MsgReadWriter
 //this allows for quick and easy unit tests without
 //having to build up the complete protocol
-type dummyRW struct{}
+type dummyRW struct {
+	msg  interface{}
+	size uint32
+	code uint64
+}
 
 func (d *dummyRW) WriteMsg(msg p2p.Msg) error {
 	return nil
 }
 
 func (d *dummyRW) ReadMsg() (p2p.Msg, error) {
+	enc := bytes.NewReader(d.getDummyMsg())
 	return p2p.Msg{
-		Code:       0,
+		Code:       d.code,
-		Size:       5,
+		Size:       d.size,
-		Payload:    bytes.NewReader(getDummyMsg()),
+		Payload:    enc,
 		ReceivedAt: time.Now(),
 	}, nil
 }
 
-func getDummyMsg() []byte {
+func (d *dummyRW) getDummyMsg() []byte {
-	msg := &dummy{content: "test"}
+	r, _ := rlp.EncodeToBytes(d.msg)
-	r, _ := rlp.EncodeToBytes(msg)
-
 	var b bytes.Buffer
 	wmsg := WrappedMsg{
 		Context: b.Bytes(),
 		Size:    uint32(len(r)),
 		Payload: r,
 	}
-
 	rr, _ := rlp.EncodeToBytes(wmsg)
+	d.size = uint32(len(rr))
 	return rr
 }
 
@@ -75,54 +206,10 @@ func createTestSpec() *Spec {
 		Version:    42,
 		MaxMsgSize: 10 * 1024,
 		Messages: []interface{}{
-			dummy{},
+			perBytesMsg{},
+			perUnitMsg{},
+			zeroMsg{},
 		},
 	}
 	return spec
 }
-
-type dummyBalanceMgr struct{}
-type dummyPriceOracle struct{}
-
-func (d *dummyPriceOracle) Price(uint32, interface{}) (EntryDirection, uint64) {
-	return ChargeSender, 99
-}
-
-func (d *dummyBalanceMgr) Credit(peer *Peer, amount uint64) error {
-	return wouldHaveAccounted
-}
-
-func (d *dummyBalanceMgr) Debit(peer *Peer, amount uint64) error {
-	return wouldHaveAccounted
-}
-
-// Test that passing a nil hook doesn't affect sending
-func TestProtocolNilHook(t *testing.T) {
-	spec := createTestSpec()
-	id := adapters.RandomNodeConfig().ID
-	p := p2p.NewPeer(id, "testPeer", nil)
-	peer := NewPeer(p, &dummyRW{}, spec)
-
-	peer.Send(context.Background(), dummy{})
-	peer.handleIncoming(func(ctx context.Context, msg interface{}) error {
-		return nil
-	})
-}
-
-func TestProtocolHook(t *testing.T) {
-	spec := createTestSpec()
-	spec.Hook = NewAccountingHook(&dummyBalanceMgr{}, &dummyPriceOracle{})
-	id := adapters.RandomNodeConfig().ID
-	p := p2p.NewPeer(id, "testPeer", nil)
-	peer := NewPeer(p, &dummyRW{}, spec)
-
-	err := peer.Send(context.Background(), dummy{})
-	if err == nil || err != wouldHaveAccounted {
-		t.Fatal("Expected fake accounting to happen, but didn't")
-	}
-
-	err = peer.handleIncoming(nil)
-	if err == nil || err != wouldHaveAccounted {
-		t.Fatal("Expected fake accounting to happen, but didn't")
-	}
-}

p2p/protocols/protocol_test.go

@@ -197,6 +197,139 @@ func TestProtoHandshakeSuccess(t *testing.T) {
 	runProtoHandshake(t, &protoHandshake{42, "420"})
 }
 
+type dummyHook struct {
+	peer  *Peer
+	size  uint32
+	msg   interface{}
+	send  bool
+	err   error
+	waitC chan struct{}
+}
+
+type dummyMsg struct {
+	Content string
+}
+
+func (d *dummyHook) Send(peer *Peer, size uint32, msg interface{}) error {
+	d.peer = peer
+	d.size = size
+	d.msg = msg
+	d.send = true
+	return d.err
+}
+
+func (d *dummyHook) Receive(peer *Peer, size uint32, msg interface{}) error {
+	d.peer = peer
+	d.size = size
+	d.msg = msg
+	d.send = false
+	d.waitC <- struct{}{}
+	return d.err
+}
+
+func TestProtocolHook(t *testing.T) {
+	testHook := &dummyHook{
+		waitC: make(chan struct{}, 1),
+	}
+	spec := &Spec{
+		Name:       "test",
+		Version:    42,
+		MaxMsgSize: 10 * 1024,
+		Messages: []interface{}{
+			dummyMsg{},
+		},
+		Hook: testHook,
+	}
+
+	runFunc := func(p *p2p.Peer, rw p2p.MsgReadWriter) error {
+		peer := NewPeer(p, rw, spec)
+		ctx := context.TODO()
+		peer.Send(ctx, &dummyMsg{
+			Content: "handshake"})
+
+		handle := func(ctx context.Context, msg interface{}) error {
+			return nil
+		}
+
+		return peer.Run(handle)
+	}
+
+	conf := adapters.RandomNodeConfig()
+	tester := p2ptest.NewProtocolTester(t, conf.ID, 2, runFunc)
+	err := tester.TestExchanges(p2ptest.Exchange{
+		Expects: []p2ptest.Expect{
+			{
+				Code: 0,
+				Msg:  &dummyMsg{Content: "handshake"},
+				Peer: tester.Nodes[0].ID(),
+			},
+		},
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if testHook.msg == nil || testHook.msg.(*dummyMsg).Content != "handshake" {
+		t.Fatal("Expected msg to be set, but it is not")
+	}
+	if testHook.send != true {
+		t.Fatal("Expected a send message, but it is not")
+	}
+	if testHook.peer == nil || testHook.peer.ID() != tester.Nodes[0].ID() {
+		t.Fatal("Expected peer ID to be set correctly, but it is not")
+	}
+	if testHook.size != 11 { //11 is the length of the encoded message
+		t.Fatalf("Expected size to be %d, but it is %d ", 1, testHook.size)
+	}
+
+	err = tester.TestExchanges(p2ptest.Exchange{
+		Triggers: []p2ptest.Trigger{
+			{
+				Code: 0,
+				Msg:  &dummyMsg{Content: "response"},
+				Peer: tester.Nodes[1].ID(),
+			},
+		},
+	})
+
+	<-testHook.waitC
+
+	if err != nil {
+		t.Fatal(err)
+	}
+	if testHook.msg == nil || testHook.msg.(*dummyMsg).Content != "response" {
+		t.Fatal("Expected msg to be set, but it is not")
+	}
+	if testHook.send != false {
+		t.Fatal("Expected a send message, but it is not")
+	}
+	if testHook.peer == nil || testHook.peer.ID() != tester.Nodes[1].ID() {
+		t.Fatal("Expected peer ID to be set correctly, but it is not")
+	}
+	if testHook.size != 10 { //11 is the length of the encoded message
+		t.Fatalf("Expected size to be %d, but it is %d ", 1, testHook.size)
+	}
+
+	testHook.err = fmt.Errorf("dummy error")
+	err = tester.TestExchanges(p2ptest.Exchange{
+		Triggers: []p2ptest.Trigger{
+			{
+				Code: 0,
+				Msg:  &dummyMsg{Content: "response"},
+				Peer: tester.Nodes[1].ID(),
+			},
+		},
+	})
+
+	<-testHook.waitC
+
+	time.Sleep(100 * time.Millisecond)
+	err = tester.TestDisconnected(&p2ptest.Disconnect{tester.Nodes[1].ID(), testHook.err})
+	if err != nil {
+		t.Fatalf("Expected a specific disconnect error, but got different one: %v", err)
+	}
+
+}
+
 func moduleHandshakeExchange(id enode.ID, resp uint) []p2ptest.Exchange {
 
 	return []p2ptest.Exchange{

swarm/network/stream/stream.go

@@ -62,8 +62,8 @@ type Registry struct {
 	intervalsStore state.Store
 	doRetrieve     bool
 	maxPeerServers int
-	balanceMgr     protocols.BalanceManager
+	balanceMgr     protocols.Balance
-	priceOracle    protocols.PriceOracle
+	prices         protocols.Prices
 	spec           *protocols.Spec
 }
 
@@ -77,7 +77,7 @@ type RegistryOptions struct {
 }
 
 // NewRegistry is Streamer constructor
-func NewRegistry(localID enode.ID, delivery *Delivery, syncChunkStore storage.SyncChunkStore, intervalsStore state.Store, options *RegistryOptions, balanceMgr protocols.BalanceManager) *Registry {
+func NewRegistry(localID enode.ID, delivery *Delivery, syncChunkStore storage.SyncChunkStore, intervalsStore state.Store, options *RegistryOptions, balanceMgr protocols.Balance) *Registry {
 	if options == nil {
 		options = &RegistryOptions{}
 	}
@@ -194,8 +194,8 @@ func NewRegistry(localID enode.ID, delivery *Delivery, syncChunkStore storage.Sy
 func (r *Registry) setupSpec() {
 	r.createSpec()
 	r.createPriceOracle()
-	if !reflect.ValueOf(r.balanceMgr).IsNil() {
+	if r.balanceMgr != nil && !reflect.ValueOf(r.balanceMgr).IsNil() {
-		r.spec.Hook = protocols.NewAccountingHook(r.balanceMgr, r.priceOracle)
+		r.spec.Hook = protocols.NewAccounting(r.balanceMgr, r.prices)
 	}
 }
 
@@ -718,65 +718,53 @@ func (r *Registry) createSpec() {
 
 //An accountable message needs some meta information attached to it
 //in order to evaluate the correct price
-type priceTag struct {
+type StreamerPrices struct {
-	price     uint64
+	priceMatrix map[uint64]*protocols.Price
-	sizeBased bool
-	direction protocols.EntryDirection
-}
-type StreamerPriceOracle struct {
-	priceMatrix map[uint64]*priceTag
 	registry    *Registry
 }
 
-func (spo *StreamerPriceOracle) Price(size uint32, msg interface{}) (direction protocols.EntryDirection, price uint64) {
+func (spo *StreamerPrices) Price(msg interface{}) *protocols.Price {
 	code, ok := spo.registry.spec.GetCode(msg)
-	direction = protocols.ChargeNone
 	if !ok {
 		//this could be handled more severely (panic) but let's be gentle?
-		return
+		return nil
 	}
 
-	tag, ok := spo.priceMatrix[code]
+	price, ok := spo.priceMatrix[code]
 	if !ok {
-		return
+		return nil
 	}
-	price = tag.price
+
-	if tag.sizeBased {
+	return price
-		price = price * uint64(size)
-	}
-	direction = tag.direction
-	return
 }
 
 func (r *Registry) createPriceOracle() {
 
-	po := &StreamerPriceOracle{
+	po := &StreamerPrices{
 		registry: r,
 	}
-	po.priceMatrix = make(map[uint64]*priceTag)
+	po.priceMatrix = make(map[uint64]*protocols.Price)
 
 	deliveryCode, ok := r.spec.GetCode(ChunkDeliveryMsgRetrieval{})
 	if !ok {
 		panic("Attempting to get message code for an expected message type, but code not found")
 	}
-	tag := &priceTag{
+	price := &protocols.Price{
-		price:     uint64(100),
+		Value:   int64(-100),
-		sizeBased: true,
+		PerByte: true,
-		direction: protocols.ChargeReceiver,
 	}
-	po.priceMatrix[deliveryCode] = tag
+	po.priceMatrix[deliveryCode] = price
 
 	retrieveReqCode, ok := r.spec.GetCode(RetrieveRequestMsg{})
 	if !ok {
 		panic("Attempting to get message code for an expected message type, but code not found")
 	}
-	tag = &priceTag{
+	price = &protocols.Price{
-		price:     uint64(10),
+		Value:   int64(10),
-		sizeBased: false,
+		PerByte: false,
-		direction: protocols.ChargeSender,
 	}
-	po.priceMatrix[retrieveReqCode] = tag
+	po.priceMatrix[retrieveReqCode] = price
-	r.priceOracle = po
+	r.prices = po
 }
 
 func (r *Registry) Protocols() []p2p.Protocol {

swarm/swap/swap.go

@@ -39,13 +39,14 @@ type Swap struct {
 }
 
 //Credit us and debit remote
-func (s *Swap) Credit(peer *protocols.Peer, amount uint64) (err error) {
+func (s *Swap) Add(amount int64, peer *protocols.Peer) (err error) {
 	s.lock.Lock()
 	defer s.lock.Unlock()
 
 	s.loadState(peer)
 
-	s.balances[peer.ID()] += int64(amount)
+	s.balances[peer.ID()] += amount
+
 	peerBalance := s.balances[peer.ID()]
 	s.stateStore.Put(peer.ID().String(), &peerBalance)
 
@@ -53,22 +54,6 @@ func (s *Swap) Credit(peer *protocols.Peer, amount uint64) (err error) {
 	return err
 }
 
-//Debit us and credit remote
-func (s *Swap) Debit(peer *protocols.Peer, amount uint64) (err error) {
-	s.lock.Lock()
-	defer s.lock.Unlock()
-
-	s.loadState(peer)
-
-	//local node is being debited (in favor of remote peer), so its balance decreases
-	s.balances[peer.ID()] -= int64(amount)
-	peerBalance := s.balances[peer.ID()]
-	s.stateStore.Put(peer.ID().String(), &peerBalance)
-
-	log.Debug(fmt.Sprintf("balance for peer %s: %s", peer.ID().String(), strconv.FormatInt(peerBalance, 10)))
-	return nil
-}
-
 //get a peer's balance
 func (swap *Swap) GetPeerBalance(peer enode.ID) (int64, error) {
 	swap.lock.RLock()

swarm/swap/swap_test.go

@@ -17,21 +17,17 @@
 package swap
 
 import (
-	"context"
-	"crypto/rand"
 	"flag"
 	"fmt"
 	"io/ioutil"
 	mrand "math/rand"
 	"os"
 	"testing"
-	"time"
 
 	"github.com/ethereum/go-ethereum/log"
 	"github.com/ethereum/go-ethereum/p2p"
 	"github.com/ethereum/go-ethereum/p2p/protocols"
 	"github.com/ethereum/go-ethereum/p2p/simulations/adapters"
-	"github.com/ethereum/go-ethereum/rlp"
 	"github.com/ethereum/go-ethereum/swarm/state"
 	colorable "github.com/mattn/go-colorable"
 )
@@ -40,76 +36,6 @@ var (
 	loglevel = flag.Int("loglevel", 2, "verbosity of logs")
 )
 
-type testPriceOracle struct{}
-
-func (tpo *testPriceOracle) Accountable(msg interface{}) bool {
-	_, ok := testSpec.GetCode(msg)
-	return ok
-}
-
-func (tpo *testPriceOracle) Price(size uint32, msg interface{}) (protocols.EntryDirection, uint64) {
-	switch msg := msg.(type) {
-	case *testCheapMsg:
-		return msg.Price(size)
-	case *testSizeBasedMsg:
-		return msg.Price(size)
-	case *testChargeRecvMsg:
-		return msg.Price(size)
-	}
-	return protocols.ChargeNone, 0
-}
-
-var testSpec = &protocols.Spec{
-	Name:       "swapTestSpec",
-	Version:    1,
-	MaxMsgSize: 10 * 1024 * 1024,
-	Messages: []interface{}{
-		testCheapMsg{},
-		testSizeBasedMsg{},
-		testChargeRecvMsg{},
-	},
-}
-
-//dummy implementation of a MsgReadWriter
-//this allows for quick and easy unit tests without
-//having to build up the complete protocol
-type dummyRW struct{}
-
-func (d *dummyRW) WriteMsg(msg p2p.Msg) error {
-	return nil
-}
-
-func (d *dummyRW) ReadMsg() (p2p.Msg, error) {
-	return p2p.Msg{
-		Code:       0,
-		Size:       0,
-		Payload:    nil,
-		ReceivedAt: time.Now(),
-	}, nil
-}
-
-//define a couple of messages for tests
-type testCheapMsg struct{}
-type testSizeBasedMsg struct {
-	Data []byte
-}
-type testChargeRecvMsg struct{}
-
-//a message with an arbitrary cost
-func (tmsg *testCheapMsg) Price(size uint32) (protocols.EntryDirection, uint64) {
-	return protocols.ChargeSender, uint64(10)
-}
-
-//a message which needs to be charged based on price
-func (tmsg *testSizeBasedMsg) Price(size uint32) (protocols.EntryDirection, uint64) {
-	return protocols.ChargeSender, uint64(size) * uint64(100)
-}
-
-//a message which needs to be charged based on price
-func (tmsg *testChargeRecvMsg) Price(size uint32) (protocols.EntryDirection, uint64) {
-	return protocols.ChargeReceiver, uint64(999)
-}
-
 func init() {
 	flag.Parse()
 
@@ -127,7 +53,7 @@ func TestRepeatedBookings(t *testing.T) {
 	amount := mrand.Intn(100)
 	cnt := 1 + mrand.Intn(10)
 	for i := 0; i < cnt; i++ {
-		swap.Credit(testPeer.Peer, uint64(amount))
+		swap.Add(int64(amount), testPeer.Peer)
 	}
 	expectedBalance := int64(cnt * amount)
 	realBalance := swap.balances[testPeer.ID()]
@@ -139,7 +65,7 @@ func TestRepeatedBookings(t *testing.T) {
 	amount = mrand.Intn(100)
 	cnt = 1 + mrand.Intn(10)
 	for i := 0; i < cnt; i++ {
-		swap.Debit(testPeer2.Peer, uint64(amount))
+		swap.Add(0-int64(amount), testPeer2.Peer)
 	}
 	expectedBalance = int64(0 - (cnt * amount))
 	realBalance = swap.balances[testPeer2.ID()]
@@ -149,13 +75,13 @@ func TestRepeatedBookings(t *testing.T) {
 
 	//mixed debits and credits
 	amount1 := mrand.Intn(100)
-	amount2 := mrand.Intn(100)
+	amount2 := mrand.Intn(55)
-	amount3 := mrand.Intn(100)
+	amount3 := mrand.Intn(999)
-	swap.Credit(testPeer2.Peer, uint64(amount1))
+	swap.Add(int64(amount1), testPeer2.Peer)
-	swap.Credit(testPeer2.Peer, uint64(amount2))
+	swap.Add(int64(0-amount2), testPeer2.Peer)
-	swap.Debit(testPeer2.Peer, uint64(amount3))
+	swap.Add(int64(0-amount3), testPeer2.Peer)
 
-	expectedBalance = expectedBalance + int64(amount1+amount2-amount3)
+	expectedBalance = expectedBalance + int64(amount1-amount2-amount3)
 	realBalance = swap.balances[testPeer2.ID()]
 
 	if expectedBalance != realBalance {
@@ -163,25 +89,6 @@ func TestRepeatedBookings(t *testing.T) {
 	}
 }
 
-//send a message with cost,
-//then check that the balance has the expected amount
-func TestSendCheapMessage(t *testing.T) {
-	//create a test swap account
-	swap, testDir := createTestSwap(t)
-	defer os.RemoveAll(testDir)
-
-	msg := &testCheapMsg{}
-	ctx := context.Background()
-	testPeer := newDummyPeer()
-	testPeer.Send(ctx, msg)
-
-	//check the new balance
-	_, price := msg.Price(0)
-	if swap.balances[testPeer.ID()] != int64(0-price) {
-		t.Fatalf("Expected balance to be %d, but is %d", price, swap.balances[testPeer.ID()])
-	}
-}
-
 //try restoring a balance from state store
 //this is simulated by creating a node,
 //assigning it an arbitrary balance,
@@ -193,25 +100,11 @@ func TestRestoreBalanceFromStateStore(t *testing.T) {
 	swap, testDir := createTestSwap(t)
 	defer os.RemoveAll(testDir)
 
-	msg := &testCheapMsg{}
-	ctx := context.Background()
 	testPeer := newDummyPeer()
-	testPeer.Send(ctx, msg)
+	swap.balances[testPeer.ID()] = -8888
 
-	//check the new balance
+	tmpBalance := swap.balances[testPeer.ID()]
-	_, price := msg.Price(0)
+	swap.stateStore.Put(testPeer.ID().String(), &tmpBalance)
-	expectedBalance := int64(0 - price)
-	if swap.balances[testPeer.ID()] != expectedBalance {
-		t.Fatalf("Expected balance to be %d, but is %d", price, swap.balances[testPeer.ID()])
-	}
-
-	var tmpBalance int64
-	swap.stateStore.Get(testPeer.ID().String(), &tmpBalance)
-	//compare the balances
-	if expectedBalance != tmpBalance {
-		t.Fatal(fmt.Sprintf("Unexpected balance value in stateStore after sending cheap message test - stateStore should have saved it. Expected balance: %d, balance is: %d",
-			expectedBalance, tmpBalance))
-	}
 
 	swap.stateStore.Close()
 	swap.stateStore = nil
@@ -225,63 +118,9 @@ func TestRestoreBalanceFromStateStore(t *testing.T) {
 	stateStore.Get(testPeer.ID().String(), &newBalance)
 
 	//compare the balances
-	if expectedBalance != newBalance {
+	if tmpBalance != newBalance {
 		t.Fatal(fmt.Sprintf("Unexpected balance value after sending cheap message test. Expected balance: %d, balance is: %d",
-			expectedBalance, newBalance))
+			tmpBalance, newBalance))
-	}
-}
-
-//send a message with cost,
-//then check that the balance has the expected amount
-//the message is charged size based
-func TestSendSizeBasedMsg(t *testing.T) {
-	//create a test swap account
-	swap, testDir := createTestSwap(t)
-	defer os.RemoveAll(testDir)
-
-	msg := &testSizeBasedMsg{}
-	msg.Data = make([]byte, 16)
-	rand.Read(msg.Data)
-	ctx := context.Background()
-	testPeer := newDummyPeer()
-	testPeer.Send(ctx, msg)
-
-	//check the new balance
-	r, err := rlp.EncodeToBytes(msg)
-	if err != nil {
-		t.Fatal(err)
-	}
-
-	size := uint32(len(r))
-	_, price := msg.Price(size)
-	expectedBalance := int64(0 - price)
-	expectedPrice := uint64(size * 100)
-	if price != expectedPrice {
-		t.Fatalf("Expected price to be %d, but is %d", expectedPrice, price)
-	}
-	if swap.balances[testPeer.ID()] != expectedBalance {
-		t.Fatalf("Expected balance to be %d, but is %d", expectedBalance, swap.balances[testPeer.ID()])
-	}
-}
-
-//charge as being the receiver of a receiver-pays message
-//as this test for simplicity only sends messages,
-//it means that the balance is changed in positive,
-//instead of negative as with all other tests
-func TestChargeAsReceiver(t *testing.T) {
-	//create a test swap account
-	swap, testDir := createTestSwap(t)
-	defer os.RemoveAll(testDir)
-
-	msg := &testChargeRecvMsg{}
-	ctx := context.Background()
-	testPeer := newDummyPeer()
-	testPeer.Send(ctx, msg)
-
-	//check the new balance
-	_, price := msg.Price(0)
-	if swap.balances[testPeer.ID()] != int64(price) {
-		t.Fatalf("Expected balance to be %d, but is %d", price, swap.balances[testPeer.ID()])
 	}
 }
 
@@ -297,28 +136,17 @@ func createTestSwap(t *testing.T) (*Swap, string) {
 		t.Fatal(err2)
 	}
 	swap := New(stateStore)
-	testSpec.Hook = protocols.NewAccountingHook(swap, &testPriceOracle{})
 	return swap, dir
 }
 
-//dummy message handler (needed or we will have a panic in the accounting)
-func dummyMsgHandler(ctx context.Context, msg interface{}) error {
-	return nil
-}
-
 type dummyPeer struct {
 	*protocols.Peer
-	testFunc func(error)
-}
-
-func (dp *dummyPeer) Drop(err error) {
-	dp.testFunc(err)
 }
 
 //creates a dummy protocols.Peer with dummy MsgReadWriter
 func newDummyPeer() *dummyPeer {
 	id := adapters.RandomNodeConfig().ID
-	protoPeer := protocols.NewPeer(p2p.NewPeer(id, "testPeer", nil), &dummyRW{}, testSpec)
+	protoPeer := protocols.NewPeer(p2p.NewPeer(id, "testPeer", nil), nil, nil)
 	dummy := &dummyPeer{
 		Peer: protoPeer,
 	}