<commit_before>package main

import (
	"context"
	"github.com/opennetworkinglab/onos-warden/warden"
	"google.golang.org/grpc"
	"google.golang.org/grpc/grpclog"
	"io"
	"fmt"
	"os"
	"os/user"
	"flag"
	"math/rand"
	"time"
	"github.com/opennetworkinglab/onos-warden/util"
)


func init() {
	rand.Seed(time.Now().Unix())
}

func sendRequest(stream warden.ClusterClientService_ServerClustersClient, request warden.ClusterRequest, cId, cType string) {
	request.ClusterId = cId
	request.ClusterType = cType
	stream.Send(&request)
}

func reserveCluster(stream warden.ClusterClientService_ServerClustersClient, baseRequest warden.ClusterRequest) *warden.ClusterAdvertisement {
	var cluster *warden.ClusterAdvertisement
	availableClusters := make([]*warden.ClusterAdvertisement, 0)
	for {
		ad, err := stream.Recv()
		if err == io.EOF {
			// stream closed
			grpclog.Fatalln("Stream closed unexpectedly")
			os.Exit(1)
		}
		if err != nil {
			grpclog.Fatalf("Failed to receive: %v", err)
		}
		grpclog.Println("Got message:", ad)
		switch ad.State {
		case warden.ClusterAdvertisement_AVAILABLE:
			if cluster == nil {
				// pick the first available cluster
				cluster = ad
				sendRequest(stream, baseRequest, ad.ClusterId, ad.ClusterType)
			} else {
				// store the cluster away for later, just in case the current one doesn't work out
				availableClusters = append(availableClusters, ad)
			}
		case warden.ClusterAdvertisement_RESERVED:
			if cluster != nil {
				if ad.RequestId != baseRequest.RequestId {
					// cluster was assigned to someone else; try again...
					avail := availableClusters[0]
					availableClusters = availableClusters[1:]
					cluster = avail
					sendRequest(stream, baseRequest, avail.ClusterId, avail.ClusterType)
				} else if ad.State == warden.ClusterAdvertisement_RESERVED {
					// cluster is ready!
					fmt.Println("Got cluster:", ad);
					return cluster
				}
			}
		}
	}
}

// Request the first available cell, return when cell reservation is completed
func main() {
	currUser, err := user.Current()
	if err != nil {
		panic(err)
	}
	username := flag.String("user", currUser.Username, "username for reservation; defaults to $USER")
	key := flag.String("key", "", "public key for SSH")
	duration := flag.Int64("duration", -1, "duration of reservation in minutes; -1 is unlimited")
	nodes := flag.Uint64("nodes", 3, "number of nodes in cell; defaults to 3")

	reqId := *username //TODO just using the username for now
	request := warden.ClusterRequest{
		Type:        warden.ClusterRequest_RESERVE,
		Duration:    int32(*duration),
		RequestId:   reqId,
		Spec: &warden.ClusterRequest_Spec{
			ControllerNodes: uint32(*nodes),
			UserName:        *username,
			UserKey:         *key,
		},
	}

	//TODO change this to a flag
	conn, err := grpc.Dial("127.0.0.1:1234", grpc.WithInsecure())
	if err != nil {
		grpclog.Fatalf("fail to dial: %v", err)
	}
	defer conn.Close()
	client := warden.NewClusterClientServiceClient(conn)

	stream, err := client.ServerClusters(context.Background())
	if err != nil {
		grpclog.Fatalf("%v.ServerClusters(_) = _, %v", client, err)
	}
	defer stream.CloseSend()

	cluster := reserveCluster(stream, request)
	util.WaitForInterrupt()
	request.Type = warden.ClusterRequest_RETURN
	sendRequest(stream, request, cluster.ClusterId, cluster.ClusterType)
}<commit_msg>Simplifying client<commit_after>package main

import (
	"context"
	"github.com/opennetworkinglab/onos-warden/warden"
	"google.golang.org/grpc"
	"google.golang.org/grpc/grpclog"
	"io"
	"fmt"
	"os"
	"os/user"
	"flag"
	"os/signal"
)

