| 1 | package daemon |
| 2 | |
| 3 | import ( |
| 4 | "context" |
| 5 | "errors" |
| 6 | "sync" |
| 7 | ) |
| 8 | |
| 9 | // Admission errors distinguish exclusive maintenance from tenant shutdown. |
| 10 | // Callers may retry maintenance refusals after the active operation finishes. |
| 11 | var ( |
| 12 | ErrMaintenance = errors.New("retryable: tenant maintenance in progress" ) |
| 13 | ErrMaintenanceBusy = errors.New("retryable: tenant maintenance is already active" ) |
| 14 | ErrTenantClosing = errors.New("tenant is closing" ) |
| 15 | ) |
| 16 | |
| 17 | // maintenanceGate counts admitted operations and grants exclusive maintenance |
| 18 | // after those operations drain. Context markers carry admission across nested |
| 19 | // service calls. A returned completion function releases each admission once. |
| 20 | type maintenanceGate struct { |
| 21 | mu sync.Mutex |
| 22 | active int |
| 23 | maintenance bool |
| 24 | closing bool |
| 25 | changed chan struct{} |
| 26 | } |
| 27 | |
| 28 | type maintenanceContextKey struct{} |
| 29 | |
| 30 | type maintenanceContext struct { |
| 31 | gate *maintenanceGate |
| 32 | maintenance bool |
| 33 | } |
| 34 | |
| 35 | // beginOperation admits ordinary tenant work. The caller passes the returned |
| 36 | // context to nested operations and defers the completion function on success. |
| 37 | func (t *Tenant) beginOperation(ctx context.Context) (context.Context, func(), error) { |
| 38 | return t.maintenance.beginOperation(ctx) |
| 39 | } |
| 40 | |
| 41 | // beginMaintenance grants exclusive tenant access after active work drains. |
| 42 | // The returned context admits maintenance's nested calls until completion. |
| 43 | func (t *Tenant) beginMaintenance(ctx context.Context) (context.Context, func(), error) { |
| 44 | return t.maintenance.beginMaintenance(ctx) |
| 45 | } |
| 46 | |
| 47 | func (g *maintenanceGate) waitIdle(ctx context.Context) error { |
| 48 | remaining := 0 |
| 49 | ownMaintenance := false |
| 50 | if marker, ok := ctx.Value(maintenanceContextKey{}).(maintenanceContext); ok && marker.gate == g { |
| 51 | ownMaintenance = marker.maintenance |
| 52 | if !marker.maintenance { |
| 53 | remaining = 1 |
| 54 | } |
| 55 | } |
| 56 | g.mu.Lock() |
| 57 | defer g.mu.Unlock() |
| 58 | for g.active > remaining || g.maintenance && !ownMaintenance { |
| 59 | if g.changed == nil { |
| 60 | g.changed = make(chan struct{}) |
| 61 | } |
| 62 | changed := g.changed |
| 63 | g.mu.Unlock() |
| 64 | select { |
| 65 | case <-ctx.Done(): |
| 66 | g.mu.Lock() |
| 67 | return ctx.Err() |
| 68 | case <-changed: |
| 69 | } |
| 70 | g.mu.Lock() |
| 71 | } |
| 72 | return nil |
| 73 | } |
| 74 | |
| 75 | func (g *maintenanceGate) beginOperation(ctx context.Context) (context.Context, func(), error) { |
| 76 | if ctx == nil { |
| 77 | ctx = context.Background() |
| 78 | } |
| 79 | if marker, ok := ctx.Value(maintenanceContextKey{}).(maintenanceContext); ok && marker.gate == g { |
| 80 | return ctx, func() {}, nil |
| 81 | } |
| 82 | if err := ctx.Err(); err != nil { |
| 83 | return ctx, nil, err |
| 84 | } |
| 85 | g.mu.Lock() |
| 86 | if g.closing { |
| 87 | g.mu.Unlock() |
| 88 | return ctx, nil, ErrTenantClosing |
| 89 | } |
| 90 | if g.maintenance { |
| 91 | g.mu.Unlock() |
| 92 | return ctx, nil, ErrMaintenance |
| 93 | } |
| 94 | g.active++ |
| 95 | g.mu.Unlock() |
| 96 | marked := context.WithValue(ctx, maintenanceContextKey{}, maintenanceContext{gate: g}) |
| 97 | return marked, g.endOperation(), nil |
| 98 | } |
| 99 | |
| 100 | func (g *maintenanceGate) beginMaintenance(ctx context.Context) (context.Context, func(), error) { |
| 101 | if ctx == nil { |
| 102 | ctx = context.Background() |
| 103 | } |
| 104 | if marker, ok := ctx.Value(maintenanceContextKey{}).(maintenanceContext); ok && marker.gate == g { |
| 105 | if marker.maintenance { |
| 106 | return ctx, func() {}, nil |
| 107 | } |
| 108 | return ctx, nil, ErrMaintenanceBusy |
| 109 | } |
| 110 | if err := ctx.Err(); err != nil { |
| 111 | return ctx, nil, err |
| 112 | } |
| 113 | g.mu.Lock() |
| 114 | if g.closing { |
| 115 | g.mu.Unlock() |
| 116 | return ctx, nil, ErrTenantClosing |
| 117 | } |
| 118 | if g.maintenance { |
| 119 | g.mu.Unlock() |
| 120 | return ctx, nil, ErrMaintenanceBusy |
| 121 | } |
| 122 | g.maintenance = true |
| 123 | g.signalLocked() |
| 124 | for g.active > 0 { |
| 125 | changed := g.changed |
| 126 | g.mu.Unlock() |
| 127 | select { |
| 128 | case <-ctx.Done(): |
| 129 | g.mu.Lock() |
| 130 | g.maintenance = false |
| 131 | g.signalLocked() |
| 132 | g.mu.Unlock() |
| 133 | return ctx, nil, ctx.Err() |
| 134 | case <-changed: |
| 135 | } |
| 136 | g.mu.Lock() |
| 137 | } |
| 138 | g.mu.Unlock() |
| 139 | marked := context.WithValue(ctx, maintenanceContextKey{}, maintenanceContext{gate: g, maintenance: true}) |
| 140 | return marked, g.endMaintenance(), nil |
| 141 | } |
| 142 | |
| 143 | func (g *maintenanceGate) markClosing() { |
| 144 | g.mu.Lock() |
| 145 | g.closing = true |
| 146 | g.signalLocked() |
| 147 | g.mu.Unlock() |
| 148 | } |
| 149 | |
| 150 | func (g *maintenanceGate) endOperation() func() { |
| 151 | var once sync.Once |
| 152 | return func() { |
| 153 | once.Do(func() { |
| 154 | g.mu.Lock() |
| 155 | if g.active > 0 { |
| 156 | g.active-- |
| 157 | } |
| 158 | g.signalLocked() |
| 159 | g.mu.Unlock() |
| 160 | }) |
| 161 | } |
| 162 | } |
| 163 | |
| 164 | func (g *maintenanceGate) endMaintenance() func() { |
| 165 | var once sync.Once |
| 166 | return func() { |
| 167 | once.Do(func() { |
| 168 | g.mu.Lock() |
| 169 | g.maintenance = false |
| 170 | g.signalLocked() |
| 171 | g.mu.Unlock() |
| 172 | }) |
| 173 | } |
| 174 | } |
| 175 | |
| 176 | // watch returns a channel that closes on the next gate change, so a |
| 177 | // long-lived reader can stop when maintenance or shutdown begins. |
| 178 | func (g *maintenanceGate) watch() <-chan struct{} { |
| 179 | g.mu.Lock() |
| 180 | defer g.mu.Unlock() |
| 181 | if g.changed == nil { |
| 182 | g.changed = make(chan struct{}) |
| 183 | } |
| 184 | return g.changed |
| 185 | } |
| 186 | |
| 187 | // blocked reports whether new operations are being refused. |
| 188 | func (g *maintenanceGate) blocked() bool { |
| 189 | g.mu.Lock() |
| 190 | defer g.mu.Unlock() |
| 191 | return g.maintenance || g.closing |
| 192 | } |
| 193 | |
| 194 | func (g *maintenanceGate) signalLocked() { |
| 195 | if g.changed != nil { |
| 196 | close(g.changed) |
| 197 | } |
| 198 | g.changed = make(chan struct{}) |
| 199 | } |