refactor() message API

performing CR in the form of a PR. Let me know what you think.

License: MIT
Signed-off-by: Brian Tiger Chow <brian@perfmode.com>

bitswap.go

@@ -169,7 +169,7 @@ func (bs *bitswap) sendWantListTo(ctx context.Context, peers <-chan peer.Peer) e
 	}
 	message := bsmsg.New()
 	for _, wanted := range bs.wantlist.Entries() {
-		message.AddEntry(wanted.Value, wanted.Priority, false)
+		message.AddEntry(wanted.Value, wanted.Priority)
 	}
 	wg := sync.WaitGroup{}
 	for peerToQuery := range peers {
@@ -207,7 +207,7 @@ func (bs *bitswap) sendWantlistToProviders(ctx context.Context, wantlist *wl.Wan
 	message := bsmsg.New()
 	message.SetFull(true)
 	for _, e := range bs.wantlist.Entries() {
-		message.AddEntry(e.Value, e.Priority, false)
+		message.AddEntry(e.Value, e.Priority)
 	}
 
 	ps := pset.NewPeerSet()
@@ -335,7 +335,7 @@ func (bs *bitswap) cancelBlocks(ctx context.Context, bkeys []u.Key) {
 	message := bsmsg.New()
 	message.SetFull(false)
 	for _, k := range bkeys {
-		message.AddEntry(k, 0, true)
+		message.Cancel(k)
 	}
 	for _, p := range bs.ledgermanager.Peers() {
 		err := bs.send(ctx, p, message)

message/message.go

@@ -24,7 +24,9 @@ type BitSwapMessage interface {
 	Blocks() []*blocks.Block
 
 	// AddEntry adds an entry to the Wantlist.
-	AddEntry(key u.Key, priority int, cancel bool)
+	AddEntry(key u.Key, priority int)
+
+	Cancel(key u.Key)
 
 	// Sets whether or not the contained wantlist represents the entire wantlist
 	// true = full wantlist
@@ -50,6 +52,10 @@ type impl struct {
 }
 
 func New() BitSwapMessage {
+	return newMsg()
+}
+
+func newMsg() *impl {
 	return &impl{
 		blocks:   make(map[u.Key]*blocks.Block),
 		wantlist: make(map[u.Key]*Entry),
@@ -64,10 +70,10 @@ type Entry struct {
 }
 
 func newMessageFromProto(pbm pb.Message) BitSwapMessage {
-	m := New()
+	m := newMsg()
 	m.SetFull(pbm.GetWantlist().GetFull())
 	for _, e := range pbm.GetWantlist().GetEntries() {
-		m.AddEntry(u.Key(e.GetBlock()), int(e.GetPriority()), e.GetCancel())
+		m.addEntry(u.Key(e.GetBlock()), int(e.GetPriority()), e.GetCancel())
 	}
 	for _, d := range pbm.GetBlocks() {
 		b := blocks.NewBlock(d)
@@ -100,7 +106,15 @@ func (m *impl) Blocks() []*blocks.Block {
 	return bs
 }
 
-func (m *impl) AddEntry(k u.Key, priority int, cancel bool) {
+func (m *impl) Cancel(k u.Key) {
+	m.addEntry(k, 0, true)
+}
+
+func (m *impl) AddEntry(k u.Key, priority int) {
+	m.addEntry(k, priority, false)
+}
+
+func (m *impl) addEntry(k u.Key, priority int, cancel bool) {
 	e, exists := m.wantlist[k]
 	if exists {
 		e.Priority = priority

message/message_test.go

@@ -14,7 +14,7 @@ import (
 func TestAppendWanted(t *testing.T) {
 	const str = "foo"
 	m := New()
-	m.AddEntry(u.Key(str), 1, false)
+	m.AddEntry(u.Key(str), 1)
 
 	if !wantlistContains(m.ToProto().GetWantlist(), str) {
 		t.Fail()
@@ -63,7 +63,7 @@ func TestWantlist(t *testing.T) {
 	keystrs := []string{"foo", "bar", "baz", "bat"}
 	m := New()
 	for _, s := range keystrs {
-		m.AddEntry(u.Key(s), 1, false)
+		m.AddEntry(u.Key(s), 1)
 	}
 	exported := m.Wantlist()
 
@@ -86,7 +86,7 @@ func TestCopyProtoByValue(t *testing.T) {
 	const str = "foo"
 	m := New()
 	protoBeforeAppend := m.ToProto()
-	m.AddEntry(u.Key(str), 1, false)
+	m.AddEntry(u.Key(str), 1)
 	if wantlistContains(protoBeforeAppend.GetWantlist(), str) {
 		t.Fail()
 	}
@@ -94,11 +94,11 @@ func TestCopyProtoByValue(t *testing.T) {
 
 func TestToNetFromNetPreservesWantList(t *testing.T) {
 	original := New()
-	original.AddEntry(u.Key("M"), 1, false)
-	original.AddEntry(u.Key("B"), 1, false)
-	original.AddEntry(u.Key("D"), 1, false)
-	original.AddEntry(u.Key("T"), 1, false)
-	original.AddEntry(u.Key("F"), 1, false)
+	original.AddEntry(u.Key("M"), 1)
+	original.AddEntry(u.Key("B"), 1)
+	original.AddEntry(u.Key("D"), 1)
+	original.AddEntry(u.Key("T"), 1)
+	original.AddEntry(u.Key("F"), 1)
 
 	var buf bytes.Buffer
 	if err := original.ToNet(&buf); err != nil {
@@ -174,8 +174,8 @@ func TestDuplicates(t *testing.T) {
 	b := blocks.NewBlock([]byte("foo"))
 	msg := New()
 
-	msg.AddEntry(b.Key(), 1, false)
-	msg.AddEntry(b.Key(), 1, false)
+	msg.AddEntry(b.Key(), 1)
+	msg.AddEntry(b.Key(), 1)
 	if len(msg.Wantlist()) != 1 {
 		t.Fatal("Duplicate in BitSwapMessage")
 	}

strategy/ledgermanager_test.go

@@ -64,7 +64,7 @@ func TestBlockRecordedAsWantedAfterMessageReceived(t *testing.T) {
 	block := blocks.NewBlock([]byte("data wanted by beggar"))
 
 	messageFromBeggarToChooser := message.New()
-	messageFromBeggarToChooser.AddEntry(block.Key(), 1, false)
+	messageFromBeggarToChooser.AddEntry(block.Key(), 1)
 
 	chooser.ls.MessageReceived(beggar.Peer, messageFromBeggarToChooser)
 	// for this test, doesn't matter if you record that beggar sent