|
| 1 | +// Licensed to Apache Software Foundation (ASF) under one or more contributor |
| 2 | +// license agreements. See the NOTICE file distributed with |
| 3 | +// this work for additional information regarding copyright |
| 4 | +// ownership. Apache Software Foundation (ASF) licenses this file to you under |
| 5 | +// the Apache License, Version 2.0 (the "License"); you may |
| 6 | +// not use this file except in compliance with the License. |
| 7 | +// You may obtain a copy of the License at |
| 8 | +// |
| 9 | +// http://www.apache.org/licenses/LICENSE-2.0 |
| 10 | +// |
| 11 | +// Unless required by applicable law or agreed to in writing, |
| 12 | +// software distributed under the License is distributed on an |
| 13 | +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| 14 | +// KIND, either express or implied. See the License for the |
| 15 | +// specific language governing permissions and limitations |
| 16 | +// under the License. |
| 17 | + |
| 18 | +// Package protector provides a set of protectors that stop the query services when the resource usage exceeds the limit. |
| 19 | +package protector |
| 20 | + |
| 21 | +import ( |
| 22 | + "context" |
| 23 | + "errors" |
| 24 | + "fmt" |
| 25 | + "runtime/metrics" |
| 26 | + "sync/atomic" |
| 27 | + "time" |
| 28 | + |
| 29 | + "github.com/dustin/go-humanize" |
| 30 | + |
| 31 | + "github.com/apache/skywalking-banyandb/banyand/observability" |
| 32 | + "github.com/apache/skywalking-banyandb/pkg/cgroups" |
| 33 | + "github.com/apache/skywalking-banyandb/pkg/logger" |
| 34 | + "github.com/apache/skywalking-banyandb/pkg/meter" |
| 35 | + "github.com/apache/skywalking-banyandb/pkg/run" |
| 36 | +) |
| 37 | + |
| 38 | +var scope = observability.RootScope.SubScope("memory_protector") |
| 39 | + |
| 40 | +// Memory is a protector that stops the query services when the memory usage exceeds the limit. |
| 41 | +type Memory struct { |
| 42 | + omr observability.MetricsRegistry |
| 43 | + limitGauge meter.Gauge |
| 44 | + usageGauge meter.Gauge |
| 45 | + l *logger.Logger |
| 46 | + closed chan struct{} |
| 47 | + blockedChan chan struct{} |
| 48 | + allowedPercent int |
| 49 | + allowedBytes run.Bytes |
| 50 | + limit uint64 |
| 51 | + usage uint64 |
| 52 | +} |
| 53 | + |
| 54 | +// NewMemory creates a new Memory protector. |
| 55 | +func NewMemory(omr observability.MetricsRegistry) *Memory { |
| 56 | + queueSize := cgroups.CPUs() |
| 57 | + factory := omr.With(scope) |
| 58 | + |
| 59 | + return &Memory{ |
| 60 | + omr: omr, |
| 61 | + blockedChan: make(chan struct{}, queueSize), |
| 62 | + closed: make(chan struct{}), |
| 63 | + |
| 64 | + limitGauge: factory.NewGauge("limit"), |
| 65 | + usageGauge: factory.NewGauge("usage"), |
| 66 | + } |
| 67 | +} |
| 68 | + |
| 69 | +// AcquireResource attempts to acquire a `size` amount of memory. |
| 70 | +func (m *Memory) AcquireResource(ctx context.Context, size uint64) error { |
| 71 | + if m.limit == 0 { |
| 72 | + return nil |
| 73 | + } |
| 74 | + start := time.Now() |
| 75 | + |
| 76 | + select { |
| 77 | + case m.blockedChan <- struct{}{}: |
| 78 | + defer func() { <-m.blockedChan }() |
| 79 | + case <-ctx.Done(): |
| 80 | + return fmt.Errorf("context canceled while waiting for blocked queue slot: %w", ctx.Err()) |
| 81 | + } |
| 82 | + |
| 83 | + for { |
| 84 | + currentUsage := atomic.LoadUint64(&m.usage) |
| 85 | + if currentUsage+size <= m.limit { |
| 86 | + return nil |
| 87 | + } |
| 88 | + |
| 89 | + select { |
| 90 | + case <-time.After(100 * time.Millisecond): |
| 91 | + continue |
| 92 | + case <-ctx.Done(): |
| 93 | + return fmt.Errorf( |
| 94 | + "context canceled: memory acquisition failed (currentUsage: %d, limit: %d, size: %d, blockedDuration: %v): %w", |
| 95 | + currentUsage, m.limit, size, time.Since(start), ctx.Err(), |
| 96 | + ) |
| 97 | + } |
| 98 | + } |
| 99 | +} |
| 100 | + |
| 101 | +// Name returns the name of the protector. |
| 102 | +func (m *Memory) Name() string { |
| 103 | + return "memory-protector" |
| 104 | +} |
| 105 | + |
| 106 | +// FlagSet returns the flag set for the protector. |
| 107 | +func (m *Memory) FlagSet() *run.FlagSet { |
| 108 | + flagS := run.NewFlagSet(m.Name()) |
| 109 | + flagS.IntVarP(&m.allowedPercent, "allowed-percent", "", 75, |
| 110 | + "Allowed bytes of memory usage. If the memory usage exceeds this value, the query services will stop. "+ |
| 111 | + "Setting a large value may evict data from the OS page cache, causing high disk I/O.") |
| 112 | + flagS.VarP(&m.allowedBytes, "allowed-bytes", "", "Allowed percentage of total memory usage. If usage exceeds this value, the query services will stop. "+ |
| 113 | + "This takes effect only if `allowed-bytes` is 0. If usage is too high, it may cause OS page cache eviction.") |
| 114 | + return flagS |
| 115 | +} |
| 116 | + |
| 117 | +// Validate validates the protector's flags. |
| 118 | +func (m *Memory) Validate() error { |
| 119 | + if m.allowedPercent <= 0 || m.allowedPercent > 100 { |
| 120 | + if m.allowedBytes <= 0 { |
| 121 | + return errors.New("allowed-bytes must be greater than 0") |
| 122 | + } |
| 123 | + return errors.New("allowed-percent must be in the range (0, 100]") |
| 124 | + } |
| 125 | + return nil |
| 126 | +} |
| 127 | + |
| 128 | +// PreRun initializes the protector. |
| 129 | +func (m *Memory) PreRun(context.Context) error { |
| 130 | + m.l = logger.GetLogger(m.Name()) |
| 131 | + if m.allowedBytes > 0 { |
| 132 | + m.limit = uint64(m.allowedBytes) |
| 133 | + m.l.Info(). |
| 134 | + Str("limit", humanize.Bytes(m.limit)). |
| 135 | + Msg("memory protector enabled") |
| 136 | + } else { |
| 137 | + cgLimit, err := cgroups.MemoryLimit() |
| 138 | + if err != nil { |
| 139 | + m.l.Warn().Err(err).Msg("failed to get memory limit from cgroups, disable memory protector") |
| 140 | + return nil |
| 141 | + } |
| 142 | + if cgLimit <= 0 || cgLimit > 1e18 { |
| 143 | + m.l.Warn().Int64("cgroup_memory_limit", cgLimit).Msg("cgroup memory limit is invalid, disable memory protector") |
| 144 | + return nil |
| 145 | + } |
| 146 | + m.limit = uint64(cgLimit) * uint64(m.allowedPercent) / 100 |
| 147 | + m.l.Info(). |
| 148 | + Str("limit", humanize.Bytes(m.limit)). |
| 149 | + Str("cgroup_limit", humanize.Bytes(uint64(cgLimit))). |
| 150 | + Int("percent", m.allowedPercent). |
| 151 | + Msg("memory protector enabled") |
| 152 | + } |
| 153 | + m.limitGauge.Set(float64(m.limit)) |
| 154 | + return nil |
| 155 | +} |
| 156 | + |
| 157 | +// GracefulStop stops the protector. |
| 158 | +func (m *Memory) GracefulStop() { |
| 159 | + close(m.closed) |
| 160 | +} |
| 161 | + |
| 162 | +// Serve starts the protector. |
| 163 | +func (m *Memory) Serve() run.StopNotify { |
| 164 | + if m.limit == 0 { |
| 165 | + return m.closed |
| 166 | + } |
| 167 | + go func() { |
| 168 | + ticker := time.NewTicker(5 * time.Second) |
| 169 | + defer ticker.Stop() |
| 170 | + |
| 171 | + for { |
| 172 | + select { |
| 173 | + case <-m.closed: |
| 174 | + return |
| 175 | + case <-ticker.C: |
| 176 | + samples := []metrics.Sample{ |
| 177 | + {Name: "/memory/classes/total:bytes"}, |
| 178 | + } |
| 179 | + metrics.Read(samples) |
| 180 | + usedBytes := samples[0].Value.Uint64() |
| 181 | + |
| 182 | + atomic.StoreUint64(&m.usage, usedBytes) |
| 183 | + |
| 184 | + if usedBytes > m.limit { |
| 185 | + m.l.Warn().Str("used", humanize.Bytes(usedBytes)).Str("limit", humanize.Bytes(m.limit)).Msg("memory usage exceeds limit") |
| 186 | + } |
| 187 | + } |
| 188 | + } |
| 189 | + }() |
| 190 | + return m.closed |
| 191 | +} |
0 commit comments