Initial implementation of buildkite-schduler.

cmd/buildkite-scheduler/buildkiteapi.go

@@ -0,0 +1,89 @@
+package main
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"net/http"
+	"strings"
+
+	"github.com/buildkite/go-buildkite/v3/buildkite"
+	"github.com/shurcooL/graphql"
+)
+
+type AgentMetrics struct {
+	// XXX: There are more fields
+	Organization struct {
+		Slug string `json:"slug,omitempty"`
+	} `json:"organization"`
+}
+
+// buildkite API docs: https://buildkite.com/docs/apis/agent-api/metrics
+func getMetrics(ctx context.Context, client *http.Client, agentToken string) (*AgentMetrics, error) {
+	req, err := http.NewRequestWithContext(ctx, "GET", "https://agent.buildkite.com/v3/metrics", nil)
+	if err != nil {
+		return nil, err
+	}
+	req.Header.Add("Authorization", "Token "+agentToken)
+	resp, err := client.Do(req)
+	if err != nil {
+		return nil, err
+	}
+	defer resp.Body.Close()
+	if resp.StatusCode != 200 {
+		return nil, fmt.Errorf("metrics query returned status %d", resp.StatusCode)
+	}
+	var m AgentMetrics
+	err = json.NewDecoder(resp.Body).Decode(&m)
+	return &m, err
+}
+
+func queryJobs(ctx context.Context, qlClient *graphql.Client, orgSlug string) ([]*buildkite.Job, error) {
+	var q struct {
+		Organization struct {
+			Pipelines struct {
+				Edges []struct {
+					Node struct {
+						Slug   string
+						Builds struct {
+							Edges []struct {
+								Node struct {
+									Jobs struct {
+										Edges []struct {
+											Node struct {
+												JobTypeCommand struct {
+													UUID  string
+													State string
+												} `graphql:"... on JobTypeCommand"`
+											}
+										}
+									} `graphql:"jobs(state: SCHEDULED, first: 500)"`
+								}
+							}
+						} `graphql:"builds(first: 500)"`
+					}
+				}
+			} `graphql:"pipelines(first: 500)"`
+		} `graphql:"organization(slug: $orgSlug)"`
+	}
+	vars := map[string]interface{}{
+		"orgSlug": orgSlug,
+	}
+	if err := qlClient.Query(ctx, &q, vars); err != nil {
+		return nil, err
+	}
+	jobs := []*buildkite.Job{}
+	for _, p := range q.Organization.Pipelines.Edges {
+		for _, b := range p.Node.Builds.Edges {
+			for _, j := range b.Node.Jobs.Edges {
+				state := strings.ToLower(j.Node.JobTypeCommand.State)
+				jobs = append(jobs, &buildkite.Job{
+					// XXX: load more fields
+					ID:    &j.Node.JobTypeCommand.UUID,
+					State: &state,
+				})
+			}
+		}
+	}
+	return jobs, nil
+}

cmd/buildkite-scheduler/main.go

@@ -0,0 +1,84 @@
+package main
+
+import (
+	"context"
+	"crypto/rand"
+	"encoding/hex"
+	"fmt"
+	"log"
+	"net"
+	"net/http"
+
+	"github.com/spf13/cobra"
+	"github.com/spf13/pflag"
+	"golang.org/x/sync/errgroup"
+	"namespacelabs.dev/foundation/internal/cli/fncobra"
+	"namespacelabs.dev/foundation/internal/providers/nscloud/api"
+)
+
+var (
+	rendezvouzEndpoint = pflag.String("rendezvous_endpoint", "rendezvous.namespace.so:5000", "namespace proxy endpoint")
+	webhookSecret      = pflag.String("webhook_secret", "", "webhook secret")
+	agentToken         = pflag.String("agent_token", "", "agent token")
+	apiToken           = pflag.String("api_token", "", "api token")
+)
+
+func init() {
+	log.SetFlags(log.Lshortfile | log.Lmicroseconds)
+	api.SetupFlags("", pflag.CommandLine, false)
+	api.Register()
+}
+
+func main() {
+	fncobra.DoMain(fncobra.MainOpts{
+		Name: "buildkite-scheduler",
+		RegisterCommands: func(root *cobra.Command) {
+			api.SetupFlags("", root.PersistentFlags(), false)
+			root.RunE = func(cmd *cobra.Command, args []string) error {
+				eg, ctx := errgroup.WithContext(cmd.Context())
+
+				sched := newScheduler(*agentToken, *apiToken)
+				eg.Go(func() error {
+					return sched.runWorker(ctx)
+				})
+
+				if *apiToken != "" {
+					eg.Go(func() error {
+						return sched.runPoller(ctx)
+					})
+				}
+
+				eg.Go(func() error {
+					listener := StartProxyListener(ctx, *rendezvouzEndpoint, func(endpoint string) {
+						fmt.Printf("Set webhook URL to http://%s/webhook\n", endpoint)
+					})
+					if *webhookSecret == "" {
+						if secret, err := generateSecret(); err != nil {
+							log.Fatalf("could not generate webhook token: %v", err)
+						} else {
+							*webhookSecret = secret
+							fmt.Printf("Set webhook token to %s\n", secret)
+						}
+					}
+					webhookSecret := []byte(*webhookSecret)
+					handler := webhookHandler(webhookSecret, sched.onWebHook)
+					srv := &http.Server{
+						Handler:     handler,
+						BaseContext: func(net.Listener) context.Context { return ctx },
+					}
+					return srv.Serve(listener)
+				})
+
+				return eg.Wait()
+			}
+		},
+	})
+}
+
+func generateSecret() (string, error) {
+	bytes := make([]byte, 16)
+	if _, err := rand.Read(bytes); err != nil {
+		return "", err
+	}
+	return hex.EncodeToString(bytes), nil
+}

cmd/buildkite-scheduler/proxylistener.go

@@ -0,0 +1,54 @@
+package main
+
+import (
+	"context"
+	"io"
+	"net"
+
+	"namespacelabs.dev/breakpoint/pkg/quicproxyclient"
+)
+
+func StartProxyListener(baseCtx context.Context, rendezvouz string, onAllocation func(endpoint string)) net.Listener {
+	ctx, cancel := context.WithCancel(baseCtx)
+	l := &proxyListener{make(chan net.Conn), cancel}
+
+	go func() {
+		quicproxyclient.Serve(ctx, *rendezvouzEndpoint, nil, quicproxyclient.Handlers{
+			OnAllocation: onAllocation,
+			Proxy: func(conn net.Conn) error {
+				l.ch <- conn
+				return nil
+			},
+		})
+		close(l.ch)
+	}()
+	return l
+}
+
+type proxyListener struct {
+	ch     chan net.Conn
+	cancel func()
+}
+
+// Accept waits for and returns the next connection to the listener.
+func (l *proxyListener) Accept() (net.Conn, error) {
+	conn, ok := <-l.ch
+	if !ok {
+		return nil, io.ErrClosedPipe
+	}
+	return conn, nil
+}
+
+func (l *proxyListener) Close() error {
+	l.cancel()
+	return nil
+}
+
+func (l *proxyListener) Addr() net.Addr {
+	return noneAddr{}
+}
+
+type noneAddr struct{}
+
+func (noneAddr) Network() string { return "none" }
+func (noneAddr) String() string  { return "none" }