responsemanager.go 8.9 KB
Newer Older
1 2 3 4
package responsemanager

import (
	"context"
5
	"sync"
6 7
	"time"

8
	"github.com/ipfs/go-graphsync"
9
	"github.com/ipfs/go-graphsync/ipldutil"
10
	gsmsg "github.com/ipfs/go-graphsync/message"
11
	"github.com/ipfs/go-graphsync/responsemanager/loader"
12
	"github.com/ipfs/go-graphsync/responsemanager/peerresponsemanager"
13
	"github.com/ipfs/go-peertaskqueue/peertask"
14 15
	ipld "github.com/ipld/go-ipld-prime"
	cidlink "github.com/ipld/go-ipld-prime/linking/cid"
16
	"github.com/ipld/go-ipld-prime/traversal"
17
	"github.com/libp2p/go-libp2p-core/peer"
18 19 20 21 22 23 24 25 26 27
)

const (
	maxInProcessRequests = 6
	thawSpeed            = time.Millisecond * 100
)

type inProgressResponseStatus struct {
	ctx      context.Context
	cancelFn func()
28
	request  gsmsg.GraphSyncRequest
29 30 31 32
}

type responseKey struct {
	p         peer.ID
33
	requestID graphsync.RequestID
34 35 36
}

type responseTaskData struct {
37 38
	ctx     context.Context
	request gsmsg.GraphSyncRequest
39 40
}

41
type requestHook struct {
42
	key  uint64
43
	hook graphsync.OnRequestReceivedHook
44 45
}

46 47 48
// QueryQueue is an interface that can receive new selector query tasks
// and prioritize them as needed, and pop them off later
type QueryQueue interface {
49 50 51 52
	PushTasks(to peer.ID, tasks ...peertask.Task)
	PopTasks(targetMinWork int) (peer.ID, []*peertask.Task, int)
	Remove(topic peertask.Topic, p peer.ID)
	TasksDone(to peer.ID, tasks ...*peertask.Task)
53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69
	ThawRound()
}

// PeerManager is an interface that returns sender interfaces for peer responses.
type PeerManager interface {
	SenderForPeer(p peer.ID) peerresponsemanager.PeerResponseSender
}

type responseManagerMessage interface {
	handle(rm *ResponseManager)
}

// ResponseManager handles incoming requests from the network, initiates selector
// traversals, and transmits responses
type ResponseManager struct {
	ctx         context.Context
	cancelFn    context.CancelFunc
70
	loader      ipld.Loader
71 72 73 74 75 76 77
	peerManager PeerManager
	queryQueue  QueryQueue

	messages            chan responseManagerMessage
	workSignal          chan struct{}
	ticker              *time.Ticker
	inProgressResponses map[responseKey]inProgressResponseStatus
78 79
	requestHooksLk      sync.RWMutex
	requestHookNextKey  uint64
80
	requestHooks        []requestHook
81 82 83 84 85
}

// New creates a new response manager from the given context, loader,
// bridge to IPLD interface, peerManager, and queryQueue.
func New(ctx context.Context,
86
	loader ipld.Loader,
87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108
	peerManager PeerManager,
	queryQueue QueryQueue) *ResponseManager {
	ctx, cancelFn := context.WithCancel(ctx)
	return &ResponseManager{
		ctx:                 ctx,
		cancelFn:            cancelFn,
		loader:              loader,
		peerManager:         peerManager,
		queryQueue:          queryQueue,
		messages:            make(chan responseManagerMessage, 16),
		workSignal:          make(chan struct{}, 1),
		ticker:              time.NewTicker(thawSpeed),
		inProgressResponses: make(map[responseKey]inProgressResponseStatus),
	}
}

type processRequestMessage struct {
	p        peer.ID
	requests []gsmsg.GraphSyncRequest
}

