kraken

1
// Copyright (c) 2016-2019 Uber Technologies, Inc.
2
//
3
// Licensed under the Apache License, Version 2.0 (the "License");
4
// you may not use this file except in compliance with the License.
5
// You may obtain a copy of the License at
6
//
7
//     http://www.apache.org/licenses/LICENSE-2.0
8
//
9
// Unless required by applicable law or agreed to in writing, software
10
// distributed under the License is distributed on an "AS IS" BASIS,
11
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
// See the License for the specific language governing permissions and
13
// limitations under the License.
14
package networkevent
15

16
import (
17
	"encoding/json"
18
	"time"
19

20
	"github.com/uber/kraken/core"
21
	"github.com/uber/kraken/utils/log"
22

23
	"github.com/willf/bitset"
24
)
25

26
// Name defines event names.
27
type Name string
28

29
// Possible event names.
30
const (
31
	AddTorrent       Name = "add_torrent"
32
	AddActiveConn    Name = "add_active_conn"
33
	DropActiveConn   Name = "drop_active_conn"
34
	BlacklistConn    Name = "blacklist_conn"
35
	RequestPiece     Name = "request_piece"
36
	ReceivePiece     Name = "receive_piece"
37
	TorrentComplete  Name = "torrent_complete"
38
	TorrentCancelled Name = "torrent_cancelled"
39
)
40

41
// Event consolidates all possible event fields.
42
type Event struct {
43
	Name    Name      `json:"event"`
44
	Torrent string    `json:"torrent"`
45
	Self    string    `json:"self"`
46
	Time    time.Time `json:"ts"`
47

48
	// Optional fields.
49
	Peer         string `json:"peer,omitempty"`
50
	Piece        int    `json:"piece,omitempty"`
51
	Bitfield     []bool `json:"bitfield,omitempty"`
52
	DurationMS   int64  `json:"duration_ms,omitempty"`
53
	ConnCapacity int    `json:"conn_capacity,omitempty"`
54
}
55

56
func baseEvent(name Name, h core.InfoHash, self core.PeerID) *Event {
57
	return &Event{
58
		Name:    name,
59
		Torrent: h.String(),
60
		Self:    self.String(),
61
		Time:    time.Now(),
62
	}
63
}
64

65
// JSON converts event into a json string primarely for logging purposes
66
func (e *Event) JSON() string {
67
	b, err := json.Marshal(e)
68
	if err != nil {
69
		log.Errorf("json marshal error %s", err)
70
		return ""
71
	}
72
	return string(b)
73
}
74

75
// AddTorrentEvent returns an event for an added torrent with initial bitfield.
76
func AddTorrentEvent(h core.InfoHash, self core.PeerID, b *bitset.BitSet, connCapacity int) *Event {
77
	e := baseEvent(AddTorrent, h, self)
78
	bools := make([]bool, b.Len())
79
	for i := uint(0); i < b.Len(); i++ {
80
		bools[i] = b.Test(i)
81
	}
82
	e.Bitfield = bools
83
	e.ConnCapacity = connCapacity
84
	return e
85
}
86

87
// AddActiveConnEvent returns an event for an added active conn from self to peer.
88
func AddActiveConnEvent(h core.InfoHash, self core.PeerID, peer core.PeerID) *Event {
89
	e := baseEvent(AddActiveConn, h, self)
90
	e.Peer = peer.String()
91
	return e
92
}
93

94
// DropActiveConnEvent returns an event for a dropped active conn from self to peer.
95
func DropActiveConnEvent(h core.InfoHash, self core.PeerID, peer core.PeerID) *Event {
96
	e := baseEvent(DropActiveConn, h, self)
97
	e.Peer = peer.String()
98
	return e
99
}
100

101
// BlacklistConnEvent returns an event for a blacklisted connection.
102
func BlacklistConnEvent(h core.InfoHash, self core.PeerID, peer core.PeerID, dur time.Duration) *Event {
103
	e := baseEvent(BlacklistConn, h, self)
104
	e.Peer = peer.String()
105
	e.DurationMS = int64(dur.Seconds() * 1000)
106
	return e
107
}
108

109
// RequestPieceEvent returns an event for a piece request sent to a peer.
110
func RequestPieceEvent(h core.InfoHash, self core.PeerID, peer core.PeerID, piece int) *Event {
111
	e := baseEvent(RequestPiece, h, self)
112
	e.Peer = peer.String()
113
	e.Piece = piece
114
	return e
115
}
116

117
// ReceivePieceEvent returns an event for a piece received from a peer.
118
func ReceivePieceEvent(h core.InfoHash, self core.PeerID, peer core.PeerID, piece int) *Event {
119
	e := baseEvent(ReceivePiece, h, self)
120
	e.Peer = peer.String()
121
	e.Piece = piece
122
	return e
123
}
124

125
// TorrentCompleteEvent returns an event for a completed torrent.
126
func TorrentCompleteEvent(h core.InfoHash, self core.PeerID) *Event {
127
	return baseEvent(TorrentComplete, h, self)
128
}
129

130
// TorrentCancelledEvent returns an event for a cancelled torrent.
131
func TorrentCancelledEvent(h core.InfoHash, self core.PeerID) *Event {
132
	return baseEvent(TorrentCancelled, h, self)
133
}
134