wait

Seven callers share GPU batches of 16 floats, in arrival order. A batch is a slice whose length is the part reserved so far. The open batch is the only item in turn, so holding it is your turn to reserve. A caller that doesn't fit takes a place in next before it passes the turn on: smaller callers behind it fill the batch, and it goes first in the next one. Each caller's view is a subslice of its batch, not a copy.

package main

import (
	"context"
	"fmt"
	"log"
	"sync"

	"blake.io/wait"
)

func main() {
	ctx := context.Background()
	var turn, next wait.List[[]float32]
	turn.Add(make([]float32, 0, 16))

	views := make([][]float32, 7)
	var wg sync.WaitGroup
	submit := func(i, n int) {
		t := turn.Take(ctx) // Take never blocks, so the call order is the arrival order
		wg.Go(func() {
			l := &turn
			b, err := t.Value()
			for err == nil && n > cap(b)-len(b) {
				l = &next
				aside := l.Take(ctx) // keep a place before passing the turn on
				t.Done()
				t = aside
				b, err = t.Value()
			}
			if err != nil {
				log.Fatal(err)
			}
			views[i] = b[len(b) : len(b)+n]
			l.Add(b[:len(b)+n]) // pass the turn on, n more reserved
		})
	}

	// run waits its turn, so everyone ahead of it has reserved or stepped
	// aside, fills the batch, and opens the next one: first to those who
	// stepped aside, then to everyone else.
	run := func(k int) {
		b, err := turn.Take(ctx).Value()
		if err != nil {
			log.Fatal(err)
		}
		mine := next.Take(ctx) // behind everyone who stepped aside
		fmt.Printf("batch %d runs %d floats\n", k, len(b))
		for i := range b {
			b[i] = float32(100*k + i)
		}
		next.Add(make([]float32, 0, 16))
		if b, err = mine.Value(); err != nil {
			log.Fatal(err)
		}
		turn.Add(b)
	}

	for i, n := range []int{6, 8, 12, 2, 5, 3} {
		submit(i, n)
	}
	run(1)
	submit(6, 4) // arrives after batch 1
	run(2)
	run(3)
	wg.Wait()
	for i, v := range views {
		fmt.Printf("caller %d: %v\n", i+1, v)
	}

}

Output:
batch 1 runs 16 floats
batch 2 runs 15 floats
batch 3 runs 9 floats
caller 1: [100 101 102 103 104 105]
caller 2: [106 107 108 109 110 111 112 113]
caller 3: [200 201 202 203 204 205 206 207 208 209 210 211]
caller 4: [114 115]
caller 5: [300 301 302 303 304]
caller 6: [212 213 214]
caller 7: [305 306 307 308]

Downloads share a link's 8 connections and 100 Mbps, strictly in the order they ask. The second download's connections fit beside the first's, but its bandwidth doesn't, so it waits for the first to finish.

The ctx here is never done, so Value never fails. Code whose ctx can end must pass the turn on, with what came back, when it gives up.

package main

import (
	"context"
	"fmt"

	"blake.io/wait"
)

func main() {
	type demand struct{ conns, mbps int }
	ctx := context.Background()

	link := demand{8, 100}

	// The link's free capacity is the only item in free, so holding it is
	// your turn.
	var free, returned wait.List[demand]
	free.Add(link)

	// acquire reports false for more than the whole link, which would hold
	// the turn forever, and everyone behind it with it.
	acquire := func(d demand) bool {
		if d.conns > link.conns || d.mbps > link.mbps {
			return false
		}
		f, _ := free.Take(ctx).Value() // wait your turn
		for d.conns > f.conns || d.mbps > f.mbps {
			r, _ := returned.Take(ctx).Value() // hold the turn until enough is back
			f.conns, f.mbps = f.conns+r.conns, f.mbps+r.mbps
		}
		free.Add(demand{f.conns - d.conns, f.mbps - d.mbps}) // pass the turn on
		return true
	}

	if !acquire(demand{16, 10}) {
		fmt.Println("16 connections is more than the link has")
	}
	for _, d := range []demand{{4, 60}, {2, 50}} {
		acquire(d)
		go func() {
			defer returned.Add(d)
			fmt.Printf("downloading on %d connections at %d Mbps\n", d.conns, d.mbps)
		}()
	}
	acquire(link) // the whole link: waits for both downloads

}

