mirror of
https://github.com/gravitational/teleport.git
synced 2026-11-01 21:52:02 +08:00
Marco André Dinis
455632e28a
* AutoDiscover: add support for EICE nodes * ignore openssh nodes when updating labels from serverinfo * downcase imports * improve comment on why we need account/instance and subkind * add applyovertime method * skip nodes that wont expire in the near future * create helper functions
114 lines
3.0 KiB
Go
114 lines
3.0 KiB
Go
/*
|
|
* Teleport
|
|
* Copyright (C) 2024 Gravitational, Inc.
|
|
*
|
|
* This program is free software: you can redistribute it and/or modify
|
|
* it under the terms of the GNU Affero General Public License as published by
|
|
* the Free Software Foundation, either version 3 of the License, or
|
|
* (at your option) any later version.
|
|
*
|
|
* This program is distributed in the hope that it will be useful,
|
|
* but WITHOUT ANY WARRANTY; without even the implied warranty of
|
|
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
|
* GNU Affero General Public License for more details.
|
|
*
|
|
* You should have received a copy of the GNU Affero General Public License
|
|
* along with this program. If not, see <http://www.gnu.org/licenses/>.
|
|
*/
|
|
|
|
package spreadwork
|
|
|
|
import (
|
|
"context"
|
|
"time"
|
|
|
|
"github.com/gravitational/trace"
|
|
"github.com/jonboulle/clockwork"
|
|
)
|
|
|
|
// ApplyOverTimeConfig is the configuration values for the ApplyOverTime method.
|
|
type ApplyOverTimeConfig struct {
|
|
// MaxDuration is the max duration for the operation to complete.
|
|
// It does not account for time spent in API calls.
|
|
MaxDuration time.Duration
|
|
|
|
// BatchInterval is the time spent between API calls to prevent server pressure.
|
|
// Defaults to 1 second.
|
|
BatchInterval time.Duration
|
|
|
|
// MinBatchSize is the min batch size per iteration.
|
|
// Defaults to 5 items.
|
|
MinBatchSize int
|
|
|
|
// clock is used for tests only.
|
|
clock clockwork.Clock
|
|
}
|
|
|
|
// CheckAndSetDefaults checks that the MaxDuration was set.
|
|
func (c *ApplyOverTimeConfig) CheckAndSetDefaults() error {
|
|
if c.MaxDuration == 0 {
|
|
return trace.BadParameter("max duration is required")
|
|
}
|
|
|
|
if c.BatchInterval == 0 {
|
|
c.BatchInterval = time.Second
|
|
}
|
|
|
|
if c.MaxDuration < c.BatchInterval {
|
|
return trace.BadParameter("max interval must be greater than batch interval")
|
|
}
|
|
|
|
if c.MinBatchSize == 0 {
|
|
c.MinBatchSize = 5
|
|
}
|
|
|
|
if c.clock == nil {
|
|
c.clock = clockwork.NewRealClock()
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
// ApplyOverTime applies the applyFn to all elements of items during a period of time.
|
|
func ApplyOverTime[T any](ctx context.Context, conf ApplyOverTimeConfig, items []T, applyFn func(T)) error {
|
|
if err := conf.CheckAndSetDefaults(); err != nil {
|
|
return trace.Wrap(err)
|
|
}
|
|
ticker := conf.clock.NewTicker(conf.BatchInterval)
|
|
defer ticker.Stop()
|
|
|
|
maxBatches := int(conf.MaxDuration / conf.BatchInterval)
|
|
dynamicBatchSize := (len(items) / maxBatches) + 1
|
|
if dynamicBatchSize < conf.MinBatchSize {
|
|
dynamicBatchSize = conf.MinBatchSize
|
|
}
|
|
|
|
for {
|
|
if dynamicBatchSize > len(items) {
|
|
dynamicBatchSize = len(items)
|
|
}
|
|
chunk := items[0:dynamicBatchSize:dynamicBatchSize]
|
|
items = items[dynamicBatchSize:]
|
|
|
|
for _, c := range chunk {
|
|
// If ctx is Done, return without processing the other elements of the current chunk.
|
|
select {
|
|
case <-ctx.Done():
|
|
return trace.Wrap(ctx.Err())
|
|
default:
|
|
applyFn(c)
|
|
}
|
|
}
|
|
if len(items) == 0 {
|
|
return nil
|
|
}
|
|
|
|
select {
|
|
// If ctx is Done, return without processing the remaining chunks.
|
|
case <-ctx.Done():
|
|
return trace.Wrap(ctx.Err())
|
|
case <-ticker.Chan():
|
|
}
|
|
}
|
|
}
|