1
package transport
2
 
3
import (
4
	"context"
5
	"sync"
6
 
7
	"go.miragespace.co/specter/spec/protocol"
8
 
9
	"github.com/zhangyunhao116/skipmap"
10
	"go.uber.org/zap"
11
)
12
 
13
type StreamHandler func(delegate *StreamDelegate)
14
 
15
type StreamRouter struct {
16
	logger                *zap.Logger
17
	virtualChordHandlers  *skipmap.Int32Map[*sync.Map] // map[protocol.Stream_Type]StreamHandler
18
	physicalChordHandlers sync.Map                     // map[protocol.Stream_Type]StreamHandler
19
	tunnelHandlers        sync.Map                     // map[protocol.Stream_Type]StreamHandler
20
	chordStream           <-chan *StreamDelegate
21
	tunnelStream          <-chan *StreamDelegate
22
}
23
 
24
func NewStreamRouter(logger *zap.Logger, chordTransport, tunnelTransport Transport) *StreamRouter {
25
	router := &StreamRouter{
26
		logger:               logger,
27
		virtualChordHandlers: skipmap.NewInt32[*sync.Map](),
28
	}
29
	if chordTransport != nil {
30
		router.chordStream = chordTransport.AcceptStream()
31
	}
32
	if tunnelTransport != nil {
33
		router.tunnelStream = tunnelTransport.AcceptStream()
34
	}
35
	return router
36
}
37
 
38
func (s *StreamRouter) HandleChord(kind protocol.Stream_Type, target *protocol.Node, handler StreamHandler) {
39
	if target == nil {
40
		// physical node handler
41
		s.physicalChordHandlers.Store(kind, handler)
42
	} else {
43
		// virtual node handler
44
		m, _ := s.virtualChordHandlers.LoadOrStoreLazy(int32(kind), func() *sync.Map {
45
			return &sync.Map{}
46
		})
47
		m.Store(target.GetId(), handler)
48
	}
49
}
50
 
51
func (s *StreamRouter) HandleTunnel(kind protocol.Stream_Type, handler StreamHandler) {
52
	s.tunnelHandlers.Store(kind, handler)
53
}
54
 
55
func (s *StreamRouter) acceptChord(ctx context.Context) {
56
	var (
57
		m       *sync.Map
58
		handler any
59
		ok      bool
60
	)
61
	for {
62
		select {
63
		case <-ctx.Done():
64
			return
65
		case delegate := <-s.chordStream:
66
			m, ok = s.virtualChordHandlers.Load(int32(delegate.Kind))
67
			if ok {
68
				// prioritize specific virtual node handler
69
				handler, ok = m.Load(delegate.Identity.GetId())
70
				if !ok {
71
					// fallback to root handler
72
					handler, ok = s.physicalChordHandlers.Load(delegate.Kind)
73
				}
74
			} else {
75
				// otherwise physical node handler
76
				handler, ok = s.physicalChordHandlers.Load(delegate.Kind)
77
			}
78
			if !ok {
79
				s.logger.Warn("No handler found for chord transport delegate",
80
					zap.Object("peer", delegate.Identity),
81
					zap.String("kind", delegate.Kind.String()),
82
				)
83
				delegate.Close()
84
				continue
85
			}
86
			go handler.(StreamHandler)(delegate)
87
		}
88
	}
89
}
90
 
91
func (s *StreamRouter) acceptTunnel(ctx context.Context) {
92
	for {
93
		select {
94
		case <-ctx.Done():
95
			return
96
		case delegate := <-s.tunnelStream:
97
			handler, ok := s.tunnelHandlers.Load(delegate.Kind)
98
			if !ok {
99
				s.logger.Warn("No handler found for tunnel transport delegate",
100
					zap.Object("peer", delegate.Identity),
101
					zap.String("kind", delegate.Kind.String()),
102
				)
103
				delegate.Close()
104
				continue
105
			}
106
			go handler.(StreamHandler)(delegate)
107
		}
108
	}
109
}
110
 
111
func (s *StreamRouter) Accept(ctx context.Context) {
112
	if s.chordStream != nil {
113
		go s.acceptChord(ctx)
114
		go s.acceptChord(ctx)
115
	}
116
	if s.tunnelStream != nil {
117
		go s.acceptTunnel(ctx)
118
		go s.acceptTunnel(ctx)
119
	}
120
}