type client struct {
	stream warden.ClusterClientService_ServerClustersClient
	ads chan *warden.ClusterAdvertisement
}

func New(addr string) *client {
	c := client {
		ads: make(chan *warden.ClusterAdvertisement),
	}

	conn, err := grpc.Dial(addr, grpc.WithInsecure())
	if err != nil {
		grpclog.Fatalf("fail to dial: %v", err)
	}
	client := warden.NewClusterClientServiceClient(conn)

	c.stream, err = client.ServerClusters(context.Background())
	if err != nil {
		grpclog.Fatalf("%v.ServerClusters(_) = _, %v", client, err)
	}

	go func() {
		for {
			ad, err := c.stream.Recv()
			if err == io.EOF {
				// stream closed
				grpclog.Fatalln("Stream closed unexpectedly")
			}
			if err != nil {
				grpclog.Fatalf("Failed to receive: %v", err)
			}
			//TODO remove this:
			grpclog.Println("Got message:", ad)
			c.ads <- ad
		}
	}()
	return &c
}

func (c *client) sendRequest(baseRequest warden.ClusterRequest, t warden.ClusterRequest_RequestType) {
	baseRequest.Type = t
	c.stream.Send(&baseRequest)
}

func (c *client) returnClusterAndExit(req warden.ClusterRequest, code int) {
	c.sendRequest(req, warden.ClusterRequest_RESERVE)
	os.Exit(code)
}

// Request the first available cell, return when cell reservation is completed or cell becomes unavailable
func main() {
	currUser, err := user.Current()
	if err != nil {
		panic(err)
	}
	username := flag.String("user", currUser.Username, "username for reservation; defaults to $USER")
	key := flag.String("key", "", "public key for SSH")
	duration := flag.Int64("duration", -1, "duration of reservation in minutes; -1 is unlimited")
	nodes := flag.Uint64("nodes", 3, "number of nodes in cell; defaults to 3")
	addr := flag.String("addr", "127.0.0.1:1234", "address of warden")

	//Note: Request ids must be unique
	// ClusterId and ClusterType are optional and we won't be filling those in
	reqId := *username //TODO just using the username for now
	baseRequest := warden.ClusterRequest{
		Duration:    int32(*duration),
		RequestId:   reqId,
		Spec: &warden.ClusterRequest_Spec{
			ControllerNodes: uint32(*nodes),
			UserName:        *username,
			UserKey:         *key,
		},
	}
	c := New(*addr)
	c.sendRequest(baseRequest, warden.ClusterRequest_RESERVE)

	intrChan := make(chan os.Signal)
	signal.Notify(intrChan, os.Interrupt, os.Kill)
	var cluster *warden.ClusterAdvertisement
	for {
		select {
		case <-intrChan:
			c.returnClusterAndExit(baseRequest,0)
		case ad := <-c.ads:
			switch ad.State {
			case warden.ClusterAdvertisement_RESERVED:
				if baseRequest.RequestId == ad.RequestId {
					// cluster is ready!
					fmt.Println("Got cluster:", ad);
					cluster = ad
				}
			case warden.ClusterAdvertisement_UNAVAILABLE:
				if cluster.State == warden.ClusterAdvertisement_RESERVED &&
					cluster.ClusterId == ad.ClusterId &&
					cluster.ClusterType == ad.ClusterType &&
					cluster.RequestId == ad.RequestId {
					// our cluster is no longer available
					fmt.Println("Returning cluster, then exit error")
					c.returnClusterAndExit(baseRequest, 1)
				}
			}
		}
	}
}
{'fraction_non_alphanumeric': 0.12814298697364435, 'fraction_numerical': 0.007573462587094819, 'mean_word_length': 3.898367952522255, 'pattern_counts': {'":': 0, '<': 6, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 2, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5675862', 'n_tokens_mistral': 2471, 'n_tokens_neox': 2185, 'n_words': 543}