api, metrics: add number of bytes received and sent from/to all entities (#1235)

* API: number of bytes received/sent from/to RTSP connections
* API: number of bytes received/sent from/to RTSP sessions
* API: number of bytes received/sent from/to RTMP connections
* API: number of bytes sent to HLS connections
* API: number of bytes received from paths
* metrics of all the above
This commit is contained in:
Alessandro Ros
2022-11-11 11:59:52 +01:00
committed by GitHub
parent 71ef9b47ab
commit 8bee4af86a
25 changed files with 338 additions and 277 deletions

View File

@@ -433,7 +433,7 @@ Full documentation of the API is available on the [dedicated site](https://aler9
### Metrics ### Metrics
A metrics exporter, compatible with Prometheus, can be enabled with the parameter `metrics: yes`; then the server can be queried for metrics with Prometheus or with a simple HTTP request: A metrics exporter, compatible with [Prometheus](https://prometheus.io/), can be enabled with the parameter `metrics: yes`; then the server can be queried for metrics with Prometheus or with a simple HTTP request:
``` ```
wget -qO- localhost:9998/metrics wget -qO- localhost:9998/metrics
@@ -441,34 +441,40 @@ wget -qO- localhost:9998/metrics
Obtaining: Obtaining:
``` ```ini
paths{name="<path_name>",state="ready"} 1 # metrics of every path
rtsp_conns 1 paths{name="[path_name]",state="[state]"} 1
rtsp_sessions{state="idle"} 0 paths_bytes_received{name="[path_name]",state="[state]"} 1234
rtsp_sessions{state="read"} 0
rtsp_sessions{state="publish"} 1
rtsps_sessions{state="idle"} 0
rtsps_sessions{state="read"} 0
rtsps_sessions{state="publish"} 0
rtmp_conns{state="idle"} 0
rtmp_conns{state="read"} 0
rtmp_conns{state="publish"} 1
hls_muxers{name="<name>"} 1
```
where: # metrics of every RTSP connection
rtsp_conns{id="[id]"} 1
rtsp_conns_bytes_received{id="[id]"} 1234
rtsp_conns_bytes_sent{id="[id]"} 187
* `paths{name="<path_name>",state="ready"} 1` is replicated for every path and shows the name and state of every path # metrics of every RTSP session
* `rtsp_sessions{state="idle"}` is the count of RTSP sessions that are idle rtsp_sessions{id="[id]",state="idle"} 1
* `rtsp_sessions{state="read"}` is the count of RTSP sessions that are reading rtsp_sessions_bytes_received{id="[id]",state="[state]"} 1234
* `rtsp_sessions{state="publish"}` is the counf ot RTSP sessions that are publishing rtsp_sessions_bytes_sent{id="[id]",state="[state]"} 187
* `rtsps_sessions{state="idle"}` is the count of RTSPS sessions that are idle
* `rtsps_sessions{state="read"}` is the count of RTSPS sessions that are reading # metrics of every RTSPS connection
* `rtsps_sessions{state="publish"}` is the counf ot RTSPS sessions that are publishing rtsps_conns{id="[id]"} 1
* `rtmp_conns{state="idle"}` is the count of RTMP connections that are idle rtsps_conns_bytes_received{id="[id]"} 1234
* `rtmp_conns{state="read"}` is the count of RTMP connections that are reading rtsps_conns_bytes_sent{id="[id]"} 187
* `rtmp_conns{state="publish"}` is the count of RTMP connections that are publishing
* `hls_muxers{name="<name>"}` is replicated for every HLS muxer and shows the name and state of every HLS muxer # metrics of every RTSPS session
rtsps_sessions{id="[id]",state="[state]"} 1
rtsps_sessions_bytes_received{id="[id]",state="[state]"} 1234
rtsps_sessions_bytes_sent{id="[id]",state="[state]"} 187
# metrics of every RTMP connection
rtmp_conns{id="[id]",state="[state]"} 1
rtmp_conns_bytes_received{id="[id]",state="[state]"} 1234
rtmp_conns_bytes_sent{id="[id]",state="[state]"} 187
# metrics of every HLS muxer
hls_muxers{name="[name]"} 1
hls_muxers_bytes_sent{name="[name]"} 187
```
### pprof ### pprof

View File

@@ -268,6 +268,8 @@ components:
- PCMU - PCMU
- VP8 - VP8
- VP9 - VP9
bytesReceived:
type: number
readers: readers:
type: array type: array
items: items:
@@ -382,6 +384,10 @@ components:
type: string type: string
remoteAddr: remoteAddr:
type: string type: string
bytesReceived:
type: number
bytesSent:
type: number
RTSPSession: RTSPSession:
type: object type: object
@@ -393,6 +399,10 @@ components:
state: state:
type: string type: string
enum: [idle, read, publish] enum: [idle, read, publish]
bytesReceived:
type: number
bytesSent:
type: number
RTMPConn: RTMPConn:
type: object type: object
@@ -404,17 +414,10 @@ components:
state: state:
type: string type: string
enum: [idle, read, publish] enum: [idle, read, publish]
bytesReceived:
RTMPSConn: type: number
type: object bytesSent:
properties: type: number
created:
type: string
remoteAddr:
type: string
state:
type: string
enum: [idle, read, publish]
HLSMuxer: HLSMuxer:
type: object type: object
@@ -423,6 +426,8 @@ components:
type: string type: string
lastRequest: lastRequest:
type: string type: string
bytesSent:
type: number
PathsList: PathsList:
type: object type: object
@@ -464,14 +469,6 @@ components:
additionalProperties: additionalProperties:
$ref: '#/components/schemas/RTMPConn' $ref: '#/components/schemas/RTMPConn'
RTMPSConnsList:
type: object
properties:
items:
type: object
additionalProperties:
$ref: '#/components/schemas/RTMPSConn'
HLSMuxersList: HLSMuxersList:
type: object type: object
properties: properties:
@@ -592,7 +589,7 @@ paths:
/v1/paths/list: /v1/paths/list:
get: get:
operationId: pathsList operationId: pathsList
summary: returns all active paths. summary: returns all paths.
description: '' description: ''
responses: responses:
'200': '200':
@@ -609,7 +606,7 @@ paths:
/v1/rtspconns/list: /v1/rtspconns/list:
get: get:
operationId: rtspConnsList operationId: rtspConnsList
summary: returns all active RTSP connections. summary: returns all RTSP connections.
description: '' description: ''
responses: responses:
'200': '200':
@@ -626,7 +623,7 @@ paths:
/v1/rtspsessions/list: /v1/rtspsessions/list:
get: get:
operationId: rtspSessionsList operationId: rtspSessionsList
summary: returns all active RTSP sessions. summary: returns all RTSP sessions.
description: '' description: ''
responses: responses:
'200': '200':
@@ -643,7 +640,7 @@ paths:
/v1/rtspsconns/list: /v1/rtspsconns/list:
get: get:
operationId: rtspsConnsList operationId: rtspsConnsList
summary: returns all active RTSPS connections. summary: returns all RTSPS connections.
description: '' description: ''
responses: responses:
'200': '200':
@@ -680,7 +677,7 @@ paths:
/v1/rtspssessions/list: /v1/rtspssessions/list:
get: get:
operationId: rtspsSessionsList operationId: rtspsSessionsList
summary: returns all active RTSPS sessions. summary: returns all RTSPS sessions.
description: '' description: ''
responses: responses:
'200': '200':
@@ -717,7 +714,7 @@ paths:
/v1/rtmpconns/list: /v1/rtmpconns/list:
get: get:
operationId: rtmpConnsList operationId: rtmpConnsList
summary: returns all active RTMP connections. summary: returns all RTMP connections.
description: '' description: ''
responses: responses:
'200': '200':
@@ -754,7 +751,7 @@ paths:
/v1/rtmpsconns/list: /v1/rtmpsconns/list:
get: get:
operationId: rtmpsConnsList operationId: rtmpsConnsList
summary: returns all active RTMPS connections. summary: returns all RTMPS connections.
description: '' description: ''
responses: responses:
'200': '200':
@@ -762,7 +759,7 @@ paths:
content: content:
application/json: application/json:
schema: schema:
$ref: '#/components/schemas/RTMPSConnsList' $ref: '#/components/schemas/RTMPConnsList'
'400': '400':
description: invalid request. description: invalid request.
'500': '500':
@@ -791,7 +788,7 @@ paths:
/v1/hlsmuxers/list: /v1/hlsmuxers/list:
get: get:
operationId: hlsMuxersList operationId: hlsMuxersList
summary: returns all active HLS muxers. summary: returns all HLS muxers.
description: '' description: ''
responses: responses:
'200': '200':

2
go.mod
View File

@@ -5,7 +5,7 @@ go 1.18
require ( require (
code.cloudfoundry.org/bytefmt v0.0.0-20211005130812-5bb3c17173e5 code.cloudfoundry.org/bytefmt v0.0.0-20211005130812-5bb3c17173e5
github.com/abema/go-mp4 v0.8.0 github.com/abema/go-mp4 v0.8.0
github.com/aler9/gortsplib v0.0.0-20221105162652-b1ed0a8abb48 github.com/aler9/gortsplib v0.0.0-20221110211534-12c8845fef0d
github.com/asticode/go-astits v1.10.1-0.20220319093903-4abe66a9b757 github.com/asticode/go-astits v1.10.1-0.20220319093903-4abe66a9b757
github.com/fsnotify/fsnotify v1.4.9 github.com/fsnotify/fsnotify v1.4.9
github.com/gin-gonic/gin v1.8.1 github.com/gin-gonic/gin v1.8.1

4
go.sum
View File

@@ -6,8 +6,8 @@ github.com/alecthomas/template v0.0.0-20190718012654-fb15b899a751 h1:JYp7IbQjafo
github.com/alecthomas/template v0.0.0-20190718012654-fb15b899a751/go.mod h1:LOuyumcjzFXgccqObfd/Ljyb9UuFJ6TxHnclSeseNhc= github.com/alecthomas/template v0.0.0-20190718012654-fb15b899a751/go.mod h1:LOuyumcjzFXgccqObfd/Ljyb9UuFJ6TxHnclSeseNhc=
github.com/alecthomas/units v0.0.0-20190924025748-f65c72e2690d h1:UQZhZ2O0vMHr2cI+DC1Mbh0TJxzA3RcLoMsFw+aXw7E= github.com/alecthomas/units v0.0.0-20190924025748-f65c72e2690d h1:UQZhZ2O0vMHr2cI+DC1Mbh0TJxzA3RcLoMsFw+aXw7E=
github.com/alecthomas/units v0.0.0-20190924025748-f65c72e2690d/go.mod h1:rBZYJk541a8SKzHPHnH3zbiI+7dagKZ0cgpgrD7Fyho= github.com/alecthomas/units v0.0.0-20190924025748-f65c72e2690d/go.mod h1:rBZYJk541a8SKzHPHnH3zbiI+7dagKZ0cgpgrD7Fyho=
github.com/aler9/gortsplib v0.0.0-20221105162652-b1ed0a8abb48 h1:SVcUqrsR+RXT8cxopOahKQaj8kb02adaEeRexLLKcMc= github.com/aler9/gortsplib v0.0.0-20221110211534-12c8845fef0d h1:fRx79L1YMXaoiSMkB32xgVCUMbOcmQ4JfySaUv7XZpc=
github.com/aler9/gortsplib v0.0.0-20221105162652-b1ed0a8abb48/go.mod h1:BOWNZ/QBkY/eVcRqUzJbPFEsRJshwxaxBT01K260Jeo= github.com/aler9/gortsplib v0.0.0-20221110211534-12c8845fef0d/go.mod h1:BOWNZ/QBkY/eVcRqUzJbPFEsRJshwxaxBT01K260Jeo=
github.com/aler9/writerseeker v0.0.0-20220601075008-6f0e685b9c82 h1:9WgSzBLo3a9ToSVV7sRTBYZ1GGOZUpq4+5H3SN0UZq4= github.com/aler9/writerseeker v0.0.0-20220601075008-6f0e685b9c82 h1:9WgSzBLo3a9ToSVV7sRTBYZ1GGOZUpq4+5H3SN0UZq4=
github.com/aler9/writerseeker v0.0.0-20220601075008-6f0e685b9c82/go.mod h1:qsMrZCbeBf/mCLOeF16KDkPu4gktn/pOWyaq1aYQE7U= github.com/aler9/writerseeker v0.0.0-20220601075008-6f0e685b9c82/go.mod h1:qsMrZCbeBf/mCLOeF16KDkPu4gktn/pOWyaq1aYQE7U=
github.com/asticode/go-astikit v0.20.0 h1:+7N+J4E4lWx2QOkRdOf6DafWJMv6O4RRfgClwQokrH8= github.com/asticode/go-astikit v0.20.0 h1:+7N+J4E4lWx2QOkRdOf6DafWJMv6O4RRfgClwQokrH8=

View File

@@ -15,6 +15,7 @@ import (
"github.com/aler9/gortsplib" "github.com/aler9/gortsplib"
"github.com/aler9/gortsplib/pkg/mpeg4audio" "github.com/aler9/gortsplib/pkg/mpeg4audio"
"github.com/pion/rtp"
"github.com/stretchr/testify/require" "github.com/stretchr/testify/require"
"github.com/aler9/rtsp-simple-server/internal/rtmp" "github.com/aler9/rtsp-simple-server/internal/rtmp"
@@ -170,9 +171,10 @@ func TestAPIPathsList(t *testing.T) {
} }
type path struct { type path struct {
Source pathSource `json:"source"` Source pathSource `json:"source"`
SourceReady bool `json:"sourceReady"` SourceReady bool `json:"sourceReady"`
Tracks []string `json:"tracks"` Tracks []string `json:"tracks"`
BytesReceived uint64 `json:"bytesReceived"`
} }
type pathList struct { type pathList struct {
@@ -186,30 +188,38 @@ func TestAPIPathsList(t *testing.T) {
require.Equal(t, true, ok) require.Equal(t, true, ok)
defer p.Close() defer p.Close()
tracks := gortsplib.Tracks{
&gortsplib.TrackH264{
PayloadType: 96,
SPS: []byte{0x01, 0x02, 0x03, 0x04},
PPS: []byte{0x01, 0x02, 0x03, 0x04},
},
&gortsplib.TrackMPEG4Audio{
PayloadType: 97,
Config: &mpeg4audio.Config{
Type: 2,
SampleRate: 44100,
ChannelCount: 2,
},
SizeLength: 13,
IndexLength: 3,
IndexDeltaLength: 3,
},
}
source := gortsplib.Client{} source := gortsplib.Client{}
err := source.StartPublishing("rtsp://localhost:8554/mypath", tracks) err := source.StartPublishing(
"rtsp://localhost:8554/mypath",
gortsplib.Tracks{
&gortsplib.TrackH264{
PayloadType: 96,
SPS: []byte{0x01, 0x02, 0x03, 0x04},
PPS: []byte{0x01, 0x02, 0x03, 0x04},
},
&gortsplib.TrackMPEG4Audio{
PayloadType: 96,
Config: &mpeg4audio.Config{
Type: 2,
SampleRate: 44100,
ChannelCount: 2,
},
SizeLength: 13,
IndexLength: 3,
IndexDeltaLength: 3,
},
})
require.NoError(t, err) require.NoError(t, err)
defer source.Close() defer source.Close()
source.WritePacketRTP(0, &rtp.Packet{
Header: rtp.Header{
Version: 2,
PayloadType: 96,
},
Payload: []byte{0x01, 0x02, 0x03, 0x04},
}, true)
var out pathList var out pathList
err = httpRequest(http.MethodGet, "http://localhost:9997/v1/paths/list", nil, &out) err = httpRequest(http.MethodGet, "http://localhost:9997/v1/paths/list", nil, &out)
require.NoError(t, err) require.NoError(t, err)
@@ -219,8 +229,9 @@ func TestAPIPathsList(t *testing.T) {
Source: pathSource{ Source: pathSource{
Type: "rtspSession", Type: "rtspSession",
}, },
SourceReady: true, SourceReady: true,
Tracks: []string{"H264", "MPEG4Audio"}, Tracks: []string{"H264", "MPEG4Audio"},
BytesReceived: 16,
}, },
}, },
}, out) }, out)

View File

@@ -94,11 +94,16 @@ window.addEventListener('DOMContentLoaded', create);
</html> </html>
` `
type hlsMuxerResponse struct {
muxer *hlsMuxer
cb func() *hls.MuxerFileResponse
}
type hlsMuxerRequest struct { type hlsMuxerRequest struct {
dir string dir string
file string file string
ctx *gin.Context ctx *gin.Context
res chan func() *hls.MuxerFileResponse res chan hlsMuxerResponse
} }
type hlsMuxerPathManager interface { type hlsMuxerPathManager interface {
@@ -133,6 +138,7 @@ type hlsMuxer struct {
lastRequestTime *int64 lastRequestTime *int64
muxer *hls.Muxer muxer *hls.Muxer
requests []*hlsMuxerRequest requests []*hlsMuxerRequest
bytesSent *uint64
// in // in
chRequest chan *hlsMuxerRequest chRequest chan *hlsMuxerRequest
@@ -179,6 +185,7 @@ func newHLSMuxer(
v := time.Now().UnixNano() v := time.Now().UnixNano()
return &v return &v
}(), }(),
bytesSent: new(uint64),
chRequest: make(chan *hlsMuxerRequest), chRequest: make(chan *hlsMuxerRequest),
chAPIHLSMuxersList: make(chan hlsServerAPIMuxersListSubReq), chAPIHLSMuxersList: make(chan hlsServerAPIMuxersListSubReq),
} }
@@ -235,7 +242,10 @@ func (m *hlsMuxer) run() {
case req := <-m.chRequest: case req := <-m.chRequest:
if isReady { if isReady {
req.res <- m.handleRequest(req) req.res <- hlsMuxerResponse{
muxer: m,
cb: m.handleRequest(req),
}
} else { } else {
m.requests = append(m.requests, req) m.requests = append(m.requests, req)
} }
@@ -244,13 +254,17 @@ func (m *hlsMuxer) run() {
req.data.Items[m.name] = hlsServerAPIMuxersListItem{ req.data.Items[m.name] = hlsServerAPIMuxersListItem{
Created: m.created, Created: m.created,
LastRequest: time.Unix(0, atomic.LoadInt64(m.lastRequestTime)), LastRequest: time.Unix(0, atomic.LoadInt64(m.lastRequestTime)),
BytesSent: atomic.LoadUint64(m.bytesSent),
} }
close(req.res) close(req.res)
case <-innerReady: case <-innerReady:
isReady = true isReady = true
for _, req := range m.requests { for _, req := range m.requests {
req.res <- m.handleRequest(req) req.res <- hlsMuxerResponse{
muxer: m,
cb: m.handleRequest(req),
}
} }
m.requests = nil m.requests = nil
@@ -264,8 +278,11 @@ func (m *hlsMuxer) run() {
m.ctxCancel() m.ctxCancel()
for _, req := range m.requests { for _, req := range m.requests {
req.res <- func() *hls.MuxerFileResponse { req.res <- hlsMuxerResponse{
return &hls.MuxerFileResponse{Status: http.StatusNotFound} muxer: m,
cb: func() *hls.MuxerFileResponse {
return &hls.MuxerFileResponse{Status: http.StatusNotFound}
},
} }
} }
@@ -547,13 +564,20 @@ func (m *hlsMuxer) authenticate(ctx *gin.Context) error {
return nil return nil
} }
func (m *hlsMuxer) addSentBytes(n uint64) {
atomic.AddUint64(m.bytesSent, n)
}
// request is called by hlsserver.Server (forwarded from ServeHTTP). // request is called by hlsserver.Server (forwarded from ServeHTTP).
func (m *hlsMuxer) request(req *hlsMuxerRequest) { func (m *hlsMuxer) request(req *hlsMuxerRequest) {
select { select {
case m.chRequest <- req: case m.chRequest <- req:
case <-m.ctx.Done(): case <-m.ctx.Done():
req.res <- func() *hls.MuxerFileResponse { req.res <- hlsMuxerResponse{
return &hls.MuxerFileResponse{Status: http.StatusInternalServerError} muxer: m,
cb: func() *hls.MuxerFileResponse {
return &hls.MuxerFileResponse{Status: http.StatusInternalServerError}
},
} }
} }
} }

View File

@@ -17,7 +17,6 @@ import (
"github.com/gin-gonic/gin" "github.com/gin-gonic/gin"
"github.com/aler9/rtsp-simple-server/internal/conf" "github.com/aler9/rtsp-simple-server/internal/conf"
"github.com/aler9/rtsp-simple-server/internal/hls"
"github.com/aler9/rtsp-simple-server/internal/logger" "github.com/aler9/rtsp-simple-server/internal/logger"
) )
@@ -30,6 +29,7 @@ func (nilWriter) Write(p []byte) (int, error) {
type hlsServerAPIMuxersListItem struct { type hlsServerAPIMuxersListItem struct {
Created time.Time `json:"created"` Created time.Time `json:"created"`
LastRequest time.Time `json:"lastRequest"` LastRequest time.Time `json:"lastRequest"`
BytesSent uint64 `json:"bytesSent"`
} }
type hlsServerAPIMuxersListData struct { type hlsServerAPIMuxersListData struct {
@@ -311,19 +311,17 @@ func (s *hlsServer) onRequest(ctx *gin.Context) {
dir = strings.TrimSuffix(dir, "/") dir = strings.TrimSuffix(dir, "/")
cres := make(chan func() *hls.MuxerFileResponse)
hreq := &hlsMuxerRequest{ hreq := &hlsMuxerRequest{
dir: dir, dir: dir,
file: fname, file: fname,
ctx: ctx, ctx: ctx,
res: cres, res: make(chan hlsMuxerResponse),
} }
select { select {
case s.request <- hreq: case s.request <- hreq:
cb := <-cres res1 := <-hreq.res
res := res1.cb()
res := cb()
for k, v := range res.Header { for k, v := range res.Header {
ctx.Writer.Header().Set(k, v) ctx.Writer.Header().Set(k, v)
@@ -332,7 +330,8 @@ func (s *hlsServer) onRequest(ctx *gin.Context) {
ctx.Writer.WriteHeader(res.Status) ctx.Writer.WriteHeader(res.Status)
if res.Body != nil { if res.Body != nil {
io.Copy(ctx.Writer, res.Body) n, _ := io.Copy(ctx.Writer, res.Body)
res1.muxer.addSentBytes(uint64(n))
} }
case <-s.ctx.Done(): case <-s.ctx.Done():

View File

@@ -93,83 +93,68 @@ func (m *metrics) onMetrics(ctx *gin.Context) {
res := m.pathManager.apiPathsList() res := m.pathManager.apiPathsList()
if res.err == nil { if res.err == nil {
for name, p := range res.data.Items { for name, i := range res.data.Items {
if p.SourceReady { var state string
out += metric("paths{name=\""+name+"\",state=\"ready\"}", 1) if i.SourceReady {
state = "ready"
} else { } else {
out += metric("paths{name=\""+name+"\",state=\"notReady\"}", 1) state = "notReady"
} }
tags := "{name=\"" + name + "\",state=\"" + state + "\"}"
out += metric("paths"+tags, 1)
out += metric("paths_bytes_received"+tags, int64(i.BytesReceived))
} }
} }
if !interfaceIsEmpty(m.rtspServer) { if !interfaceIsEmpty(m.rtspServer) { //nolint:dupl
func() { func() {
res := m.rtspServer.apiConnsList() res := m.rtspServer.apiConnsList()
if res.err == nil { if res.err == nil {
out += metric("rtsp_conns", int64(len(res.data.Items))) for id, i := range res.data.Items {
tags := "{id=\"" + id + "\"}"
out += metric("rtsp_conns"+tags, 1)
out += metric("rtsp_conns_bytes_received"+tags, int64(i.BytesReceived))
out += metric("rtsp_conns_bytes_sent"+tags, int64(i.BytesSent))
}
} }
}() }()
func() { func() {
res := m.rtspServer.apiSessionsList() res := m.rtspServer.apiSessionsList()
if res.err == nil { if res.err == nil {
idleCount := int64(0) for id, i := range res.data.Items {
readCount := int64(0) tags := "{id=\"" + id + "\",state=\"" + i.State + "\"}"
publishCount := int64(0) out += metric("rtsp_sessions"+tags, 1)
out += metric("rtsp_sessions_bytes_received"+tags, int64(i.BytesReceived))
for _, i := range res.data.Items { out += metric("rtsp_sessions_bytes_sent"+tags, int64(i.BytesSent))
switch i.State {
case "idle":
idleCount++
case "read":
readCount++
case "publish":
publishCount++
}
} }
out += metric("rtsp_sessions{state=\"idle\"}",
idleCount)
out += metric("rtsp_sessions{state=\"read\"}",
readCount)
out += metric("rtsp_sessions{state=\"publish\"}",
publishCount)
} }
}() }()
} }
if !interfaceIsEmpty(m.rtspsServer) { if !interfaceIsEmpty(m.rtspsServer) { //nolint:dupl
func() { func() {
res := m.rtspsServer.apiConnsList() res := m.rtspsServer.apiConnsList()
if res.err == nil { if res.err == nil {
out += metric("rtsps_conns", int64(len(res.data.Items))) for id, i := range res.data.Items {
tags := "{id=\"" + id + "\"}"
out += metric("rtsps_conns"+tags, 1)
out += metric("rtsps_conns_bytes_received"+tags, int64(i.BytesReceived))
out += metric("rtsps_conns_bytes_sent"+tags, int64(i.BytesSent))
}
} }
}() }()
func() { func() {
res := m.rtspsServer.apiSessionsList() res := m.rtspsServer.apiSessionsList()
if res.err == nil { if res.err == nil {
idleCount := int64(0) for id, i := range res.data.Items {
readCount := int64(0) tags := "{id=\"" + id + "\",state=\"" + i.State + "\"}"
publishCount := int64(0) out += metric("rtsps_sessions"+tags, 1)
out += metric("rtsps_sessions_bytes_received"+tags, int64(i.BytesReceived))
for _, i := range res.data.Items { out += metric("rtsps_sessions_bytes_sent"+tags, int64(i.BytesSent))
switch i.State {
case "idle":
idleCount++
case "read":
readCount++
case "publish":
publishCount++
}
} }
out += metric("rtsps_sessions{state=\"idle\"}",
idleCount)
out += metric("rtsps_sessions{state=\"read\"}",
readCount)
out += metric("rtsps_sessions{state=\"publish\"}",
publishCount)
} }
}() }()
} }
@@ -177,35 +162,22 @@ func (m *metrics) onMetrics(ctx *gin.Context) {
if !interfaceIsEmpty(m.rtmpServer) { if !interfaceIsEmpty(m.rtmpServer) {
res := m.rtmpServer.apiConnsList() res := m.rtmpServer.apiConnsList()
if res.err == nil { if res.err == nil {
idleCount := int64(0) for id, i := range res.data.Items {
readCount := int64(0) tags := "{id=\"" + id + "\",state=\"" + i.State + "\"}"
publishCount := int64(0) out += metric("rtmp_conns"+tags, 1)
out += metric("rtmp_conns_bytes_received"+tags, int64(i.BytesReceived))
for _, i := range res.data.Items { out += metric("rtmp_conns_bytes_sent"+tags, int64(i.BytesSent))
switch i.State {
case "idle":
idleCount++
case "read":
readCount++
case "publish":
publishCount++
}
} }
out += metric("rtmp_conns{state=\"idle\"}",
idleCount)
out += metric("rtmp_conns{state=\"read\"}",
readCount)
out += metric("rtmp_conns{state=\"publish\"}",
publishCount)
} }
} }
if !interfaceIsEmpty(m.hlsServer) { if !interfaceIsEmpty(m.hlsServer) {
res := m.hlsServer.apiHLSMuxersList() res := m.hlsServer.apiHLSMuxersList()
if res.err == nil { if res.err == nil {
for name := range res.data.Items { for name, i := range res.data.Items {
out += metric("hls_muxers{name=\""+name+"\"}", 1) tags := "{name=\"" + name + "\"}"
out += metric("hls_muxers"+tags, 1)
out += metric("hls_muxers_bytes_sent"+tags, int64(i.BytesSent))
} }
} }
} }

View File

@@ -1,12 +1,12 @@
package core package core
import ( import (
"crypto/tls"
"io" "io"
"net" "net"
"net/http" "net/http"
"net/url" "net/url"
"os" "os"
"strings"
"testing" "testing"
"github.com/aler9/gortsplib" "github.com/aler9/gortsplib"
@@ -40,12 +40,17 @@ func TestMetrics(t *testing.T) {
} }
source := gortsplib.Client{} source := gortsplib.Client{}
err = source.StartPublishing("rtsp://localhost:8554/rtsp_path", err = source.StartPublishing("rtsp://localhost:8554/rtsp_path",
gortsplib.Tracks{track}) gortsplib.Tracks{track})
require.NoError(t, err) require.NoError(t, err)
defer source.Close() defer source.Close()
source2 := gortsplib.Client{TLSConfig: &tls.Config{InsecureSkipVerify: true}}
err = source2.StartPublishing("rtsps://localhost:8322/rtsps_path",
gortsplib.Tracks{track})
require.NoError(t, err)
defer source2.Close()
u, err := url.Parse("rtmp://localhost:1935/rtmp_path") u, err := url.Parse("rtmp://localhost:1935/rtmp_path")
require.NoError(t, err) require.NoError(t, err)
@@ -88,27 +93,29 @@ func TestMetrics(t *testing.T) {
bo, err := io.ReadAll(res.Body) bo, err := io.ReadAll(res.Body)
require.NoError(t, err) require.NoError(t, err)
vals := make(map[string]string) require.Regexp(t,
lines := strings.Split(string(bo), "\n") `^paths\{name=".*?",state="ready"\} 1`+"\n"+
for _, l := range lines[:len(lines)-1] { `paths_bytes_received\{name=".*?",state="ready"\} 0`+"\n"+
fields := strings.Split(l, " ") `paths\{name=".*?",state="ready"\} 1`+"\n"+
vals[fields[0]] = fields[1] `paths_bytes_received\{name=".*?",state="ready"\} 0`+"\n"+
} `paths\{name=".*?",state="ready"\} 1`+"\n"+
`paths_bytes_received\{name=".*?",state="ready"\} 0`+"\n"+
require.Equal(t, map[string]string{ `rtsp_conns\{id=".*?"\} 1`+"\n"+
"hls_muxers{name=\"rtsp_path\"}": "1", `rtsp_conns_bytes_received\{id=".*?"\} [0-9]+`+"\n"+
"paths{name=\"rtsp_path\",state=\"ready\"}": "1", `rtsp_conns_bytes_sent\{id=".*?"\} [0-9]+`+"\n"+
"paths{name=\"rtmp_path\",state=\"ready\"}": "1", `rtsp_sessions\{id=".*?",state="publish"\} 1`+"\n"+
"rtmp_conns{state=\"idle\"}": "0", `rtsp_sessions_bytes_received\{id=".*?",state="publish"\} 0`+"\n"+
"rtmp_conns{state=\"publish\"}": "1", `rtsp_sessions_bytes_sent\{id=".*?",state="publish"\} [0-9]+`+"\n"+
"rtmp_conns{state=\"read\"}": "0", `rtsps_conns\{id=".*?"\} 1`+"\n"+
"rtsp_conns": "1", `rtsps_conns_bytes_received\{id=".*?"\} [0-9]+`+"\n"+
"rtsp_sessions{state=\"idle\"}": "0", `rtsps_conns_bytes_sent\{id=".*?"\} [0-9]+`+"\n"+
"rtsp_sessions{state=\"publish\"}": "1", `rtsps_sessions\{id=".*?",state="publish"\} 1`+"\n"+
"rtsp_sessions{state=\"read\"}": "0", `rtsps_sessions_bytes_received\{id=".*?",state="publish"\} 0`+"\n"+
"rtsps_conns": "0", `rtsps_sessions_bytes_sent\{id=".*?",state="publish"\} [0-9]+`+"\n"+
"rtsps_sessions{state=\"idle\"}": "0", `rtmp_conns\{id=".*?",state="publish"\} 1`+"\n"+
"rtsps_sessions{state=\"publish\"}": "0", `rtmp_conns_bytes_received\{id=".*?",state="publish"\} [0-9]+`+"\n"+
"rtsps_sessions{state=\"read\"}": "0", `rtmp_conns_bytes_sent\{id=".*?",state="publish"\} [0-9]+`+"\n"+
}, vals) `hls_muxers\{name="rtsp_path"\} 1`+"\n"+
`hls_muxers_bytes_sent\{name="rtsp_path"\} [0-9]+`+"\n"+"$",
string(bo))
} }

View File

@@ -7,6 +7,7 @@ import (
"strconv" "strconv"
"strings" "strings"
"sync" "sync"
"sync/atomic"
"time" "time"
"github.com/aler9/gortsplib" "github.com/aler9/gortsplib"
@@ -178,12 +179,13 @@ type pathPublisherStopReq struct {
} }
type pathAPIPathsListItem struct { type pathAPIPathsListItem struct {
ConfName string `json:"confName"` ConfName string `json:"confName"`
Conf *conf.PathConf `json:"conf"` Conf *conf.PathConf `json:"conf"`
Source interface{} `json:"source"` Source interface{} `json:"source"`
SourceReady bool `json:"sourceReady"` SourceReady bool `json:"sourceReady"`
Tracks []string `json:"tracks"` Tracks []string `json:"tracks"`
Readers []interface{} `json:"readers"` BytesReceived uint64 `json:"bytesReceived"`
Readers []interface{} `json:"readers"`
} }
type pathAPIPathsListData struct { type pathAPIPathsListData struct {
@@ -221,6 +223,7 @@ type path struct {
ctx context.Context ctx context.Context
ctxCancel func() ctxCancel func()
source source source source
bytesReceived *uint64
stream *stream stream *stream
readers map[reader]pathReaderState readers map[reader]pathReaderState
describeRequestsOnHold []pathDescribeReq describeRequestsOnHold []pathDescribeReq
@@ -279,6 +282,7 @@ func newPath(
parent: parent, parent: parent,
ctx: ctx, ctx: ctx,
ctxCancel: ctxCancel, ctxCancel: ctxCancel,
bytesReceived: new(uint64),
readers: make(map[reader]pathReaderState), readers: make(map[reader]pathReaderState),
onDemandStaticSourceReadyTimer: newEmptyTimer(), onDemandStaticSourceReadyTimer: newEmptyTimer(),
onDemandStaticSourceCloseTimer: newEmptyTimer(), onDemandStaticSourceCloseTimer: newEmptyTimer(),
@@ -663,8 +667,8 @@ func (pa *path) onDemandPublisherStop() {
} }
} }
func (pa *path) sourceSetReady(tracks gortsplib.Tracks, generateRTPPackets bool) error { func (pa *path) sourceSetReady(tracks gortsplib.Tracks, allocateEncoder bool) error {
stream, err := newStream(tracks, generateRTPPackets) stream, err := newStream(tracks, allocateEncoder, pa.bytesReceived)
if err != nil { if err != nil {
return err return err
} }
@@ -957,6 +961,7 @@ func (pa *path) handleAPIPathsList(req pathAPIPathsListSubReq) {
} }
return sourceTrackNames(pa.stream.tracks()) return sourceTrackNames(pa.stream.tracks())
}(), }(),
BytesReceived: atomic.LoadUint64(pa.bytesReceived),
Readers: func() []interface{} { Readers: func() []interface{} {
ret := []interface{}{} ret := []interface{}{}
for r := range pa.readers { for r := range pa.readers {

View File

@@ -14,9 +14,11 @@ import (
) )
type rtmpServerAPIConnsListItem struct { type rtmpServerAPIConnsListItem struct {
Created time.Time `json:"created"` Created time.Time `json:"created"`
RemoteAddr string `json:"remoteAddr"` RemoteAddr string `json:"remoteAddr"`
State string `json:"state"` State string `json:"state"`
BytesReceived uint64 `json:"bytesReceived"`
BytesSent uint64 `json:"bytesSent"`
} }
type rtmpServerAPIConnsListData struct { type rtmpServerAPIConnsListData struct {
@@ -236,6 +238,8 @@ outer:
} }
return "idle" return "idle"
}(), }(),
BytesReceived: c.conn.BytesReceived(),
BytesSent: c.conn.BytesSent(),
} }
} }

View File

@@ -19,8 +19,10 @@ import (
) )
type rtspServerAPIConnsListItem struct { type rtspServerAPIConnsListItem struct {
Created time.Time `json:"created"` Created time.Time `json:"created"`
RemoteAddr string `json:"remoteAddr"` RemoteAddr string `json:"remoteAddr"`
BytesReceived uint64 `json:"bytesReceived"`
BytesSent uint64 `json:"bytesSent"`
} }
type rtspServerAPIConnsListData struct { type rtspServerAPIConnsListData struct {
@@ -33,9 +35,11 @@ type rtspServerAPIConnsListRes struct {
} }
type rtspServerAPISessionsListItem struct { type rtspServerAPISessionsListItem struct {
Created time.Time `json:"created"` Created time.Time `json:"created"`
RemoteAddr string `json:"remoteAddr"` RemoteAddr string `json:"remoteAddr"`
State string `json:"state"` State string `json:"state"`
BytesReceived uint64 `json:"bytesReceived"`
BytesSent uint64 `json:"bytesSent"`
} }
type rtspServerAPISessionsListData struct { type rtspServerAPISessionsListData struct {
@@ -374,8 +378,10 @@ func (s *rtspServer) apiConnsList() rtspServerAPIConnsListRes {
for _, c := range s.conns { for _, c := range s.conns {
data.Items[c.uuid.String()] = rtspServerAPIConnsListItem{ data.Items[c.uuid.String()] = rtspServerAPIConnsListItem{
Created: c.created, Created: c.created,
RemoteAddr: c.remoteAddr().String(), RemoteAddr: c.remoteAddr().String(),
BytesReceived: c.conn.BytesReceived(),
BytesSent: c.conn.BytesSent(),
} }
} }
@@ -413,6 +419,8 @@ func (s *rtspServer) apiSessionsList() rtspServerAPISessionsListRes {
} }
return "idle" return "idle"
}(), }(),
BytesReceived: s.session.BytesReceived(),
BytesSent: s.session.BytesSent(),
} }
} }

View File

@@ -34,7 +34,7 @@ type rtspSessionParent interface {
type rtspSession struct { type rtspSession struct {
isTLS bool isTLS bool
protocols map[conf.Protocol]struct{} protocols map[conf.Protocol]struct{}
ss *gortsplib.ServerSession session *gortsplib.ServerSession
author *gortsplib.ServerConn author *gortsplib.ServerConn
externalCmdPool *externalcmd.Pool externalCmdPool *externalcmd.Pool
pathManager rtspSessionPathManager pathManager rtspSessionPathManager
@@ -52,7 +52,7 @@ type rtspSession struct {
func newRTSPSession( func newRTSPSession(
isTLS bool, isTLS bool,
protocols map[conf.Protocol]struct{}, protocols map[conf.Protocol]struct{},
ss *gortsplib.ServerSession, session *gortsplib.ServerSession,
sc *gortsplib.ServerConn, sc *gortsplib.ServerConn,
externalCmdPool *externalcmd.Pool, externalCmdPool *externalcmd.Pool,
pathManager rtspSessionPathManager, pathManager rtspSessionPathManager,
@@ -61,7 +61,7 @@ func newRTSPSession(
s := &rtspSession{ s := &rtspSession{
isTLS: isTLS, isTLS: isTLS,
protocols: protocols, protocols: protocols,
ss: ss, session: session,
author: sc, author: sc,
externalCmdPool: externalCmdPool, externalCmdPool: externalCmdPool,
pathManager: pathManager, pathManager: pathManager,
@@ -77,7 +77,7 @@ func newRTSPSession(
// Close closes a Session. // Close closes a Session.
func (s *rtspSession) close() { func (s *rtspSession) close() {
s.ss.Close() s.session.Close()
} }
// isRTSPSession implements pathRTSPSession. // isRTSPSession implements pathRTSPSession.
@@ -100,7 +100,7 @@ func (s *rtspSession) log(level logger.Level, format string, args ...interface{}
// onClose is called by rtspServer. // onClose is called by rtspServer.
func (s *rtspSession) onClose(err error) { func (s *rtspSession) onClose(err error) {
if s.ss.State() == gortsplib.ServerSessionStatePlay { if s.session.State() == gortsplib.ServerSessionStatePlay {
if s.onReadCmd != nil { if s.onReadCmd != nil {
s.onReadCmd.Close() s.onReadCmd.Close()
s.onReadCmd = nil s.onReadCmd = nil
@@ -108,7 +108,7 @@ func (s *rtspSession) onClose(err error) {
} }
} }
switch s.ss.State() { switch s.session.State() {
case gortsplib.ServerSessionStatePrePlay, gortsplib.ServerSessionStatePlay: case gortsplib.ServerSessionStatePrePlay, gortsplib.ServerSessionStatePlay:
s.path.readerRemove(pathReaderRemoveReq{author: s}) s.path.readerRemove(pathReaderRemoveReq{author: s})
@@ -181,7 +181,7 @@ func (s *rtspSession) onSetup(c *rtspConn, ctx *gortsplib.ServerHandlerOnSetupCt
} }
} }
switch s.ss.State() { switch s.session.State() {
case gortsplib.ServerSessionStateInitial, gortsplib.ServerSessionStatePrePlay: // play case gortsplib.ServerSessionStateInitial, gortsplib.ServerSessionStatePrePlay: // play
res := s.pathManager.readerAdd(pathReaderAddReq{ res := s.pathManager.readerAdd(pathReaderAddReq{
author: s, author: s,
@@ -247,19 +247,19 @@ func (s *rtspSession) onSetup(c *rtspConn, ctx *gortsplib.ServerHandlerOnSetupCt
func (s *rtspSession) onPlay(ctx *gortsplib.ServerHandlerOnPlayCtx) (*base.Response, error) { func (s *rtspSession) onPlay(ctx *gortsplib.ServerHandlerOnPlayCtx) (*base.Response, error) {
h := make(base.Header) h := make(base.Header)
if s.ss.State() == gortsplib.ServerSessionStatePrePlay { if s.session.State() == gortsplib.ServerSessionStatePrePlay {
s.path.readerStart(pathReaderStartReq{author: s}) s.path.readerStart(pathReaderStartReq{author: s})
tracks := make(gortsplib.Tracks, len(s.ss.SetuppedTracks())) tracks := make(gortsplib.Tracks, len(s.session.SetuppedTracks()))
n := 0 n := 0
for id := range s.ss.SetuppedTracks() { for id := range s.session.SetuppedTracks() {
tracks[n] = s.stream.tracks()[id] tracks[n] = s.stream.tracks()[id]
n++ n++
} }
s.log(logger.Info, "is reading from path '%s', with %s, %s", s.log(logger.Info, "is reading from path '%s', with %s, %s",
s.path.Name(), s.path.Name(),
s.ss.SetuppedTransport(), s.session.SetuppedTransport(),
sourceTrackInfo(tracks)) sourceTrackInfo(tracks))
if s.path.Conf().RunOnRead != "" { if s.path.Conf().RunOnRead != "" {
@@ -289,7 +289,7 @@ func (s *rtspSession) onPlay(ctx *gortsplib.ServerHandlerOnPlayCtx) (*base.Respo
func (s *rtspSession) onRecord(ctx *gortsplib.ServerHandlerOnRecordCtx) (*base.Response, error) { func (s *rtspSession) onRecord(ctx *gortsplib.ServerHandlerOnRecordCtx) (*base.Response, error) {
res := s.path.publisherStart(pathPublisherStartReq{ res := s.path.publisherStart(pathPublisherStartReq{
author: s, author: s,
tracks: s.ss.AnnouncedTracks(), tracks: s.session.AnnouncedTracks(),
generateRTPPackets: false, generateRTPPackets: false,
}) })
if res.err != nil { if res.err != nil {
@@ -300,8 +300,8 @@ func (s *rtspSession) onRecord(ctx *gortsplib.ServerHandlerOnRecordCtx) (*base.R
s.log(logger.Info, "is publishing to path '%s', with %s, %s", s.log(logger.Info, "is publishing to path '%s', with %s, %s",
s.path.Name(), s.path.Name(),
s.ss.SetuppedTransport(), s.session.SetuppedTransport(),
sourceTrackInfo(s.ss.AnnouncedTracks())) sourceTrackInfo(s.session.AnnouncedTracks()))
s.stream = res.stream s.stream = res.stream
@@ -316,7 +316,7 @@ func (s *rtspSession) onRecord(ctx *gortsplib.ServerHandlerOnRecordCtx) (*base.R
// onPause is called by rtspServer. // onPause is called by rtspServer.
func (s *rtspSession) onPause(ctx *gortsplib.ServerHandlerOnPauseCtx) (*base.Response, error) { func (s *rtspSession) onPause(ctx *gortsplib.ServerHandlerOnPauseCtx) (*base.Response, error) {
switch s.ss.State() { switch s.session.State() {
case gortsplib.ServerSessionStatePlay: case gortsplib.ServerSessionStatePlay:
if s.onReadCmd != nil { if s.onReadCmd != nil {
s.log(logger.Info, "runOnRead command stopped") s.log(logger.Info, "runOnRead command stopped")
@@ -381,7 +381,7 @@ func (s *rtspSession) apiSourceDescribe() interface{} {
func (s *rtspSession) onPacketRTP(ctx *gortsplib.ServerHandlerOnPacketRTPCtx) { func (s *rtspSession) onPacketRTP(ctx *gortsplib.ServerHandlerOnPacketRTPCtx) {
var err error var err error
switch s.ss.AnnouncedTracks()[ctx.TrackID].(type) { switch s.session.AnnouncedTracks()[ctx.TrackID].(type) {
case *gortsplib.TrackH264: case *gortsplib.TrackH264:
err = s.stream.writeData(&dataH264{ err = s.stream.writeData(&dataH264{
trackID: ctx.TrackID, trackID: ctx.TrackID,

View File

@@ -2,6 +2,7 @@ package core
import ( import (
"sync" "sync"
"sync/atomic"
"github.com/aler9/gortsplib" "github.com/aler9/gortsplib"
) )
@@ -51,13 +52,19 @@ func (m *streamNonRTSPReadersMap) hasReaders() bool {
} }
type stream struct { type stream struct {
bytesReceived *uint64
nonRTSPReaders *streamNonRTSPReadersMap nonRTSPReaders *streamNonRTSPReadersMap
rtspStream *gortsplib.ServerStream rtspStream *gortsplib.ServerStream
streamTracks []streamTrack streamTracks []streamTrack
} }
func newStream(tracks gortsplib.Tracks, generateRTPPackets bool) (*stream, error) { func newStream(
tracks gortsplib.Tracks,
generateRTPPackets bool,
bytesReceived *uint64,
) (*stream, error) {
s := &stream{ s := &stream{
bytesReceived: bytesReceived,
nonRTSPReaders: newStreamNonRTSPReadersMap(), nonRTSPReaders: newStreamNonRTSPReadersMap(),
rtspStream: gortsplib.NewServerStream(tracks), rtspStream: gortsplib.NewServerStream(tracks),
} }
@@ -104,6 +111,7 @@ func (s *stream) writeData(data data) error {
// forward RTP packets to RTSP readers // forward RTP packets to RTSP readers
for _, pkt := range data.getRTPPackets() { for _, pkt := range data.getRTPPackets() {
atomic.AddUint64(s.bytesReceived, uint64(pkt.MarshalSize()))
s.rtspStream.WritePacketRTP(data.getTrackID(), pkt, data.getPTSEqualsDTS()) s.rtspStream.WritePacketRTP(data.getTrackID(), pkt, data.getPTSEqualsDTS())
} }

View File

@@ -70,13 +70,13 @@ type streamTrackH264 struct {
func newStreamTrackH264( func newStreamTrackH264(
track *gortsplib.TrackH264, track *gortsplib.TrackH264,
generateRTPPackets bool, allocateEncoder bool,
) *streamTrackH264 { ) *streamTrackH264 {
t := &streamTrackH264{ t := &streamTrackH264{
track: track, track: track,
} }
if generateRTPPackets { if allocateEncoder {
t.encoder = &rtph264.Encoder{PayloadType: 96} t.encoder = &rtph264.Encoder{PayloadType: 96}
t.encoder.Init() t.encoder.Init()
} }

View File

@@ -15,13 +15,13 @@ type streamTrackMPEG4Audio struct {
func newStreamTrackMPEG4Audio( func newStreamTrackMPEG4Audio(
track *gortsplib.TrackMPEG4Audio, track *gortsplib.TrackMPEG4Audio,
generateRTPPackets bool, allocateEncoder bool,
) *streamTrackMPEG4Audio { ) *streamTrackMPEG4Audio {
t := &streamTrackMPEG4Audio{ t := &streamTrackMPEG4Audio{
track: track, track: track,
} }
if generateRTPPackets { if allocateEncoder {
t.encoder = &rtpmpeg4audio.Encoder{ t.encoder = &rtpmpeg4audio.Encoder{
PayloadType: 96, PayloadType: 96,
SampleRate: track.ClockRate(), SampleRate: track.ClockRate(),

View File

@@ -1,42 +1,36 @@
package bytecounter package bytecounter
import ( import (
"bufio"
"io" "io"
"sync/atomic"
) )
type readerInner struct {
r io.Reader
count uint32
}
func (r *readerInner) Read(p []byte) (int, error) {
n, err := r.r.Read(p)
r.count += uint32(n)
return n, err
}
// Reader allows to count read bytes. // Reader allows to count read bytes.
type Reader struct { type Reader struct {
ri *readerInner r io.Reader
*bufio.Reader count uint64
} }
// NewReader allocates a Reader. // NewReader allocates a Reader.
func NewReader(r io.Reader) *Reader { func NewReader(r io.Reader) *Reader {
ri := &readerInner{r: r}
return &Reader{ return &Reader{
ri: ri, r: r,
Reader: bufio.NewReader(ri),
} }
} }
// Count returns read bytes. // Read implements io.Reader.
func (r Reader) Count() uint32 { func (r *Reader) Read(p []byte) (int, error) {
return r.ri.count n, err := r.r.Read(p)
atomic.AddUint64(&r.count, uint64(n))
return n, err
}
// Count returns received bytes.
func (r *Reader) Count() uint64 {
return atomic.LoadUint64(&r.count)
} }
// SetCount sets read bytes. // SetCount sets read bytes.
func (r *Reader) SetCount(v uint32) { func (r *Reader) SetCount(v uint64) {
r.ri.count = v atomic.StoreUint64(&r.count, v)
} }

View File

@@ -19,5 +19,5 @@ func TestReader(t *testing.T) {
require.NoError(t, err) require.NoError(t, err)
require.Equal(t, 64, n) require.Equal(t, 64, n)
require.Equal(t, uint32(100+1024), r.Count()) require.Equal(t, uint64(100+64), r.Count())
} }

View File

@@ -2,12 +2,13 @@ package bytecounter
import ( import (
"io" "io"
"sync/atomic"
) )
// Writer allows to count written bytes. // Writer allows to count written bytes.
type Writer struct { type Writer struct {
w io.Writer w io.Writer
count uint32 count uint64
} }
// NewWriter allocates a Writer. // NewWriter allocates a Writer.
@@ -20,16 +21,16 @@ func NewWriter(w io.Writer) *Writer {
// Write implements io.Writer. // Write implements io.Writer.
func (w *Writer) Write(p []byte) (int, error) { func (w *Writer) Write(p []byte) (int, error) {
n, err := w.w.Write(p) n, err := w.w.Write(p)
w.count += uint32(n) atomic.AddUint64(&w.count, uint64(n))
return n, err return n, err
} }
// Count returns written bytes. // Count returns sent bytes.
func (w Writer) Count() uint32 { func (w *Writer) Count() uint64 {
return w.count return atomic.LoadUint64(&w.count)
} }
// SetCount sets written bytes. // SetCount sets sent bytes.
func (w *Writer) SetCount(v uint32) { func (w *Writer) SetCount(v uint64) {
w.count = v atomic.StoreUint64(&w.count, v)
} }

View File

@@ -14,5 +14,5 @@ func TestWriter(t *testing.T) {
w.SetCount(100) w.SetCount(100)
w.Write(bytes.Repeat([]byte{0x01}, 64)) w.Write(bytes.Repeat([]byte{0x01}, 64))
require.Equal(t, uint32(100+64), w.Count()) require.Equal(t, uint64(100+64), w.Count())
} }

View File

@@ -114,10 +114,19 @@ type Conn struct {
// NewConn initializes a connection. // NewConn initializes a connection.
func NewConn(rw io.ReadWriter) *Conn { func NewConn(rw io.ReadWriter) *Conn {
c := &Conn{} return &Conn{
c.bc = bytecounter.NewReadWriter(rw) bc: bytecounter.NewReadWriter(rw),
c.mrw = message.NewReadWriter(c.bc, false) }
return c }
// BytesReceived returns the number of bytes received.
func (c *Conn) BytesReceived() uint64 {
return c.bc.Reader.Count()
}
// BytesSent returns the number of bytes sent.
func (c *Conn) BytesSent() uint64 {
return c.bc.Writer.Count()
} }
func (c *Conn) readCommand() (*message.MsgCommandAMF0, error) { func (c *Conn) readCommand() (*message.MsgCommandAMF0, error) {
@@ -161,6 +170,8 @@ func (c *Conn) InitializeClient(u *url.URL, isPublishing bool) error {
return err return err
} }
c.mrw = message.NewReadWriter(c.bc, false)
err = c.mrw.Write(&message.MsgSetWindowAckSize{ err = c.mrw.Write(&message.MsgSetWindowAckSize{
Value: 2500000, Value: 2500000,
}) })
@@ -319,6 +330,8 @@ func (c *Conn) InitializeServer() (*url.URL, bool, error) {
return nil, false, err return nil, false, err
} }
c.mrw = message.NewReadWriter(c.bc, false)
cmd, err := c.readCommand() cmd, err := c.readCommand()
if err != nil { if err != nil {
return nil, false, err return nil, false, err

View File

@@ -244,6 +244,14 @@ func TestInitializeClient(t *testing.T) {
err = conn.InitializeClient(u, ca == "publish") err = conn.InitializeClient(u, ca == "publish")
require.NoError(t, err) require.NoError(t, err)
if ca == "read" {
require.Equal(t, uint64(3421), conn.BytesReceived())
require.Equal(t, uint64(3409), conn.BytesSent())
} else {
require.Equal(t, uint64(3427), conn.BytesReceived())
require.Equal(t, uint64(3466), conn.BytesSent())
}
<-done <-done
}) })
} }

View File

@@ -1,6 +1,7 @@
package rawmessage package rawmessage
import ( import (
"bufio"
"errors" "errors"
"fmt" "fmt"
"time" "time"
@@ -22,14 +23,14 @@ type readerChunkStream struct {
} }
func (rc *readerChunkStream) readChunk(c chunk.Chunk, chunkBodySize uint32) error { func (rc *readerChunkStream) readChunk(c chunk.Chunk, chunkBodySize uint32) error {
err := c.Read(rc.mr.r, chunkBodySize) err := c.Read(rc.mr.br, chunkBodySize)
if err != nil { if err != nil {
return err return err
} }
// check if an ack is needed // check if an ack is needed
if rc.mr.ackWindowSize != 0 { if rc.mr.ackWindowSize != 0 {
count := rc.mr.r.Count() count := uint32(rc.mr.r.Count())
diff := count - rc.mr.lastAckCount diff := count - rc.mr.lastAckCount
if diff > (rc.mr.ackWindowSize) { if diff > (rc.mr.ackWindowSize) {
@@ -210,6 +211,7 @@ type Reader struct {
r *bytecounter.Reader r *bytecounter.Reader
onAckNeeded func(uint32) error onAckNeeded func(uint32) error
br *bufio.Reader
chunkSize uint32 chunkSize uint32
ackWindowSize uint32 ackWindowSize uint32
lastAckCount uint32 lastAckCount uint32
@@ -223,8 +225,11 @@ type Reader struct {
// NewReader allocates a Reader. // NewReader allocates a Reader.
func NewReader(r *bytecounter.Reader, onAckNeeded func(uint32) error) *Reader { func NewReader(r *bytecounter.Reader, onAckNeeded func(uint32) error) *Reader {
br := bufio.NewReader(r)
return &Reader{ return &Reader{
r: r, r: r,
br: br,
onAckNeeded: onAckNeeded, onAckNeeded: onAckNeeded,
chunkSize: 128, chunkSize: 128,
chunkStreams: make(map[byte]*readerChunkStream), chunkStreams: make(map[byte]*readerChunkStream),
@@ -244,7 +249,7 @@ func (r *Reader) SetWindowAckSize(v uint32) {
// Read reads a Message. // Read reads a Message.
func (r *Reader) Read() (*Message, error) { func (r *Reader) Read() (*Message, error) {
for { for {
byt, err := r.r.ReadByte() byt, err := r.br.ReadByte()
if err != nil { if err != nil {
return nil, err return nil, err
} }
@@ -258,7 +263,7 @@ func (r *Reader) Read() (*Message, error) {
r.chunkStreams[chunkStreamID] = rc r.chunkStreams[chunkStreamID] = rc
} }
r.r.UnreadByte() r.br.UnreadByte()
msg, err := rc.readMessage(typ) msg, err := rc.readMessage(typ)
if err != nil { if err != nil {

View File

@@ -198,8 +198,7 @@ func TestReader(t *testing.T) {
for _, ca := range cases { for _, ca := range cases {
t.Run(ca.name, func(t *testing.T) { t.Run(ca.name, func(t *testing.T) {
var buf bytes.Buffer var buf bytes.Buffer
bcr := bytecounter.NewReader(&buf) r := NewReader(bytecounter.NewReader(&buf), func(count uint32) error {
r := NewReader(bcr, func(count uint32) error {
return nil return nil
}) })
@@ -224,14 +223,14 @@ func TestReaderAcknowledge(t *testing.T) {
onAckCalled := make(chan struct{}) onAckCalled := make(chan struct{})
var buf bytes.Buffer var buf bytes.Buffer
bcr := bytecounter.NewReader(&buf) bc := bytecounter.NewReader(&buf)
r := NewReader(bcr, func(count uint32) error { r := NewReader(bc, func(count uint32) error {
close(onAckCalled) close(onAckCalled)
return nil return nil
}) })
if ca == "overflow" { if ca == "overflow" {
bcr.SetCount(4294967096) bc.SetCount(4294967096)
r.lastAckCount = 4294967096 r.lastAckCount = 4294967096
} }

View File

@@ -20,7 +20,7 @@ type writerChunkStream struct {
func (wc *writerChunkStream) writeChunk(c chunk.Chunk) error { func (wc *writerChunkStream) writeChunk(c chunk.Chunk) error {
// check if we received an acknowledge // check if we received an acknowledge
if wc.mw.checkAcknowledge && wc.mw.ackWindowSize != 0 { if wc.mw.checkAcknowledge && wc.mw.ackWindowSize != 0 {
diff := wc.mw.w.Count() - wc.mw.ackValue diff := uint32(wc.mw.w.Count()) - wc.mw.ackValue
if diff > (wc.mw.ackWindowSize * 3 / 2) { if diff > (wc.mw.ackWindowSize * 3 / 2) {
return fmt.Errorf("no acknowledge received within window") return fmt.Errorf("no acknowledge received within window")