Pull nothing at capacity 0, take a new capacity at the next pull, and keep a node's pause on its pause key, so a paused node stays paused across a restart (issue 1887 A2b)

2026-10-23 15:16:46 +02:00
parent a1d73bdcf0
commit 89aa6fc17c
3 changed files with 168 additions and 9 deletions
@@ -2,6 +2,7 @@ package buildstep

 import (
 	"context"
+	"encoding/json"
 	"errors"
 	"fmt"
 	"sync"
@@ -62,6 +63,10 @@ type Puller struct {
 	PullWait  time.Duration
 	Heartbeat time.Duration
 	Logf      func(format string, args ...any)
+	// Load reads one key of the module's steps state (nil and no error when it is absent). When set, Pull reads the
+	// node's pause key (delivery.PauseKey) before it pulls, and Pause and Resume write it with Put, so a pause outlives
+	// a restart.
+	Load func(key string) ([]byte, error)

 	mu     sync.Mutex
 	inHand int
@@ -89,19 +94,58 @@ func (p *Puller) Wake() {
 	}
 }

-// Pause stops the puller taking asks; the asks it holds run to their end.
-func (p *Puller) Pause() {
-	p.mu.Lock()
-	p.paused = true
-	p.mu.Unlock()
-}
+// Pause stops the puller taking asks; the asks it holds run to their end. by names who paused it. With Load set the
+// pause is written on the node's pause key first, and an error there leaves the puller as it was.
+func (p *Puller) Pause(by string) error { return p.setPaused(true, by) }

-// Resume lets the puller take asks again.
-func (p *Puller) Resume() {
+// Resume lets the puller take asks again, written as Pause's is.
+func (p *Puller) Resume(by string) error { return p.setPaused(false, by) }
+
+func (p *Puller) setPaused(paused bool, by string) error {
+	if p.Load != nil {
+		key, err := delivery.PauseKey(p.Node)
+		if err != nil {
+			return err
+		}
+		raw, err := delivery.EncodePause(delivery.Pause{Paused: paused, By: by, At: now()})
+		if err != nil {
+			return err
+		}
+		if _, err := p.Put(key, json.RawMessage(raw)); err != nil {
+			return fmt.Errorf("buildstep: writing the pause of %s: %w", p.Node, err)
+		}
+	}
 	p.mu.Lock()
-	p.paused = false
+	p.paused = paused
 	p.mu.Unlock()
 	p.Wake()
+	return nil
+}
+
+// restore reads the node's pause key, when Load is set.
+func (p *Puller) restore() error {
+	if p.Load == nil {
+		return nil
+	}
+	key, err := delivery.PauseKey(p.Node)
+	if err != nil {
+		return err
+	}
+	raw, err := p.Load(key)
+	if err != nil {
+		return fmt.Errorf("buildstep: reading the pause of %s: %w", p.Node, err)
+	}
+	if raw == nil {
+		return nil
+	}
+	v, err := delivery.DecodePause(raw)
+	if err != nil {
+		return fmt.Errorf("buildstep: the pause of %s: %w", p.Node, err)
+	}
+	p.mu.Lock()
+	p.paused = v.Paused
+	p.mu.Unlock()
+	return nil
 }

 // Paused says whether the puller takes no asks.
@@ -136,6 +180,9 @@ func (p *Puller) Pull(ctx context.Context) error {
 	if p.Source == nil || p.Put == nil || p.Run == nil {
 		return errors.New("buildstep: a puller needs its source, its steps' put and its run")
 	}
+	if err := p.restore(); err != nil {
+		return err // pulling on a pause that cannot be read could build on a node paused on purpose
+	}
 	wait := p.PullWait
 	if wait <= 0 {
 		wait = DefaultPullWait
@@ -153,3 +153,81 @@ func TestPullDropsWhatIsNoAsk(t *testing.T) {
 		t.Fatal("a wrong ask ran")
 	}
 }
+
+// load reads the recorded steps state's raw values: the pause key is not a step, so it is kept apart.
+type kv struct {
+	mu   sync.Mutex
+	vals map[string][]byte
+}
+
+func (k *kv) put(key string, value any) (uint64, error) {
+	raw, _ := value.(json.RawMessage)
+	k.mu.Lock()
+	defer k.mu.Unlock()
+	if k.vals == nil {
+		k.vals = map[string][]byte{}
+	}
+	k.vals[key] = append([]byte(nil), raw...)
+	return 1, nil
+}
+
+func (k *kv) load(key string) ([]byte, error) {
+	k.mu.Lock()
+	defer k.mu.Unlock()
+	return k.vals[key], nil
+}
+
+// 1887/A2b, 1887/AC-zero: at capacity 0 the puller takes nothing; a capacity raised takes effect at the next pull.
+func TestPullAtCapacityZeroTakesNothing(t *testing.T) {
+	src, st := NewFakeAsks(), &kv{}
+	_ = src.PublishAsk(testAsk(delivery.AskDelivery))
+	var mu sync.Mutex
+	capacity := 0
+	p := &Puller{Source: src, Language: delivery.AskGo, Node: "novox", Put: st.put, PullWait: 10 * time.Millisecond,
+		Capacity: func() int { mu.Lock(); defer mu.Unlock(); return capacity },
+		Run:      func(ctx context.Context, a delivery.StepAsk) (any, error) { return nil, nil }}
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	go p.Pull(ctx)
+	time.Sleep(50 * time.Millisecond)
+	if src.Waiting(delivery.AskDelivery) != 1 {
+		t.Fatal("a node at capacity 0 took an ask")
+	}
+	mu.Lock()
+	capacity = 1
+	mu.Unlock()
+	p.Wake()
+	waitFor(t, "the ask taken at capacity 1", func() bool { a, _, _ := src.Fates(); return len(a) == 1 })
+}
+
+// 1887/A2b: a paused puller takes nothing, its pause is kept on the pause key, and a puller started again on that state
+// stays paused until resumed.
+func TestPullPauseOutlivesARestart(t *testing.T) {
+	src, st := NewFakeAsks(), &kv{}
+	run := func(ctx context.Context, a delivery.StepAsk) (any, error) { return nil, nil }
+	p := &Puller{Source: src, Language: delivery.AskGo, Node: "shanks", Put: st.put, Load: st.load, PullWait: 10 * time.Millisecond,
+		Capacity: func() int { return 2 }, Run: run}
+	if err := p.Pause("the operator"); err != nil {
+		t.Fatal(err)
+	}
+	key, _ := delivery.PauseKey("shanks")
+	raw, _ := st.load(key)
+	if v, err := delivery.DecodePause(raw); err != nil || !v.Paused || v.By != "the operator" {
+		t.Fatalf("the pause key holds %s (%v)", raw, err)
+	}
+	// The restart: a new puller on the same state.
+	again := &Puller{Source: src, Language: delivery.AskGo, Node: "shanks", Put: st.put, Load: st.load, PullWait: 10 * time.Millisecond,
+		Capacity: func() int { return 2 }, Run: run}
+	_ = src.PublishAsk(testAsk(delivery.AskDelivery))
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	go again.Pull(ctx)
+	time.Sleep(50 * time.Millisecond)
+	if !again.Paused() || src.Waiting(delivery.AskDelivery) != 1 {
+		t.Fatalf("restarted: paused %v, waiting %d; want paused and 1", again.Paused(), src.Waiting(delivery.AskDelivery))
+	}
+	if err := again.Resume("the operator"); err != nil {
+		t.Fatal(err)
+	}
+	waitFor(t, "the ask taken after the resume", func() bool { a, _, _ := src.Fates(); return len(a) == 1 })
+}
@@ -373,3 +373,37 @@ func ParseAskStepKey(key string) (node, askID string, err error) {
 	}
 	return node, askID, nil
 }
+
+// PauseKey is the key, in a build module's StepsBucket, that keeps whether the module on node takes asks:
+// `paused.<node>` (issue 1887's design limit 2: every holder pulls from one shared worker, so a node's pause is its holder
+// not fetching, kept here so a restart keeps it). A pause key is never an ask's step key: its second part is a node.
+func PauseKey(node string) (string, error) {
+	if err := checkName("node", node); err != nil {
+		return "", err
+	}
+	return "paused." + node, nil
+}
+
+// Pause is the value of a pause key: whether the node's module takes no asks, who said so and when.
+type Pause struct {
+	V      int       `json:"v"`
+	Paused bool      `json:"paused"`
+	By     string    `json:"by,omitempty"`
+	At     time.Time `json:"at"`
+}
+
+// EncodePause is p as JSON at Version.
+func EncodePause(p Pause) ([]byte, error) {
+	p.V = Version
+	p.At = p.At.UTC()
+	return json.Marshal(p)
+}
+
+// DecodePause reads a pause value; it refuses an unknown version (ErrUnknownVersion).
+func DecodePause(data []byte) (Pause, error) {
+	var p Pause
+	if err := decode(data, &p); err != nil {
+		return Pause{}, err
+	}
+	return p, nil
+}