messagequeue.go 8.46 KB
Newer Older
1 2 3 4
package messagequeue

import (
	"context"
5 6
	"errors"
	"fmt"
7 8 9
	"sync"
	"time"

Hannah Howard's avatar
Hannah Howard committed
10 11
	logging "github.com/ipfs/go-log"
	"github.com/libp2p/go-libp2p-core/peer"
12

13 14
	gsmsg "github.com/ipfs/go-graphsync/message"
	gsnet "github.com/ipfs/go-graphsync/network"
15
	"github.com/ipfs/go-graphsync/notifications"
16 17 18 19 20 21
)

var log = logging.Logger("graphsync")

const maxRetries = 10

22 23 24
// max block size is the maximum size for batching blocks in a single payload
const maxBlockSize uint64 = 512 * 1024

25 26 27 28 29 30 31 32 33 34 35 36 37
type EventName uint64

const (
	Queued EventName = iota
	Sent
	Error
)

type Event struct {
	Name EventName
	Err  error
}

38 39 40 41 42 43 44
// MessageNetwork is any network that can connect peers and generate a message
// sender.
type MessageNetwork interface {
	NewMessageSender(context.Context, peer.ID) (gsnet.MessageSender, error)
	ConnectTo(context.Context, peer.ID) error
}

45
type Allocator interface {
46
	AllocateBlockMemory(p peer.ID, amount uint64) <-chan error
47 48 49 50
	ReleasePeerMemory(p peer.ID) error
	ReleaseBlockMemory(p peer.ID, amount uint64) error
}

51 52 53 54 55 56 57 58 59 60
// MessageQueue implements queue of want messages to send to peers.
type MessageQueue struct {
	p       peer.ID
	network MessageNetwork
	ctx     context.Context

	outgoingWork chan struct{}
	done         chan struct{}

	// internal do not touch outside go routines
61 62 63 64 65 66 67
	sender              gsnet.MessageSender
	eventPublisher      notifications.Publisher
	buildersLk          sync.RWMutex
	builders            []*gsmsg.Builder
	nextBuilderTopic    gsmsg.Topic
	allocatorSubscriber *notifications.TopicDataSubscriber
	allocator           Allocator
68 69 70
}

// New creats a new MessageQueue.
71
func New(ctx context.Context, p peer.ID, network MessageNetwork, allocator Allocator) *MessageQueue {
72
	return &MessageQueue{
73 74 75 76 77 78 79 80
		ctx:                 ctx,
		network:             network,
		p:                   p,
		outgoingWork:        make(chan struct{}, 1),
		done:                make(chan struct{}),
		eventPublisher:      notifications.NewPublisher(),
		allocator:           allocator,
		allocatorSubscriber: notifications.NewTopicDataSubscriber(&allocatorSubscriber{allocator, p}),
81 82 83
	}
}

84 85 86 87 88 89 90 91 92 93
// AllocateAndBuildMessage allows you to work modify the next message that is sent in the queue.
// If blkSize > 0, message building may block until enough memory has been freed from the queues to allocate the message.
func (mq *MessageQueue) AllocateAndBuildMessage(size uint64, buildMessageFn func(*gsmsg.Builder), notifees []notifications.Notifee) {
	if size > 0 {
		select {
		case <-mq.allocator.AllocateBlockMemory(mq.p, size):
		case <-mq.ctx.Done():
			return
		}
	}
94
	if mq.buildMessage(size, buildMessageFn, notifees) {
95 96 97 98
		mq.signalWork()
	}
}

99 100 101 102 103 104 105 106 107 108 109 110
func (mq *MessageQueue) buildMessage(size uint64, buildMessageFn func(*gsmsg.Builder), notifees []notifications.Notifee) bool {
	mq.buildersLk.Lock()
	defer mq.buildersLk.Unlock()
	if shouldBeginNewResponse(mq.builders, size) {
		topic := mq.nextBuilderTopic
		mq.nextBuilderTopic++
		mq.builders = append(mq.builders, gsmsg.NewBuilder(topic))
	}
	builder := mq.builders[len(mq.builders)-1]
	buildMessageFn(builder)
	for _, notifee := range notifees {
		notifications.SubscribeWithData(mq.eventPublisher, builder.Topic(), notifee)
111
	}
112 113 114 115 116 117 118 119 120 121 122
	return !builder.Empty()
}