// ProcessRequests processes incoming requests for the given peer
109
func (rm *ResponseManager) ProcessRequests(ctx context.Context, p peer.ID, requests []gsmsg.GraphSyncRequest) {
110 111 112
	select {
	case rm.messages <- &processRequestMessage{p, requests}:
	case <-rm.ctx.Done():
113
	case <-ctx.Done():
114 115 116
	}
}

117
// RegisterHook registers an extension to process new incoming requests
118 119 120 121 122 123 124 125 126 127 128 129 130 131 132
func (rm *ResponseManager) RegisterHook(hook graphsync.OnRequestReceivedHook) graphsync.UnregisterHookFunc {
	rm.requestHooksLk.Lock()
	rh := requestHook{rm.requestHookNextKey, hook}
	rm.requestHookNextKey++
	rm.requestHooks = append(rm.requestHooks, rh)
	rm.requestHooksLk.Unlock()
	return func() {
		rm.requestHooksLk.Lock()
		defer rm.requestHooksLk.Unlock()
		for i, matchHook := range rm.requestHooks {
			if rh.key == matchHook.key {
				rm.requestHooks = append(rm.requestHooks[:i], rm.requestHooks[i+1:]...)
				return
			}
		}
133 134 135
	}
}

136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162
type synchronizeMessage struct {
	sync chan struct{}
}

// this is a test utility method to force all messages to get processed
func (rm *ResponseManager) synchronize() {
	sync := make(chan struct{})
	select {
	case rm.messages <- &synchronizeMessage{sync}:
	case <-rm.ctx.Done():
	}
	select {
	case <-sync:
	case <-rm.ctx.Done():
	}
}

type responseDataRequest struct {
	key          responseKey
	taskDataChan chan *responseTaskData
}

type finishResponseRequest struct {
	key responseKey
}

func (rm *ResponseManager) processQueriesWorker() {
163
	const targetWork = 1
164 165 166
	taskDataChan := make(chan *responseTaskData)
	var taskData *responseTaskData
	for {
167 168
		pid, tasks, _ := rm.queryQueue.PopTasks(targetWork)
		for len(tasks) == 0 {
169 170 171 172
			select {
			case <-rm.ctx.Done():
				return
			case <-rm.workSignal:
173
				pid, tasks, _ = rm.queryQueue.PopTasks(targetWork)
174 175
			case <-rm.ticker.C:
				rm.queryQueue.ThawRound()
176
				pid, tasks, _ = rm.queryQueue.PopTasks(targetWork)
177 178
			}
		}
179 180
		for _, task := range tasks {
			key := task.Topic.(responseKey)
181 182 183 184 185 186 187 188 189 190
			select {
			case rm.messages <- &responseDataRequest{key, taskDataChan}:
			case <-rm.ctx.Done():
				return
			}
			select {
			case taskData = <-taskDataChan:
			case <-rm.ctx.Done():
				return
			}
191
			rm.executeQuery(taskData.ctx, key.p, taskData.request)
192 193 194 195 196
			select {
			case rm.messages <- &finishResponseRequest{key}:
			case <-rm.ctx.Done():
			}
		}
197
		rm.queryQueue.TasksDone(pid, tasks...)
198

199 200 201 202
	}

}

203
func noopVisitor(tp traversal.Progress, n ipld.Node, tr traversal.VisitReason) error {
204 205 206
	return nil
}

207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226
type hookActions struct {
	isValidated        bool
	requestID          graphsync.RequestID
	peerResponseSender peerresponsemanager.PeerResponseSender
	err                error
}

func (ha *hookActions) SendExtensionData(ext graphsync.ExtensionData) {
	ha.peerResponseSender.SendExtensionData(ha.requestID, ext)
}

func (ha *hookActions) TerminateWithError(err error) {
	ha.err = err
	ha.peerResponseSender.FinishWithError(ha.requestID, graphsync.RequestFailedUnknown)
}

func (ha *hookActions) ValidateRequest() {
	ha.isValidated = true
}