Output:
16 connections is more than the link has
downloading on 4 connections at 60 Mbps
downloading on 2 connections at 50 Mbps

This example loads copies of a model only when callers need them, up to max. A caller that finds no idle copy starts a load if there is room for another copy, then waits in line. Each copy, whether newly loaded or given back, goes to the caller that has waited longest.

Calling TryTake before Take is fair, because TryTake cannot take a copy that a waiting caller is owed. If a copy is given back between TryTake and Take, the caller may start a load it didn't need, but there are still at most max copies.

A load can fail for reasons the process doesn't control, such as other processes filling the GPU's memory. The goroutine that started the load tries again rather than giving up its room, because callers already in line are waiting for that copy.

package main

import (
	"context"
	"errors"
	"fmt"
	"sync/atomic"
	"time"

	"blake.io/wait"
)

func main() {
	const max = 2
	ctx := context.Background()
	var copies wait.List[string]

	// load stands in for reading the model into GPU memory. Every other
	// load finds the memory full.
	var loads atomic.Int64
	load := func() error {
		if loads.Add(1)%2 == 0 {
			return errors.New("out of GPU memory")
		}
		return nil
	}

	var misses atomic.Int64 // each of the first max misses starts a load
	get := func() wait.Ticket[string] {
		if t, ok := copies.TryTake(); ok {
			return t // an idle copy
		}
		if n := misses.Add(1); n <= max {
			go func() {
				for err := load(); err != nil; err = load() {
					fmt.Printf("loading copy %d: %v; retrying\n", n, err)
					time.Sleep(100 * time.Millisecond)
				}
				copies.Add(fmt.Sprintf("copy %d", n))
			}()
		}
		return copies.Take(ctx)
	}
	use := func(caller string, t wait.Ticket[string]) {
		name, _ := t.Value() // ctx is never done, so Value never fails
		fmt.Println(caller, "running on", name)
	}

	a := get() // loads copy 1
	use("a", a)
	b := get() // copy 1 is in use: loads copy 2
	use("b", b)
	c := get() // at max: waits in line
	a.Done()   // copy 1 goes to c
	use("c", c)
	b.Done()
	c.Done()

}

Output:
a running on copy 1
loading copy 2: out of GPU memory; retrying
b running on copy 2
c running on copy 1

A caller that finds its connection broken redials it before the deferred Done gives it back, so the next caller gets a fresh one. Dropping the conn instead would shrink the List for good, and once every conn is dropped, callers wait forever.

package main

import (
	"context"
	"fmt"
	"log"
	"sync"

	"blake.io/wait"
)

// A conn is a pooled connection that redials itself. It dials in the
// background when it is made and again after redial, so whoever takes it
// from a List usually finds it connected.
type conn struct {
	addr   string
	dials  int
	dialed func() (string, error) // waits for the dial in flight
}

func newConn(addr string) *conn {
	c := &conn{addr: addr}
	c.redial()
	return c
}

func (c *conn) redial() {
	c.dials++
	session := fmt.Sprintf("%s#%d", c.addr, c.dials) // stands in for a real dial
	c.dialed = sync.OnceValues(func() (string, error) { return session, nil })
	go c.dialed()
}

// A caller that finds its connection broken redials it before the deferred
// Done gives it back, so the next caller gets a fresh one. Dropping the
// conn instead would shrink the List for good, and once every conn is
// dropped, callers wait forever.
func main() {
	var conns wait.List[*conn]
	conns.Add(newConn("db"))

	use := func(broken bool) {
		t := conns.Take(context.Background())
		defer t.Done()
		c, err := t.Value()
		if err != nil {
			log.Fatal(err)
		}
		session, err := c.dialed()
		if err != nil {
			c.redial()
			return
		}
		fmt.Println("using", session)
		if broken {
			c.redial()
		}
	}

	use(true) // db#1 breaks while in use
	use(false)

}

Output:
using db#1
using db#2