Skip to content

Commit 66a17de

Browse files
Tianbo-QiuTianbo-Qiu
committed
Confine zmq socket to a single send loop to fix concurrent access
1 parent 3d299a3 commit 66a17de

3 files changed

Lines changed: 265 additions & 11 deletions

File tree

‎datastore/zmq/zmq.go‎

Lines changed: 50 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -57,23 +57,37 @@ var (
5757
// ZMQ publisher socket.
5858
const MonitorSocketAddr = "inproc://zmq_socket_monitor.rep"
5959

60+
// recordChanBufferSize bounds how many records can wait for the send loop
61+
// before Produce applies backpressure to its caller.
62+
const recordChanBufferSize = 1000
63+
6064
// Producer implements the telemetry.Producer interface by publishing to a
6165
// bound zmq socket.
6266
type Producer struct {
6367
namespace string
6468
ctx context.Context
69+
cancel context.CancelFunc
6570
sock *zmq4.Socket
6671
logger *logrus.Logger
6772
airbrakeHandler *airbrake.Handler
6873
ackChan chan (*telemetry.Record)
6974
reliableAckTxTypes map[string]interface{}
75+
recordChan chan *telemetry.Record
76+
wg sync.WaitGroup
77+
closeOnce sync.Once
7078
}
7179

7280
// Produce the record to the socket.
7381
func (p *Producer) Produce(rec *telemetry.Record) {
74-
if p.ctx.Err() != nil {
75-
return
82+
select {
83+
case p.recordChan <- rec:
84+
case <-p.ctx.Done():
7685
}
86+
}
87+
88+
// send writes a single record to the socket.
89+
// It must only ever be called from sendLoop, which owns the socket.
90+
func (p *Producer) send(rec *telemetry.Record) {
7791
nBytes, err := p.sock.SendMessage(telemetry.BuildTopicName(p.namespace, rec.TxType), rec.Payload())
7892
if err != nil {
7993
metricsRegistry.errorCount.Inc(map[string]string{"record_type": rec.TxType})
@@ -85,21 +99,37 @@ func (p *Producer) Produce(rec *telemetry.Record) {
8599
metricsRegistry.publishCount.Inc(map[string]string{"record_type": rec.TxType})
86100
}
87101

102+
// sendLoop is the single goroutine that owns and writes to the socket.
103+
func (p *Producer) sendLoop() {
104+
defer p.wg.Done()
105+
for {
106+
select {
107+
case <-p.ctx.Done():
108+
return
109+
case rec := <-p.recordChan:
110+
p.send(rec)
111+
}
112+
}
113+
}
114+
88115
// ReportError to airbrake and logger
89116
func (p *Producer) ReportError(message string, err error, logInfo logrus.LogInfo) {
90117
p.airbrakeHandler.ReportLogMessage(logrus.ERROR, message, err, logInfo)
91118
p.logger.ErrorLog(message, err, logInfo)
92119
}
93120

94-
// Close the underlying socket.
121+
// Close stops the send loop and then closes the underlying socket.
95122
func (p *Producer) Close() error {
96-
if p.sock != nil {
97-
if err := p.sock.Close(); err != nil {
98-
return err
123+
var err error
124+
p.closeOnce.Do(func() {
125+
p.cancel()
126+
p.wg.Wait()
127+
if p.sock != nil {
128+
err = p.sock.Close()
129+
p.sock = nil
99130
}
100-
}
101-
p.sock = nil
102-
return nil
131+
})
132+
return err
103133
}
104134

105135
// ProcessReliableAck sends to ackChan if reliable ack is configured
@@ -175,15 +205,24 @@ func NewProducer(ctx context.Context, config *Config, metrics metrics.MetricColl
175205
return
176206
}
177207

178-
return &Producer{
208+
ctx, cancel := context.WithCancel(ctx)
209+
210+
p := &Producer{
179211
namespace: namespace,
180212
ctx: ctx,
213+
cancel: cancel,
181214
sock: sock,
182215
logger: logger,
183216
airbrakeHandler: airbrakeHandler,
184217
ackChan: ackChan,
185218
reliableAckTxTypes: reliableAckTxTypes,
186-
}, nil
219+
recordChan: make(chan *telemetry.Record, recordChanBufferSize),
220+
}
221+
222+
p.wg.Add(1)
223+
go p.sendLoop()
224+
225+
return p, nil
187226
}
188227

189228
// logSocketInBackground logs the socket activity in the background.

‎datastore/zmq/zmq_suite_test.go‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
package zmq_test
2+
3+
import (
4+
"testing"
5+
6+
. "github.com/onsi/ginkgo/v2"
7+
. "github.com/onsi/gomega"
8+
)
9+
10+
func TestZMQ(t *testing.T) {
11+
RegisterFailHandler(Fail)
12+
RunSpecs(t, "ZMQ Suite")
13+
}

‎datastore/zmq/zmq_test.go‎

Lines changed: 202 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,202 @@
1+
package zmq_test
2+
3+
import (
4+
"context"
5+
"sync"
6+
"time"
7+
8+
. "github.com/onsi/ginkgo/v2"
9+
. "github.com/onsi/gomega"
10+
"github.com/pebbe/zmq4"
11+
"google.golang.org/protobuf/proto"
12+
13+
"github.com/teslamotors/fleet-telemetry/datastore/zmq"
14+
logrus "github.com/teslamotors/fleet-telemetry/logger"
15+
"github.com/teslamotors/fleet-telemetry/messages"
16+
"github.com/teslamotors/fleet-telemetry/metrics"
17+
"github.com/teslamotors/fleet-telemetry/protos"
18+
"github.com/teslamotors/fleet-telemetry/server/airbrake"
19+
"github.com/teslamotors/fleet-telemetry/telemetry"
20+
)
21+
22+
var _ = Describe("ZMQ Producer", func() {
23+
var (
24+
mockLogger *logrus.Logger
25+
mockCollector metrics.MetricCollector
26+
mockAirbrake *airbrake.Handler
27+
mockConfig *zmq.Config
28+
serializer *telemetry.BinarySerializer
29+
)
30+
31+
BeforeEach(func() {
32+
mockLogger, _ = logrus.NoOpLogger()
33+
mockCollector = metrics.NewCollector(nil, mockLogger)
34+
mockAirbrake = airbrake.NewAirbrakeHandler(nil)
35+
mockConfig = &zmq.Config{Addr: "tcp://127.0.0.1:*"}
36+
37+
serializer = telemetry.NewBinarySerializer(
38+
&telemetry.RequestIdentity{
39+
DeviceID: "TEST123",
40+
SenderID: "vehicle_device.TEST123",
41+
},
42+
map[string][]telemetry.Producer{},
43+
mockLogger,
44+
)
45+
})
46+
47+
newRecord := func() *telemetry.Record {
48+
payload := &protos.Payload{Vin: "TEST123"}
49+
payloadBytes, err := proto.Marshal(payload)
50+
Expect(err).NotTo(HaveOccurred())
51+
52+
message := messages.StreamMessage{
53+
TXID: []byte("1234"),
54+
SenderID: []byte("vehicle_device.TEST123"),
55+
MessageTopic: []byte("V"),
56+
Payload: payloadBytes,
57+
}
58+
msgBytes, err := message.ToBytes()
59+
Expect(err).NotTo(HaveOccurred())
60+
record, err := telemetry.NewRecord(serializer, msgBytes, "1", true)
61+
Expect(err).NotTo(HaveOccurred())
62+
return record
63+
}
64+
65+
It("handles concurrent Produce calls safely", func() {
66+
producer, err := zmq.NewProducer(
67+
context.Background(),
68+
mockConfig,
69+
mockCollector,
70+
"test_namespace",
71+
mockAirbrake,
72+
nil,
73+
nil,
74+
mockLogger,
75+
)
76+
Expect(err).NotTo(HaveOccurred())
77+
defer func() {
78+
Expect(producer.Close()).To(Succeed())
79+
}()
80+
81+
rec := newRecord()
82+
83+
const goroutines = 50
84+
const perGoroutines = 100
85+
var wg sync.WaitGroup
86+
for range goroutines {
87+
wg.Add(1)
88+
go func() {
89+
defer wg.Done()
90+
for range perGoroutines {
91+
producer.Produce(rec)
92+
}
93+
}()
94+
}
95+
wg.Wait()
96+
})
97+
98+
It("does not block or panic when Close races with Produce", func() {
99+
producer, err := zmq.NewProducer(
100+
context.Background(),
101+
mockConfig,
102+
mockCollector,
103+
"test_namespace",
104+
mockAirbrake,
105+
nil,
106+
nil,
107+
mockLogger,
108+
)
109+
Expect(err).NotTo(HaveOccurred())
110+
111+
rec := newRecord()
112+
113+
const goroutines = 10
114+
const perGoroutines = 1000
115+
var wg sync.WaitGroup
116+
117+
for range goroutines {
118+
wg.Add(1)
119+
go func() {
120+
defer GinkgoRecover()
121+
defer wg.Done()
122+
for range perGoroutines {
123+
producer.Produce(rec)
124+
}
125+
}()
126+
}
127+
128+
Expect(producer.Close()).To(Succeed())
129+
130+
done := make(chan struct{})
131+
go func() {
132+
wg.Wait()
133+
close(done)
134+
}()
135+
Eventually(done, "5s").Should(BeClosed())
136+
})
137+
138+
It("can be closed more than once", func() {
139+
producer, err := zmq.NewProducer(
140+
context.Background(),
141+
mockConfig,
142+
mockCollector,
143+
"test_namespace",
144+
mockAirbrake,
145+
nil,
146+
nil,
147+
mockLogger,
148+
)
149+
Expect(err).NotTo(HaveOccurred())
150+
151+
Expect(producer.Close()).To(Succeed())
152+
Expect(producer.Close()).To(Succeed())
153+
})
154+
155+
It("delivers published records to a subscriber with intact framing", func() {
156+
addr := "inproc://zmq-delivery-test"
157+
cfg := &zmq.Config{Addr: addr}
158+
producer, err := zmq.NewProducer(
159+
context.Background(),
160+
cfg,
161+
mockCollector,
162+
"test_namespace",
163+
mockAirbrake,
164+
nil,
165+
nil,
166+
mockLogger,
167+
)
168+
Expect(err).NotTo(HaveOccurred())
169+
defer func() {
170+
Expect(producer.Close()).To(Succeed())
171+
}()
172+
173+
sub, err := zmq4.NewSocket(zmq4.SUB)
174+
Expect(err).NotTo(HaveOccurred())
175+
defer func() { _ = sub.Close() }()
176+
Expect(sub.SetSubscribe("")).To(Succeed())
177+
Expect(sub.Connect(addr)).To(Succeed())
178+
179+
rec := newRecord()
180+
expectedTopic := telemetry.BuildTopicName("test_namespace", rec.TxType)
181+
expectedPayload := rec.Payload()
182+
183+
Expect(sub.SetRcvtimeo(100 * time.Millisecond)).To(Succeed())
184+
185+
// warm up: publish one at a time until the first message gets through
186+
// a successful receive proves the subscription has reached the publisher
187+
Eventually(func() bool {
188+
producer.Produce(rec)
189+
_, err := sub.RecvMessageBytes(0)
190+
return err == nil
191+
}, "5s", "10ms").Should(BeTrue())
192+
193+
for range 5 {
194+
producer.Produce(rec)
195+
frames, err := sub.RecvMessageBytes(0)
196+
Expect(err).NotTo(HaveOccurred())
197+
Expect(frames).To(HaveLen(2))
198+
Expect(string(frames[0])).To(Equal(expectedTopic))
199+
Expect(frames[1]).To(Equal(expectedPayload))
200+
}
201+
})
202+
})

0 commit comments

Comments
 (0)