paulmataruso/goflow2
1
0
Fork 0
mirror of https://github.com/openobserve/goflow2.git synced 2025-11-02 21:13:27 +00:00
Code Issues Packages Projects Releases Wiki Activity

Allow Flow Routines to be cancellable (#40)

Browse Source 
* Allow Flow Routines to be cancellable
This commit is contained in:
Mario Macias
2021-11-01 00:42:07 +01:00
committed by GitHub
parent 92043a6233
commit d1e1ace318
7 changed files with 259 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files

104
utils/utils.go
View File

@@ -6,6 +6,7 @@ import (
"io"
"net"
"strconv"
"sync/atomic"
"time"
reuseport "github.com/libp2p/go-reuseport"
@@ -99,6 +100,11 @@ func (cb *DefaultErrorCallback) Callback(name string, id int, start, end time.Ti
}
func UDPRoutine(name string, decodeFunc decoder.DecoderFunc, workers int, addr string, port int, sockReuse bool, logger Logger) error {
return UDPStoppableRoutine(make(chan struct{}), name, decodeFunc, workers, addr, port, sockReuse, logger)
}
// UDPStoppableRoutine runs a UDPRoutine that can be stopped by closing the stopCh passed as argument
func UDPStoppableRoutine(stopCh <-chan struct{}, name string, decodeFunc decoder.DecoderFunc, workers int, addr string, port int, sockReuse bool, logger Logger) error {
ecb := DefaultErrorCallback{
Logger: logger,
}
@@ -146,41 +152,71 @@ func UDPRoutine(name string, decodeFunc decoder.DecoderFunc, workers int, addr s
localIP = ""
}
for {
size, pktAddr, _ := udpconn.ReadFromUDP(payload)
payloadCut := make([]byte, size)
copy(payloadCut, payload[0:size])
type udpData struct {
size int
pktAddr *net.UDPAddr
}
baseMessage := BaseMessage{
Src: pktAddr.IP,
Port: pktAddr.Port,
Payload: payloadCut,
stopped := atomic.Value{}
stopped.Store(false)
udpDataCh := make(chan udpData)
go func() {
for {
u := udpData{}
u.size, u.pktAddr, _ = udpconn.ReadFromUDP(payload)
if stopped.Load() == false {
udpDataCh <- u
} else {
return
}
}
}()
for {
select {
case u := <-udpDataCh:
process(u.size, payload, u.pktAddr, processor, localIP, addrUDP, name)
case <-stopCh:
stopped.Store(true)
udpconn.Close()
close(udpDataCh)
return nil
}
processor.ProcessMessage(baseMessage)
MetricTrafficBytes.With(
prometheus.Labels{
"remote_ip": pktAddr.IP.String(),
"local_ip": localIP,
"local_port": strconv.Itoa(addrUDP.Port),
"type": name,
}).
Add(float64(size))
MetricTrafficPackets.With(
prometheus.Labels{
"remote_ip": pktAddr.IP.String(),
"local_ip": localIP,
"local_port": strconv.Itoa(addrUDP.Port),
"type": name,
}).
Inc()
MetricPacketSizeSum.With(
prometheus.Labels{
"remote_ip": pktAddr.IP.String(),
"local_ip": localIP,
"local_port": strconv.Itoa(addrUDP.Port),
"type": name,
}).
Observe(float64(size))
}
}
func process(size int, payload []byte, pktAddr *net.UDPAddr, processor decoder.Processor, localIP string, addrUDP net.UDPAddr, name string) {
payloadCut := make([]byte, size)
copy(payloadCut, payload[0:size])
baseMessage := BaseMessage{
Src: pktAddr.IP,
Port: pktAddr.Port,
Payload: payloadCut,
}
processor.ProcessMessage(baseMessage)
MetricTrafficBytes.With(
prometheus.Labels{
"remote_ip": pktAddr.IP.String(),
"local_ip": localIP,
"local_port": strconv.Itoa(addrUDP.Port),
"type": name,
}).
Add(float64(size))
MetricTrafficPackets.With(
prometheus.Labels{
"remote_ip": pktAddr.IP.String(),
"local_ip": localIP,
"local_port": strconv.Itoa(addrUDP.Port),
"type": name,
}).
Inc()
MetricPacketSizeSum.With(
prometheus.Labels{
"remote_ip": pktAddr.IP.String(),
"local_ip": localIP,
"local_port": strconv.Itoa(addrUDP.Port),
"type": name,
}).
Observe(float64(size))
}