227 228
func (rm *ResponseManager) executeQuery(ctx context.Context,
	p peer.ID,
229
	request gsmsg.GraphSyncRequest) {
230
	peerResponseSender := rm.peerManager.SenderForPeer(p)
231
	selectorSpec := request.Selector()
232
	ha := &hookActions{false, request.ID(), peerResponseSender, nil}
233
	rm.requestHooksLk.RLock()
234 235 236
	for _, requestHook := range rm.requestHooks {
		requestHook.hook(p, request, ha)
		if ha.err != nil {
237
			rm.requestHooksLk.RUnlock()
238 239 240
			return
		}
	}
241
	rm.requestHooksLk.RUnlock()
242
	if !ha.isValidated {
243 244
		peerResponseSender.FinishWithError(request.ID(), graphsync.RequestFailedUnknown)
		return
245
	}
246
	selector, err := ipldutil.ParseSelector(selectorSpec)
247 248 249 250
	if err != nil {
		peerResponseSender.FinishWithError(request.ID(), graphsync.RequestFailedUnknown)
		return
	}
251
	rootLink := cidlink.Link{Cid: request.Root()}
252
	wrappedLoader := loader.WrapLoader(ctx, rm.loader, request.ID(), peerResponseSender)
253
	err = ipldutil.Traverse(ctx, wrappedLoader, rootLink, selector, noopVisitor)
254 255 256 257 258 259 260
	if err != nil {
		peerResponseSender.FinishWithError(request.ID(), graphsync.RequestFailedUnknown)
		return
	}
	peerResponseSender.FinishRequest(request.ID())
}

261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301
// Startup starts processing for the WantManager.
func (rm *ResponseManager) Startup() {
	go rm.run()
}

// Shutdown ends processing for the want manager.
func (rm *ResponseManager) Shutdown() {
	rm.cancelFn()
}

func (rm *ResponseManager) cleanupInProcessResponses() {
	for _, response := range rm.inProgressResponses {
		response.cancelFn()
	}
}

func (rm *ResponseManager) run() {
	defer rm.cleanupInProcessResponses()
	for i := 0; i < maxInProcessRequests; i++ {
		go rm.processQueriesWorker()
	}

	for {
		select {
		case <-rm.ctx.Done():
			return
		case message := <-rm.messages:
			message.handle(rm)
		}
	}
}

func (prm *processRequestMessage) handle(rm *ResponseManager) {
	for _, request := range prm.requests {
		key := responseKey{p: prm.p, requestID: request.ID()}
		if !request.IsCancel() {
			ctx, cancelFn := context.WithCancel(rm.ctx)
			rm.inProgressResponses[key] =
				inProgressResponseStatus{
					ctx:      ctx,
					cancelFn: cancelFn,
302
					request:  request,
303
				}
304 305
			// TODO: Use a better work estimation metric.
			rm.queryQueue.PushTasks(prm.p, peertask.Task{Topic: key, Priority: int(request.Priority()), Work: 1})
306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323
			select {
			case rm.workSignal <- struct{}{}:
			default:
			}
		} else {
			rm.queryQueue.Remove(key, key.p)
			response, ok := rm.inProgressResponses[key]
			if ok {
				response.cancelFn()
			}
		}
	}
}

func (rdr *responseDataRequest) handle(rm *ResponseManager) {
	response, ok := rm.inProgressResponses[rdr.key]
	var taskData *responseTaskData
	if ok {
324
		taskData = &responseTaskData{response.ctx, response.request}
325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348
	} else {
		taskData = nil
	}
	select {
	case <-rm.ctx.Done():
	case rdr.taskDataChan <- taskData:
	}
}

func (frr *finishResponseRequest) handle(rm *ResponseManager) {
	response, ok := rm.inProgressResponses[frr.key]
	if !ok {
		return
	}
	delete(rm.inProgressResponses, frr.key)
	response.cancelFn()
}

func (sm *synchronizeMessage) handle(rm *ResponseManager) {
	select {
	case <-rm.ctx.Done():
	case sm.sync <- struct{}{}:
	}
}