func shouldBeginNewResponse(builders []*gsmsg.Builder, blkSize uint64) bool {
	if len(builders) == 0 {
		return true
	}
	if blkSize == 0 {
		return false
	}
	return builders[len(builders)-1].BlockSize()+blkSize > maxBlockSize
123 124
}

125 126 127 128 129 130 131 132 133 134 135 136
// Startup starts the processing of messages, and creates an initial message
// based on the given initial wantlist.
func (mq *MessageQueue) Startup() {
	go mq.runQueue()
}

// Shutdown stops the processing of messages for a message queue.
func (mq *MessageQueue) Shutdown() {
	close(mq.done)
}

func (mq *MessageQueue) runQueue() {
137 138 139 140
	defer func() {
		mq.allocator.ReleasePeerMemory(mq.p)
		mq.eventPublisher.Shutdown()
	}()
141
	mq.eventPublisher.Startup()
142 143 144 145 146
	for {
		select {
		case <-mq.outgoingWork:
			mq.sendMessage()
		case <-mq.done:
147 148
			select {
			case <-mq.outgoingWork:
149 150 151 152 153 154 155
				for {
					_, topic, err := mq.extractOutgoingMessage()
					if err == nil {
						mq.eventPublisher.Publish(topic, Event{Name: Error, Err: fmt.Errorf("message queue shutdown")})
					} else {
						break
					}
156 157 158
				}
			default:
			}
159 160 161 162 163 164
			if mq.sender != nil {
				mq.sender.Close()
			}
			return
		case <-mq.ctx.Done():
			if mq.sender != nil {
Hannah Howard's avatar
Hannah Howard committed
165
				_ = mq.sender.Reset()
166 167 168 169 170 171 172 173 174 175 176 177 178
			}
			return
		}
	}
}

func (mq *MessageQueue) signalWork() {
	select {
	case mq.outgoingWork <- struct{}{}:
	default:
	}
}

179 180 181
var errEmptyMessage = errors.New("Empty Message")

func (mq *MessageQueue) extractOutgoingMessage() (gsmsg.GraphSyncMessage, gsmsg.Topic, error) {
182
	// grab outgoing message
183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206
	mq.buildersLk.Lock()
	if len(mq.builders) == 0 {
		mq.buildersLk.Unlock()
		return gsmsg.GraphSyncMessage{}, gsmsg.Topic(0), errEmptyMessage
	}
	builder := mq.builders[0]
	mq.builders = mq.builders[1:]
	// if there are more queued messages, signal we still have more work
	if len(mq.builders) > 0 {
		select {
		case mq.outgoingWork <- struct{}{}:
		default:
		}
	}
	mq.buildersLk.Unlock()
	if builder.Empty() {
		return gsmsg.GraphSyncMessage{}, gsmsg.Topic(0), errEmptyMessage
	}
	notifications.SubscribeWithData(mq.eventPublisher, builder.Topic(), notifications.Notifee{
		Data:       builder.BlockSize(),
		Subscriber: mq.allocatorSubscriber,
	})
	message, err := builder.Build()
	return message, builder.Topic(), err
207 208 209
}

