wandb

1
package server
2

3
import (
4
	"fmt"
5

6
	"github.com/wandb/wandb/core/pkg/observability"
7
	"github.com/wandb/wandb/core/pkg/service"
8
)
9

10
type Responder interface {
11
	Respond(response *service.ServerResponse)
12
}
13

14
type ResponderEntry struct {
15
	Responder Responder
16
	ID        string
17
}
18

19
type Dispatcher struct {
20
	responders map[string]Responder
21
	logger     *observability.CoreLogger
22
}
23

24
// AddResponders adds the given responders to the stream's dispatcher.
25
func (d *Dispatcher) AddResponders(entries ...ResponderEntry) {
26
	if d.responders == nil {
27
		d.responders = make(map[string]Responder)
28
	}
29
	for _, entry := range entries {
30
		responderId := entry.ID
31
		if _, ok := d.responders[responderId]; !ok {
32
			d.responders[responderId] = entry.Responder
33
		} else {
34
			d.logger.CaptureWarn("Responder already exists", "responder", responderId)
35
		}
36
	}
37
}
38

39
func (d *Dispatcher) handleRespond(result *service.Result) {
40
	responderId := result.GetControl().GetConnectionId()
41
	d.logger.Debug("dispatch: got result", "result", result)
42
	if responderId == "" {
43
		d.logger.Debug("dispatch: got result with no connection id", "result", result)
44
		return
45
	}
46
	response := &service.ServerResponse{
47
		ServerResponseType: &service.ServerResponse_ResultCommunicate{
48
			ResultCommunicate: result,
49
		},
50
	}
51
	if responder, ok := d.responders[responderId]; ok {
52
		responder.Respond(response)
53
	} else {
54
		err := fmt.Errorf("dispatch: no responder found: %s", responderId)
55
		d.logger.CaptureFatalAndPanic("dispatch: no responder found", err)
56
	}
57
}
58

59
func NewDispatcher(logger *observability.CoreLogger) *Dispatcher {
60
	return &Dispatcher{
61
		logger:     logger,
62
		responders: make(map[string]Responder),
63
	}
64
}
65