| 1 | package host |
| 2 | |
| 3 | import ( |
| 4 | "context" |
| 5 | "errors" |
| 6 | "sort" |
| 7 | "sync" |
| 8 | |
| 9 | "github.com/carlosframework/platform/internal/hostbox" |
| 10 | "github.com/carlosframework/platform/internal/hostguard" |
| 11 | ) |
| 12 | |
| 13 | type claimProducer struct { |
| 14 | cancel context.CancelFunc |
| 15 | joined chan struct{} |
| 16 | err error |
| 17 | } |
| 18 | type completeClaimProducer func(context.Context, hostguard.FenceTarget, string, hostbox.ProducerCompletion) (hostbox.ProducerCompletion, error) |
| 19 | |
| 20 | type ClaimProducers struct { |
| 21 | closeWriters func(context.Context, hostguard.FenceTarget) error |
| 22 | completed hostbox.ProducerCompletion |
| 23 | complete completeClaimProducer |
| 24 | mu sync.Mutex |
| 25 | target hostguard.FenceTarget |
| 26 | processStart string |
| 27 | fence string |
| 28 | drained bool |
| 29 | // writersStopped closes admission to every producer but the retained |
| 30 | // ones (StopWriters); the full drain still cancels those. |
| 31 | writersStopped bool |
| 32 | workers map[string]*claimProducer |
| 33 | Record func(context.Context, hostguard.FenceTarget, string) error |
| 34 | } |
| 35 | |
| 36 | func NewClaimProducers(target hostguard.FenceTarget, start string, record func(context.Context, hostguard.FenceTarget, string) error, complete ...completeClaimProducer) (*ClaimProducers, error) { |
| 37 | if target.Host == "" || target.ControllerID == "" || start == "" || record == nil || len(complete) > 1 || (len(complete) == 1 && complete[0] == nil) { |
| 38 | return nil, ErrClaimStopUnavailable |
| 39 | } |
| 40 | p := &ClaimProducers{target: target, processStart: start, Record: record, workers: map[string]*claimProducer{}} |
| 41 | if len(complete) == 1 { |
| 42 | p.complete = complete[0] |
| 43 | } |
| 44 | return p, nil |
| 45 | } |
| 46 | func (p *ClaimProducers) Launch(ctx context.Context, target hostguard.FenceTarget, id string, fn func(context.Context) error) error { |
| 47 | if p == nil || fn == nil || !stopComponent(id) { |
| 48 | return ErrClaimStopUnavailable |
| 49 | } |
| 50 | p.mu.Lock() |
| 51 | defer p.mu.Unlock() |
| 52 | if p.fence != "" || !hostbox.SameLaunchOwner(p.target, target) { |
| 53 | return ErrClaimStopIdentity |
| 54 | } |
| 55 | if p.writersStopped && !claimRetainedProducer(id) { |
| 56 | return ErrClaimWritersStopped |
| 57 | } |
| 58 | if err := p.compact(ctx); err != nil { |
| 59 | return err |
| 60 | } |
| 61 | if _, ok := p.workers[id]; ok { |
| 62 | return ErrClaimStopIdentity |
| 63 | } |
| 64 | if err := p.Record(ctx, target, id); err != nil { |
| 65 | return err |
| 66 | } |
| 67 | child, cancel := context.WithCancel(ctx) |
| 68 | w := &claimProducer{cancel: cancel, joined: make(chan struct{})} |
| 69 | p.workers[id] = w |
| 70 | go func() { p.finish(w, fn(child)) }() |
| 71 | return nil |
| 72 | } |
| 73 | func (p *ClaimProducers) DrainClaim(ctx context.Context, r StopRequest, fence string) (DrainAcknowledgement, error) { |
| 74 | if p == nil || fence == "" { |
| 75 | return DrainAcknowledgement{}, ErrClaimStopUnavailable |
| 76 | } |
| 77 | p.mu.Lock() |
| 78 | if p.target.Host != r.Host || p.target.Incarnation != r.Incarnation || p.target.Stamp != r.InstallationStamp || p.target.Session != r.WriterSession || p.target.Box != r.Box || p.target.Enrollment != r.Enrollment || p.target.BootID != r.BootID || p.target.ControllerID != r.ControllerID || (p.fence != "" && p.fence != fence && !p.drained) { |
| 79 | p.mu.Unlock() |
| 80 | return DrainAcknowledgement{}, ErrClaimStopIdentity |
| 81 | } |
| 82 | p.fence, p.drained = fence, false |
| 83 | if p.closeWriters != nil { |
| 84 | if err := p.closeWriters(ctx, p.target); err != nil { |
| 85 | p.mu.Unlock() |
| 86 | return DrainAcknowledgement{}, err |
| 87 | } |
| 88 | } |
| 89 | workers := map[string]*claimProducer{} |
| 90 | for id, w := range p.workers { |
| 91 | workers[id] = w |
| 92 | w.cancel() |
| 93 | } |
| 94 | p.mu.Unlock() |
| 95 | for _, w := range workers { |
| 96 | select { |
| 97 | case <-ctx.Done(): |
| 98 | return DrainAcknowledgement{}, ctx.Err() |
| 99 | case <-w.joined: |
| 100 | } |
| 101 | if w.err != nil && !errors.Is(w.err, context.Canceled) { |
| 102 | return DrainAcknowledgement{}, w.err |
| 103 | } |
| 104 | } |
| 105 | p.mu.Lock() |
| 106 | defer p.mu.Unlock() |
| 107 | if err := p.compact(ctx); err != nil { |
| 108 | return DrainAcknowledgement{}, err |
| 109 | } |
| 110 | ids := make([]string, 0, len(p.workers)) |
| 111 | for id := range p.workers { |
| 112 | ids = append(ids, id) |
| 113 | } |
| 114 | p.drained = true |
| 115 | return DrainClaimAcknowledgement(fence, r.ControllerID, p.processStart, ids, p.completed), nil |
| 116 | } |
| 117 | |
| 118 | func (p *ClaimProducers) Register(ctx context.Context, target hostguard.FenceTarget, id string) (context.Context, func(error), error) { |
| 119 | if p == nil || !stopComponent(id) { |
| 120 | return nil, nil, ErrClaimStopUnavailable |
| 121 | } |
| 122 | p.mu.Lock() |
| 123 | defer p.mu.Unlock() |
| 124 | if p.fence != "" || !hostbox.SameLaunchOwner(p.target, target) { |
| 125 | return nil, nil, ErrClaimStopIdentity |
| 126 | } |
| 127 | if p.writersStopped && !claimRetainedProducer(id) { |
| 128 | return nil, nil, ErrClaimWritersStopped |
| 129 | } |
| 130 | if err := p.compact(ctx); err != nil { |
| 131 | return nil, nil, err |
| 132 | } |
| 133 | if _, ok := p.workers[id]; ok { |
| 134 | return nil, nil, ErrClaimStopIdentity |
| 135 | } |
| 136 | if err := p.Record(ctx, target, id); err != nil { |
| 137 | return nil, nil, err |
| 138 | } |
| 139 | child, cancel := context.WithCancel(ctx) |
| 140 | w := &claimProducer{cancel: cancel, joined: make(chan struct{})} |
| 141 | p.workers[id] = w |
| 142 | var once sync.Once |
| 143 | finish := func(err error) { once.Do(func() { p.finish(w, err) }) } |
| 144 | return child, finish, nil |
| 145 | } |
| 146 | |
| 147 | // StopWriters is the writers-only drain (claim Plan 11, R2): it closes |
| 148 | // admission to every producer but the retained ones, cancels the rest (the |
| 149 | // role launch's producers and the schedule runner's timer deliveries) and |
| 150 | // joins them. The replicator, the generation publisher and the files |
| 151 | // publisher keep running; DrainClaim, root's full drain, still ends them. |
| 152 | // It returns the ids it joined, sorted. Repeating it is safe. |
| 153 | func (p *ClaimProducers) StopWriters(ctx context.Context, target hostguard.FenceTarget) ([]string, error) { |
| 154 | if p == nil || ctx == nil { |
| 155 | return nil, ErrClaimStopUnavailable |
| 156 | } |
| 157 | p.mu.Lock() |
| 158 | if p.fence != "" || !hostbox.SameLaunchOwner(p.target, target) { |
| 159 | p.mu.Unlock() |
| 160 | return nil, ErrClaimStopIdentity |
| 161 | } |
| 162 | p.writersStopped = true |
| 163 | if p.closeWriters != nil { |
| 164 | if err := p.closeWriters(ctx, p.target); err != nil { |
| 165 | p.mu.Unlock() |
| 166 | return nil, err |
| 167 | } |
| 168 | } |
| 169 | workers := map[string]*claimProducer{} |
| 170 | for id, w := range p.workers { |
| 171 | if !claimRetainedProducer(id) { |
| 172 | workers[id] = w |
| 173 | w.cancel() |
| 174 | } |
| 175 | } |
| 176 | p.mu.Unlock() |
| 177 | ids := make([]string, 0, len(workers)) |
| 178 | for id, w := range workers { |
| 179 | select { |
| 180 | case <-ctx.Done(): |
| 181 | return nil, ctx.Err() |
| 182 | case <-w.joined: |
| 183 | } |
| 184 | if w.err != nil && !errors.Is(w.err, context.Canceled) { |
| 185 | return nil, w.err |
| 186 | } |
| 187 | ids = append(ids, id) |
| 188 | } |
| 189 | p.mu.Lock() |
| 190 | defer p.mu.Unlock() |
| 191 | if err := p.compact(ctx); err != nil { |
| 192 | return nil, err |
| 193 | } |
| 194 | sort.Strings(ids) |
| 195 | return ids, nil |
| 196 | } |
| 197 | |
| 198 | // Running reports the producers that have not finished, sorted. |
| 199 | func (p *ClaimProducers) Running() []string { |
| 200 | p.mu.Lock() |
| 201 | defer p.mu.Unlock() |
| 202 | ids := make([]string, 0, len(p.workers)) |
| 203 | for id, w := range p.workers { |
| 204 | select { |
| 205 | case <-w.joined: |
| 206 | default: |
| 207 | ids = append(ids, id) |
| 208 | } |
| 209 | } |
| 210 | sort.Strings(ids) |
| 211 | return ids |
| 212 | } |
| 213 | |
| 214 | type ClaimProducerSet struct { |
| 215 | CloseWriters func(context.Context, hostguard.FenceTarget) error |
| 216 | mu sync.Mutex |
| 217 | items []*ClaimProducers |
| 218 | } |
| 219 | |
| 220 | func (s *ClaimProducerSet) Add(p *ClaimProducers) error { |
| 221 | if p == nil { |
| 222 | return ErrClaimStopUnavailable |
| 223 | } |
| 224 | s.mu.Lock() |
| 225 | defer s.mu.Unlock() |
| 226 | for _, old := range s.items { |
| 227 | if hostbox.SameLaunchOwner(old.target, p.target) { |
| 228 | return ErrClaimStopIdentity |
| 229 | } |
| 230 | } |
| 231 | s.items = append(s.items, p) |
| 232 | return nil |
| 233 | } |
| 234 | func (s *ClaimProducerSet) For(t hostguard.FenceTarget) (*ClaimProducers, error) { |
| 235 | s.mu.Lock() |
| 236 | defer s.mu.Unlock() |
| 237 | for _, p := range s.items { |
| 238 | if hostbox.SameLaunchOwner(p.target, t) { |
| 239 | return p, nil |
| 240 | } |
| 241 | } |
| 242 | return nil, ErrClaimStopUnavailable |
| 243 | } |
| 244 | func (s *ClaimProducerSet) DrainClaim(ctx context.Context, r StopRequest, fence string) (DrainAcknowledgement, error) { |
| 245 | t := hostguard.FenceTarget{Host: r.Host, Incarnation: r.Incarnation, Stamp: r.InstallationStamp, Session: r.WriterSession, Box: r.Box, Enrollment: r.Enrollment, BootID: r.BootID, ControllerID: r.ControllerID} |
| 246 | p, err := s.For(t) |
| 247 | if err != nil { |
| 248 | return DrainAcknowledgement{}, err |
| 249 | } |
| 250 | return p.DrainClaim(ctx, r, fence) |
| 251 | } |
| 252 | |
| 253 | func (s *ClaimProducerSet) Ensure(target hostguard.FenceTarget, start string, record func(context.Context, hostguard.FenceTarget, string) error, complete ...completeClaimProducer) (*ClaimProducers, error) { |
| 254 | s.mu.Lock() |
| 255 | defer s.mu.Unlock() |
| 256 | for _, old := range s.items { |
| 257 | if hostbox.SameLaunchOwner(old.target, target) { |
| 258 | if old.processStart != start { |
| 259 | return nil, ErrClaimStopIdentity |
| 260 | } |
| 261 | return old, nil |
| 262 | } |
| 263 | } |
| 264 | p, err := NewClaimProducers(target, start, record, complete...) |
| 265 | if err != nil { |
| 266 | return nil, err |
| 267 | } |
| 268 | p.closeWriters = s.CloseWriters |
| 269 | s.items = append(s.items, p) |
| 270 | return p, nil |
| 271 | } |
| 272 | |
| 273 | func (p *ClaimProducers) finish(w *claimProducer, err error) { |
| 274 | p.mu.Lock() |
| 275 | defer p.mu.Unlock() |
| 276 | w.err = err |
| 277 | w.cancel() |
| 278 | close(w.joined) |
| 279 | ctx, cancel := context.WithTimeout(context.Background(), hostbox.RequestLimit) |
| 280 | defer cancel() |
| 281 | // A failed durable write retains the joined worker. Admission and drain |
| 282 | // retry it before accepting more work or acknowledging completion. |
| 283 | _ = p.compact(ctx) |
| 284 | } |
| 285 | |
| 286 | // compact runs under mu so the journal and the memory summary advance together. |
| 287 | func (p *ClaimProducers) compact(ctx context.Context) error { |
| 288 | if p.complete == nil { |
| 289 | return nil |
| 290 | } |
| 291 | for id, w := range p.workers { |
| 292 | select { |
| 293 | case <-w.joined: |
| 294 | default: |
| 295 | continue |
| 296 | } |
| 297 | if w.err != nil && !errors.Is(w.err, context.Canceled) { |
| 298 | return w.err |
| 299 | } |
| 300 | next, err := p.complete(ctx, p.target, id, p.completed) |
| 301 | if err != nil { |
| 302 | return err |
| 303 | } |
| 304 | if !next.Valid() || next.Count != p.completed.Count+1 { |
| 305 | return ErrClaimStopIdentity |
| 306 | } |
| 307 | p.completed = next |
| 308 | delete(p.workers, id) |
| 309 | } |
| 310 | return nil |
| 311 | } |
| 312 | |
amadan