@ -333,18 +333,25 @@ func Test_blockPlansForGaps(t *testing.T) {
}
}
func Test_BuilderLoop ( t * testing . T ) {
const (
nTasks = 100
nBuilders = 10
)
logger := log . NewNopLogger ( )
limits := & fakeLimits { }
cfg := Config {
PlanningInterval : 1 * time . Hour ,
MaxQueuedTasksPerTenant : 10000 ,
func createTasks ( n int ) [ ] * Task {
tasks := make ( [ ] * Task , 0 , n )
// Enqueue tasks
for i := 0 ; i < n ; i ++ {
task := NewTask (
context . Background ( ) , time . Now ( ) ,
protos . NewTask ( config . NewDayTable ( config . NewDayTime ( 0 ) , "fake" ) , "fakeTenant" , v1 . NewBounds ( 0 , 10 ) , tsdbID ( 1 ) , nil ) ,
)
tasks = append ( tasks , task )
}
return tasks
}
func createPlanner (
t * testing . T ,
cfg Config ,
limits Limits ,
logger log . Logger ,
) * Planner {
schemaCfg := config . SchemaConfig {
Configs : [ ] config . PeriodConfig {
{
@ -377,77 +384,227 @@ func Test_BuilderLoop(t *testing.T) {
} ,
}
// Create planner
planner , err := New ( cfg , limits , schemaCfg , storageCfg , storage . NewClientMetrics ( ) , nil , logger , p rometheus . DefaultR egisterer )
reg := prometheus . NewPedanticRegistry ( )
planner , err := New ( cfg , limits , schemaCfg , storageCfg , storage . ClientMetrics { } , nil , logger , reg )
require . NoError ( t , err )
// Start planner
err = services . StartAndAwaitRunning ( context . Background ( ) , planner )
require . NoError ( t , err )
t . Cleanup ( func ( ) {
err := services . StopAndAwaitTerminated ( context . Background ( ) , planner )
require . NoError ( t , err )
} )
return planner
}
// Enqueue tasks
for i := 0 ; i < nTasks ; i ++ {
task := NewTask (
context . Background ( ) , time . Now ( ) ,
protos . NewTask ( config . NewDayTable ( config . NewDayTime ( 0 ) , "fake" ) , "fakeTenant" , v1 . NewBounds ( 0 , 10 ) , tsdbID ( 1 ) , nil ) ,
)
func Test_BuilderLoop ( t * testing . T ) {
const (
nTasks = 100
nBuilders = 10
)
err = planner . enqueueTask ( task )
require . NoError ( t , err )
}
for _ , tc := range [ ] struct {
name string
limits Limits
expectedBuilderLoopError error
// modifyBuilder should leave the builder in a state where it will not return or return an error
modifyBuilder func ( builder * fakeBuilder )
// resetBuilder should reset the builder to a state where it will return no errors
resetBuilder func ( builder * fakeBuilder )
} {
{
name : "success" ,
limits : & fakeLimits { } ,
expectedBuilderLoopError : errPlannerIsNotRunning ,
} ,
{
name : "error rpc" ,
limits : & fakeLimits { } ,
expectedBuilderLoopError : errPlannerIsNotRunning ,
modifyBuilder : func ( builder * fakeBuilder ) {
builder . SetReturnError ( true )
} ,
resetBuilder : func ( builder * fakeBuilder ) {
builder . SetReturnError ( false )
} ,
} ,
{
name : "error msg" ,
limits : & fakeLimits { } ,
expectedBuilderLoopError : errPlannerIsNotRunning ,
modifyBuilder : func ( builder * fakeBuilder ) {
builder . SetReturnErrorMsg ( true )
} ,
resetBuilder : func ( builder * fakeBuilder ) {
builder . SetReturnErrorMsg ( false )
} ,
} ,
{
name : "timeout" ,
limits : & fakeLimits {
timeout : 1 * time . Second ,
} ,
expectedBuilderLoopError : errPlannerIsNotRunning ,
modifyBuilder : func ( builder * fakeBuilder ) {
builder . SetWait ( true )
} ,
resetBuilder : func ( builder * fakeBuilder ) {
builder . SetWait ( false )
} ,
} ,
{
name : "context cancel" ,
limits : & fakeLimits { } ,
// Builders cancel the context when they disconnect. We forward this error to the planner.
expectedBuilderLoopError : context . Canceled ,
modifyBuilder : func ( builder * fakeBuilder ) {
builder . CancelContext ( true )
} ,
} ,
} {
t . Run ( tc . name , func ( t * testing . T ) {
logger := log . NewNopLogger ( )
//logger := log.NewLogfmtLogger(os.Stdout)
// All tasks should be pending
require . Equal ( t , nTasks , planner . totalPendingTasks ( ) )
cfg := Config {
PlanningInterval : 1 * time . Hour ,
MaxQueuedTasksPerTenant : 10000 ,
}
planner := createPlanner ( t , cfg , tc . limits , logger )
// Create builders and call planner.BuilderLoop
builders := make ( [ ] * fakeBuilder , 0 , nBuilders )
for i := 0 ; i < nBuilders ; i ++ {
builder := newMockBuilder ( fmt . Sprintf ( "builder-%d" , i ) )
builders = append ( builders , builder )
// Start planner
err := services . StartAndAwaitRunning ( context . Background ( ) , planner )
require . NoError ( t , err )
t . Cleanup ( func ( ) {
err := services . StopAndAwaitTerminated ( context . Background ( ) , planner )
require . NoError ( t , err )
} )
// Enqueue tasks
tasks := createTasks ( nTasks )
for _ , task := range tasks {
err = planner . enqueueTask ( task )
require . NoError ( t , err )
}
go func ( ) {
// We ignore the error since when the planner is stopped,
// the loop will return an error (queue closed)
_ = planner . BuilderLoop ( builder )
} ( )
}
// Create builders and call planner.BuilderLoop
builders := make ( [ ] * fakeBuilder , 0 , nBuilders )
for i := 0 ; i < nBuilders ; i ++ {
builder := newMockBuilder ( fmt . Sprintf ( "builder-%d" , i ) )
builders = append ( builders , builder )
go func ( ) {
err = planner . BuilderLoop ( builder )
require . ErrorIs ( t , err , tc . expectedBuilderLoopError )
} ( )
}
// Eventually, all tasks should be sent to builders
require . Eventually ( t , func ( ) bool {
var receivedTasks int
for _ , builder := range builders {
receivedTasks += len ( builder . ReceivedTasks ( ) )
}
return receivedTasks == nTasks
} , 5 * time . Second , 10 * time . Millisecond )
// Finally, the queue should be empty
require . Equal ( t , 0 , planner . totalPendingTasks ( ) )
if tc . modifyBuilder != nil {
// Configure builders to return errors
for _ , builder := range builders {
tc . modifyBuilder ( builder )
}
// Enqueue tasks again
for _ , task := range tasks {
err = planner . enqueueTask ( task )
require . NoError ( t , err )
}
// Tasks should not be consumed
require . Neverf (
t , func ( ) bool {
return planner . totalPendingTasks ( ) == 0
} ,
5 * time . Second , 10 * time . Millisecond ,
"all tasks were consumed but they should not be" ,
)
}
if tc . resetBuilder != nil {
// Configure builders to return no errors
for _ , builder := range builders {
tc . resetBuilder ( builder )
}
// Eventually, all tasks should be sent to builders
require . Eventually ( t , func ( ) bool {
var receivedTasks int
for _ , builder := range builders {
receivedTasks += len ( builder . ReceivedTasks ( ) )
}
return receivedTasks == nTasks
} , 15 * time . Second , 10 * time . Millisecond )
// Finally, the queue should be empty
require . Equal ( t , 0 , planner . totalPendingTasks ( ) )
// Now all tasks should be consumed
require . Eventuallyf (
t , func ( ) bool {
return planner . totalPendingTasks ( ) == 0
} ,
5 * time . Second , 10 * time . Millisecond ,
"tasks not consumed, pending: %d" , planner . totalPendingTasks ( ) ,
)
}
} )
}
}
type fakeBuilder struct {
id string
tasks [ ] * protos . Task
grpc . ServerStream
returnError bool
returnErrorMsg bool
wait bool
ctx context . Context
ctxCancel context . CancelFunc
}
func newMockBuilder ( id string ) * fakeBuilder {
return & fakeBuilder { id : id }
ctx , cancel := context . WithCancel ( context . Background ( ) )
return & fakeBuilder {
id : id ,
ctx : ctx ,
ctxCancel : cancel ,
}
}
func ( f * fakeBuilder ) ReceivedTasks ( ) [ ] * protos . Task {
return f . tasks
}
func ( f * fakeBuilder ) SetReturnError ( b bool ) {
f . returnError = b
}
func ( f * fakeBuilder ) SetReturnErrorMsg ( b bool ) {
f . returnErrorMsg = b
}
func ( f * fakeBuilder ) SetWait ( b bool ) {
f . wait = b
}
func ( f * fakeBuilder ) CancelContext ( b bool ) {
if b {
f . ctxCancel ( )
return
}
// Reset context
f . ctx , f . ctxCancel = context . WithCancel ( context . Background ( ) )
}
func ( f * fakeBuilder ) Context ( ) context . Context {
return context . Background ( )
return f . ctx
}
func ( f * fakeBuilder ) Send ( req * protos . PlannerToBuilder ) error {
if f . ctx . Err ( ) != nil {
// Context was canceled
return f . ctx . Err ( )
}
task , err := protos . FromProtoTask ( req . Task )
if err != nil {
return err
@ -458,12 +615,37 @@ func (f *fakeBuilder) Send(req *protos.PlannerToBuilder) error {
}
func ( f * fakeBuilder ) Recv ( ) ( * protos . BuilderToPlanner , error ) {
if f . returnError {
return nil , fmt . Errorf ( "fake error from %s" , f . id )
}
// Wait until `wait` is false
for f . wait {
time . Sleep ( time . Second )
}
if f . ctx . Err ( ) != nil {
// Context was canceled
return nil , f . ctx . Err ( )
}
var errMsg string
if f . returnErrorMsg {
errMsg = fmt . Sprintf ( "fake error from %s" , f . id )
}
return & protos . BuilderToPlanner {
BuilderID : f . id ,
Error : errMsg ,
} , nil
}
type fakeLimits struct {
timeout time . Duration
}
func ( f * fakeLimits ) BuilderResponseTimeout ( _ string ) time . Duration {
return f . timeout
}
func ( f * fakeLimits ) BloomCreationEnabled ( _ string ) bool {