func (mq *MessageQueue) sendMessage() {
210 211 212 213 214
	message, topic, err := mq.extractOutgoingMessage()
	if err != nil {
		if err != errEmptyMessage {
			log.Errorf("Unable to assemble GraphSync message: %s", err.Error())
		}
215 216
		return
	}
217 218
	mq.eventPublisher.Publish(topic, Event{Name: Queued, Err: nil})
	defer mq.eventPublisher.Close(topic)
219

220
	err = mq.initializeSender()
221 222 223
	if err != nil {
		log.Infof("cant open message sender to peer %s: %s", mq.p, err)
		// TODO: cant connect, what now?
224
		mq.eventPublisher.Publish(topic, Event{Name: Error, Err: fmt.Errorf("cant open message sender to peer %s: %w", mq.p, err)})
225 226 227 228
		return
	}

	for i := 0; i < maxRetries; i++ { // try to send this message until we fail.
229
		if mq.attemptSendAndRecovery(message, topic) {
230 231 232
			return
		}
	}
233
	mq.eventPublisher.Publish(topic, Event{Name: Error, Err: fmt.Errorf("expended retries on SendMsg(%s)", mq.p)})
234 235 236 237 238 239 240 241 242 243 244 245 246 247
}

func (mq *MessageQueue) initializeSender() error {
	if mq.sender != nil {
		return nil
	}
	nsender, err := openSender(mq.ctx, mq.network, mq.p)
	if err != nil {
		return err
	}
	mq.sender = nsender
	return nil
}

248
func (mq *MessageQueue) attemptSendAndRecovery(message gsmsg.GraphSyncMessage, topic gsmsg.Topic) bool {
249 250
	err := mq.sender.SendMsg(mq.ctx, message)
	if err == nil {
251
		mq.eventPublisher.Publish(topic, Event{Name: Sent})
252 253 254 255
		return true
	}

	log.Infof("graphsync send error: %s", err)
Hannah Howard's avatar
Hannah Howard committed
256
	_ = mq.sender.Reset()
257 258 259
	mq.sender = nil

	select {
260
	case <-mq.done:
261
		mq.eventPublisher.Publish(topic, Event{Name: Error, Err: errors.New("queue shutdown")})
262
		return true
263
	case <-mq.ctx.Done():
264
		mq.eventPublisher.Publish(topic, Event{Name: Error, Err: errors.New("context cancelled")})
265 266 267
		return true
	case <-time.After(time.Millisecond * 100):
		// wait 100ms in case disconnect notifications are still propogating
Hannah Howard's avatar
Hannah Howard committed
268
		log.Warn("SendMsg errored but neither 'done' nor context.Done() were set")
269 270 271 272 273 274 275 276 277
	}

	err = mq.initializeSender()
	if err != nil {
		log.Infof("couldnt open sender again after SendMsg(%s) failed: %s", mq.p, err)
		// TODO(why): what do we do now?
		// I think the *right* answer is to probably put the message we're
		// trying to send back, and then return to waiting for new work or
		// a disconnect.
278
		mq.eventPublisher.Publish(topic, Event{Name: Error, Err: fmt.Errorf("couldnt open sender again after SendMsg(%s) failed: %w", mq.p, err)})
279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302
		return true
	}

	return false
}

func openSender(ctx context.Context, network MessageNetwork, p peer.ID) (gsnet.MessageSender, error) {
	// allow ten minutes for connections this includes looking them up in the
	// dht dialing them, and handshaking
	conctx, cancel := context.WithTimeout(ctx, time.Minute*10)
	defer cancel()

	err := network.ConnectTo(conctx, p)
	if err != nil {
		return nil, err
	}

	nsender, err := network.NewMessageSender(ctx, p)
	if err != nil {
		return nil, err
	}

	return nsender, nil
}
303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322

type allocatorSubscriber struct {
	allocator Allocator
	p         peer.ID
}

func (as *allocatorSubscriber) OnNext(topic notifications.Topic, event notifications.Event) {
	blkSize, ok := topic.(uint64)
	if !ok {
		return
	}
	ev, ok := event.(Event)
	if !ok || ev.Name == Queued {
		return
	}
	_ = as.allocator.ReleaseBlockMemory(as.p, blkSize)
}

func (as *allocatorSubscriber) OnClose(topic notifications.Topic) {
}