github.com/GoogleCloudPlatform/testgrid@v0.0.174/config/queue.go (about)

     1  /*
     2  Copyright 2021 The TestGrid Authors.
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8      http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  
    17  package config
    18  
    19  import (
    20  	"context"
    21  	"sync"
    22  	"time"
    23  
    24  	"bitbucket.org/creachadair/stringset"
    25  	configpb "github.com/GoogleCloudPlatform/testgrid/pb/config"
    26  	"github.com/GoogleCloudPlatform/testgrid/util/queue"
    27  	"github.com/sirupsen/logrus"
    28  )
    29  
    30  // DashboardQueue sends dashboard names at a specific frequency.
    31  type DashboardQueue struct {
    32  	queue.Queue
    33  	dashboards map[string]*configpb.Dashboard
    34  	groups     map[string]*stringset.Set
    35  
    36  	lock sync.RWMutex
    37  }
    38  
    39  // Init (or reinit) the queue with the specified configuration.
    40  func (q *DashboardQueue) Init(log logrus.FieldLogger, dashboards []*configpb.Dashboard, when time.Time) {
    41  	n := len(dashboards)
    42  	names := make([]string, n)
    43  	namedDashboards := make(map[string]*configpb.Dashboard, n)
    44  	groups := make(map[string]*stringset.Set, n)
    45  	for i, d := range dashboards {
    46  		name := d.Name
    47  		names[i] = name
    48  		for _, tab := range d.DashboardTab {
    49  			if groups[tab.TestGroupName] == nil {
    50  				ns := stringset.New()
    51  				groups[tab.TestGroupName] = &ns
    52  			}
    53  			groups[tab.TestGroupName].Add(name)
    54  		}
    55  	}
    56  	q.lock.Lock()
    57  	q.Queue.Init(log, names, when)
    58  	q.dashboards = namedDashboards
    59  	q.groups = groups
    60  	q.lock.Unlock()
    61  }
    62  
    63  // FixTestGroups will fix all the dashboards associated with the groups.
    64  func (q *DashboardQueue) FixTestGroups(when time.Time, later bool, groups ...string) error {
    65  	q.lock.RLock()
    66  	defer q.lock.RUnlock()
    67  	dashboards := make(map[string]time.Time, len(groups))
    68  	for _, groupName := range groups {
    69  		dashes := q.groups[groupName]
    70  		if dashes == nil {
    71  			continue
    72  		}
    73  		for dashName := range *dashes {
    74  			dashboards[dashName] = when
    75  		}
    76  	}
    77  	return q.FixAll(dashboards, later)
    78  }
    79  
    80  // TestGroupQueue can send test groups to receivers at a specific frequency.
    81  //
    82  // Also contains the ability to modify the next time to send groups.
    83  // First call must be to Init().
    84  // Exported methods are safe to call concurrently.
    85  type TestGroupQueue struct {
    86  	queue.Queue
    87  	groups map[string]*configpb.TestGroup
    88  	lock   sync.RWMutex
    89  }
    90  
    91  // Init (or reinit) the queue with the specified groups, which should be updated at frequency.
    92  func (q *TestGroupQueue) Init(log logrus.FieldLogger, testGroups []*configpb.TestGroup, when time.Time) {
    93  	n := len(testGroups)
    94  	groups := make(map[string]*configpb.TestGroup, n)
    95  	names := make([]string, n)
    96  
    97  	for i, tg := range testGroups {
    98  		name := tg.Name
    99  		names[i] = name
   100  		groups[name] = tg
   101  	}
   102  
   103  	q.lock.Lock()
   104  	q.Queue.Init(log, names, when)
   105  	q.groups = groups
   106  	q.lock.Unlock()
   107  }
   108  
   109  // Status of the queue: depth, next item and when the next item is ready.
   110  func (q *TestGroupQueue) Status() (int, *configpb.TestGroup, time.Time) {
   111  	q.lock.RLock()
   112  	defer q.lock.RUnlock()
   113  	var tg *configpb.TestGroup
   114  	var when time.Time
   115  	n, who, when := q.Queue.Status()
   116  	if who != nil {
   117  		tg = q.groups[*who]
   118  	}
   119  	return n, tg, when
   120  }
   121  
   122  // Send test groups to receivers until the context expires.
   123  //
   124  // Pops items off the queue when frequency is zero.
   125  // Otherwise reschedules the item after the specified frequency has elapsed.
   126  func (q *TestGroupQueue) Send(ctx context.Context, receivers chan<- *configpb.TestGroup, frequency time.Duration) error {
   127  	ch := make(chan string)
   128  	var err error
   129  	go func() {
   130  		err = q.Queue.Send(ctx, ch, frequency)
   131  		close(ch)
   132  	}()
   133  
   134  	for who := range ch {
   135  		q.lock.RLock()
   136  		tg := q.groups[who]
   137  		q.lock.RUnlock()
   138  		if tg == nil {
   139  			continue
   140  		}
   141  		select {
   142  		case <-ctx.Done():
   143  			return ctx.Err()
   144  		case receivers <- tg:
   145  		}
   146  	}
   147  	return err
   148  }