44package store
55
66import (
7+ "context"
78 "fmt"
89 "io/ioutil"
910 "math"
1011 "math/rand"
1112 "os"
1213 "path/filepath"
1314 "testing"
15+ "time"
1416
17+ "github.com/fortytw2/leaktest"
1518 "github.com/go-kit/kit/log"
19+ "github.com/pkg/errors"
20+ "github.com/prometheus/prometheus/pkg/labels"
1621 "github.com/prometheus/prometheus/tsdb"
1722 "github.com/thanos-io/thanos/pkg/component"
1823 "github.com/thanos-io/thanos/pkg/store/storepb"
@@ -21,6 +26,8 @@ import (
2126)
2227
2328func TestMultiTSDBSeries (t * testing.T ) {
29+ defer leaktest .CheckTimeout (t , 10 * time .Second )()
30+
2431 tb := testutil .NewTB (t )
2532 storetestutil .RunSeriesInterestingCases (tb , 200e3 , 200e3 , func (t testutil.TB , samplesPerSeries , series int ) {
2633 if ok := t .Run ("headOnly" , func (t testutil.TB ) {
@@ -116,12 +123,12 @@ func benchMultiTSDBSeries(t testutil.TB, totalSamples, totalSeries int, flushToB
116123 dbs [j ] = & mockedStartTimeDB {DBReadOnly : db , startTime : int64 (j * samplesPerSeriesPerTSDB * seriesPerTSDB )}
117124 }
118125
119- tsdbs := map [string ]* TSDBStore {}
126+ tsdbs := map [string ]storepb. StoreServer {}
120127 for i , db := range dbs {
121128 tsdbs [fmt .Sprintf ("%v" , i )] = & TSDBStore {db : db , logger : logger }
122129 }
123130
124- store := NewMultiTSDBStore (logger , nil , component .Receive , func () map [string ]* TSDBStore { return tsdbs })
131+ store := NewMultiTSDBStore (logger , nil , component .Receive , func () map [string ]storepb. StoreServer { return tsdbs })
125132
126133 var expected []storepb.Series
127134 lastLabels := storepb.Series {}
@@ -154,3 +161,116 @@ func benchMultiTSDBSeries(t testutil.TB, totalSamples, totalSeries int, flushToB
154161 },
155162 )
156163}
164+
165+ type mockedStoreServer struct {
166+ storepb.StoreServer
167+
168+ responses []* storepb.SeriesResponse
169+ }
170+
171+ func (m * mockedStoreServer ) Series (_ * storepb.SeriesRequest , server storepb.Store_SeriesServer ) error {
172+ for _ , r := range m .responses {
173+ if err := server .Send (r ); err != nil {
174+ return err
175+ }
176+ }
177+ return nil
178+ }
179+
180+ // Regression test against https://github.com/thanos-io/thanos/issues/2823.
181+ func TestTenantSeriesSetServert_NotLeakingIfNotExhausted (t * testing.T ) {
182+ t .Run ("exhausted StoreSet" , func (t * testing.T ) {
183+ defer leaktest .CheckTimeout (t , 10 * time .Second )()
184+
185+ s := newTenantSeriesSetServer (context .Background (), "a" , nil )
186+
187+ resps := []* storepb.SeriesResponse {
188+ storeSeriesResponse (t , labels .FromStrings ("a" , "a" ), []sample {{0 , 0 }, {2 , 1 }, {3 , 2 }}),
189+ storeSeriesResponse (t , labels .FromStrings ("a" , "b" ), []sample {{0 , 0 }, {2 , 1 }, {3 , 2 }}),
190+ storeSeriesResponse (t , labels .FromStrings ("a" , "c" ), []sample {{0 , 0 }, {2 , 1 }, {3 , 2 }}),
191+ }
192+
193+ m := & mockedStoreServer {responses : resps }
194+
195+ go func () {
196+ s .Series (m , & storepb.SeriesRequest {PartialResponseStrategy : storepb .PartialResponseStrategy_ABORT })
197+ }()
198+
199+ testutil .Ok (t , s .Err ())
200+ i := 0
201+ for s .Next () {
202+ l , c := s .At ()
203+
204+ testutil .Equals (t , resps [i ].GetSeries ().Labels , l )
205+ testutil .Equals (t , resps [i ].GetSeries ().Chunks , c )
206+
207+ i ++
208+ }
209+ testutil .Ok (t , s .Err ())
210+ testutil .Equals (t , 3 , i )
211+ })
212+
213+ t .Run ("cancelled, not exhausted StoreSet" , func (t * testing.T ) {
214+ defer leaktest .CheckTimeout (t , 10 * time .Second )()
215+
216+ ctx , cancel := context .WithCancel (context .Background ())
217+ s := newTenantSeriesSetServer (ctx , "a" , nil )
218+
219+ m := & mockedStoreServer {responses : []* storepb.SeriesResponse {
220+ storeSeriesResponse (t , labels .FromStrings ("a" , "a" ), []sample {{0 , 0 }, {2 , 1 }, {3 , 2 }}),
221+ storeSeriesResponse (t , labels .FromStrings ("a" , "b" ), []sample {{0 , 0 }, {2 , 1 }, {3 , 2 }}),
222+ storeSeriesResponse (t , labels .FromStrings ("a" , "c" ), []sample {{0 , 0 }, {2 , 1 }, {3 , 2 }}),
223+ }}
224+ go func () {
225+ s .Series (m , & storepb.SeriesRequest {PartialResponseStrategy : storepb .PartialResponseStrategy_ABORT })
226+ }()
227+
228+ testutil .Ok (t , s .Err ())
229+ testutil .Equals (t , true , s .Next ())
230+ cancel ()
231+ })
232+ }
233+
234+ type faillingSeriesServer struct {
235+ storepb.Store_SeriesServer
236+
237+ ctx context.Context
238+ }
239+
240+ func (s * faillingSeriesServer ) Send (* storepb.SeriesResponse ) error {
241+ return errors .New ("I am broken" )
242+ }
243+ func (s * faillingSeriesServer ) Context () context.Context { return s .ctx }
244+
245+ // Regression test against https://github.com/thanos-io/thanos/issues/2823.
246+ // This is different leak than in TestTenantSeriesSetServert_NotLeakingIfNotExhausted
247+ func TestMultiTSDBStore_NotLeakingOnPrematureFinish (t * testing.T ) {
248+ defer leaktest .CheckTimeout (t , 10 * time .Second )()
249+
250+ m := NewMultiTSDBStore (log .NewNopLogger (), nil , component .Receive , func () map [string ]storepb.StoreServer {
251+ return map [string ]storepb.StoreServer {
252+ // Ensure more than 10 (internal respCh channel).
253+ "a" : & mockedStoreServer {responses : []* storepb.SeriesResponse {
254+ storeSeriesResponse (t , labels .FromStrings ("a" , "a" ), []sample {{0 , 0 }, {2 , 1 }, {3 , 2 }}),
255+ storeSeriesResponse (t , labels .FromStrings ("a" , "b" ), []sample {{0 , 0 }, {2 , 1 }, {3 , 2 }}),
256+ storeSeriesResponse (t , labels .FromStrings ("a" , "c" ), []sample {{0 , 0 }, {2 , 1 }, {3 , 2 }}),
257+ storeSeriesResponse (t , labels .FromStrings ("a" , "a" ), []sample {{0 , 0 }, {2 , 1 }, {3 , 2 }}),
258+ storeSeriesResponse (t , labels .FromStrings ("a" , "b" ), []sample {{0 , 0 }, {2 , 1 }, {3 , 2 }}),
259+ storeSeriesResponse (t , labels .FromStrings ("a" , "c" ), []sample {{0 , 0 }, {2 , 1 }, {3 , 2 }}),
260+ }},
261+ "b" : & mockedStoreServer {responses : []* storepb.SeriesResponse {
262+ storeSeriesResponse (t , labels .FromStrings ("b" , "a" ), []sample {{0 , 0 }, {2 , 1 }, {3 , 2 }}),
263+ storeSeriesResponse (t , labels .FromStrings ("b" , "b" ), []sample {{0 , 0 }, {2 , 1 }, {3 , 2 }}),
264+ storeSeriesResponse (t , labels .FromStrings ("b" , "c" ), []sample {{0 , 0 }, {2 , 1 }, {3 , 2 }}),
265+ storeSeriesResponse (t , labels .FromStrings ("a" , "a" ), []sample {{0 , 0 }, {2 , 1 }, {3 , 2 }}),
266+ storeSeriesResponse (t , labels .FromStrings ("a" , "b" ), []sample {{0 , 0 }, {2 , 1 }, {3 , 2 }}),
267+ storeSeriesResponse (t , labels .FromStrings ("a" , "c" ), []sample {{0 , 0 }, {2 , 1 }, {3 , 2 }}),
268+ }},
269+ }
270+ })
271+
272+ ctx , cancel := context .WithCancel (context .Background ())
273+ // We mimic failing series server, but practically context cancel will do the same.
274+ testutil .NotOk (t , m .Series (& storepb.SeriesRequest {PartialResponseStrategy : storepb .PartialResponseStrategy_ABORT }, & faillingSeriesServer {ctx : ctx }))
275+ cancel ()
276+ }
0 commit comments