func findParentMachine(node *pb.ResourceTopologyNodeDescriptor, resourceMap *types.ResourceMap) *pb.ResourceDescriptor {
	for {
		if node.ResourceDesc.Type == pb.ResourceDescriptor_ResourceMachine {
			return node.ResourceDesc
		}
		id, err := strconv.ParseUint(node.ParentId, 10, 64)
		if err != nil {
			log.Panicf("Could not parse parentID\n")
		}
		parentID := types.ResourceID(id)
		parentNode := resourceMap.FindPtrOrNull(parentID)
		if parentNode == nil {
			return nil
		}
		node = parentNode.TopologyNode
	}
}
FindParentMachine returns the machine descriptor if given any descendant node of the machine.
If it's not a descendant or the machine node it self then it returns nil.
package main
import (
	"flag"
	"fmt"
	"log"
	"strconv"
	"time"
	"github.com/coreos/ksched/k8s/k8sclient"
	"github.com/coreos/ksched/k8s/k8stype"
	"github.com/coreos/ksched/pkg/types"
	"github.com/coreos/ksched/pkg/types/resourcestatus"
	"github.com/coreos/ksched/pkg/util"
	"github.com/coreos/ksched/pkg/util/idgenerator"
	"github.com/coreos/ksched/pkg/util/queue"
	pb "github.com/coreos/ksched/proto"
	"github.com/coreos/ksched/scheduling/flow/flowscheduler"
)
var (
	numMachines      int
	address          string
	maxTasksPerPu    int
	batchTimeout     int
	nodeBatchTimeout int
	podChanSize      int
	fakeMachines     bool
)
func init() {
	flag.StringVar(&address, "addr", "127.0.0.1:8080", "APIServer addr")
	flag.IntVar(&maxTasksPerPu, "mt", 1000, "maximum number of tasks per machine")
	flag.IntVar(&batchTimeout, "pbt", 2, "pods batch timeout in seconds")
	flag.IntVar(&nodeBatchTimeout, "nbt", 2, "node batch timeout in seconds")
	flag.IntVar(&podChanSize, "pcs", 5000, "pod channel size in client's pod informer")
	flag.BoolVar(&fakeMachines, "fakeMachines", false, "fake the machine topology if running without a real cluster")
	flag.IntVar(&numMachines, "nm", 2, "number of machines, only needed if faking the resource topology")
	flag.Parse()
}
type k8scheduler struct {
	nodeToMachineID map[string]string
	machineToNodeID map[string]string
	podToTaskID     map[string]uint64
	taskToPodID     map[uint64]string
	oldTaskBindings map[types.TaskID]types.ResourceID
	resourceMap     *types.ResourceMap
	jobMap          *types.JobMap
	taskMap         *types.TaskMap
	rootNode        *pb.ResourceTopologyNodeDescriptor
	flowScheduler   flowscheduler.Scheduler
	client          *k8sclient.Client
	maxTasksPerPu   int
}
func New(client *k8sclient.Client, maxTasksPerPu int) *k8scheduler {
	resourceMap := types.NewResourceMap()
	jobMap := types.NewJobMap()
	taskMap := types.NewTaskMap()
	rootNode := &pb.ResourceTopologyNodeDescriptor{
		ResourceDesc: createResourceDesc(pb.ResourceDescriptor_ResourceCoordinator, 0),
	}
	flowScheduler := flowscheduler.NewScheduler(resourceMap, jobMap, taskMap, rootNode, uint64(maxTasksPerPu))
	return &k8scheduler{
		nodeToMachineID: make(map[string]string),
		machineToNodeID: make(map[string]string),
		podToTaskID:     make(map[string]uint64),
		taskToPodID:     make(map[uint64]string),
		oldTaskBindings: make(map[types.TaskID]types.ResourceID),
		resourceMap:     resourceMap,
		jobMap:          jobMap,
		taskMap:         taskMap,
		rootNode:        rootNode,
		client:          client,
		flowScheduler:   flowScheduler,
		maxTasksPerPu:   maxTasksPerPu,
	}
}
func main() {
	config := k8sclient.Config{Addr: address}
	client, err := k8sclient.New(config, podChanSize)
	if err != nil {
		panic(err)
	}
	scheduler := New(client, maxTasksPerPu)
	fmt.Printf("Initializing nodes in resource topology\n")
	if fakeMachines {
		scheduler.fakeResourceTopology(numMachines)
	} else {
		scheduler.initResourceTopology()
	}
	fmt.Printf("Done\n")
	scheduler.Run(client)
}
func (ks *k8scheduler) Run(client *k8sclient.Client) {
	jobID := addNewJob(ks.jobMap, ks.flowScheduler)
	for {
		newPods := client.GetPodBatch(time.Duration(batchTimeout) * time.Second)
		if len(newPods) == 0 {
			continue
		}
		log.Printf("Adding Pods as tasks to scheduler\n")
		for _, pod := range newPods {
			if _, ok := ks.podToTaskID[pod.ID]; ok {
				fmt.Printf("Skipping already existing Pod:%v ==> Task:%v\n", pod.ID, ks.podToTaskID[pod.ID])
				continue
			}
			taskID := addTaskToJob(jobID, ks.jobMap, ks.taskMap)
			ks.podToTaskID[pod.ID] = uint64(taskID)
			ks.taskToPodID[uint64(taskID)] = pod.ID
		}
		fmt.Printf("\nScheduling all tasks\n")
		start := time.Now()
		ks.flowScheduler.ScheduleAllJobs()
		elapsed := time.Since(start)
		fmt.Printf("Time taken to schedule all tasks:%s\n", elapsed)
		podToNodeBindings := make([]*k8stype.Binding, 0)
		taskBindings := ks.flowScheduler.GetTaskBindings()
		for taskID, resourceID := range taskBindings {
			if ks.oldTaskBindings[taskID] == resourceID {
				continue
			}
			ks.oldTaskBindings[taskID] = resourceID
			puNode := ks.resourceMap.FindPtrOrNull(resourceID).TopologyNode
			machineID := findParentMachine(puNode, ks.resourceMap).Uuid
			nodeID := ks.machineToNodeID[machineID]
			podID := ks.taskToPodID[uint64(taskID)]
			binding := &k8stype.Binding{
				PodID:  podID,
				NodeID: nodeID,
			}
			podToNodeBindings = append(podToNodeBindings, binding)
		}
		ks.client.AssignBinding(podToNodeBindings)
		fmt.Printf("%d Pod Bindings assigned\n\n", len(podToNodeBindings))
	}
}
func (ks *k8scheduler) fakeResourceTopology(numMachines int) {
	nodeIDGen := idgenerator.New(false)
	for i := 0; i < numMachines; i++ {
		nextID := strconv.FormatUint(nodeIDGen.NextID(), 10)
		machineNode := addMachine(nextID, ks.maxTasksPerPu, ks.rootNode, ks.resourceMap, ks.flowScheduler)
		ks.nodeToMachineID[nextID] = machineNode.ResourceDesc.Uuid
		ks.machineToNodeID[machineNode.ResourceDesc.Uuid] = nextID
	}
}
func (ks *k8scheduler) initResourceTopology() {
	nodeChan := ks.client.GetNodeChan()
	done := make(chan bool)
	go func() {
		<-time.After(time.Duration(nodeBatchTimeout) * time.Second)
		done <- true
	}()
	finish := false
	for !finish {
		select {
		case node := <-nodeChan:
			if _, ok := ks.nodeToMachineID[node.ID]; ok {
				continue
			}
			machineNode := addMachine(node.ID, ks.maxTasksPerPu, ks.rootNode, ks.resourceMap, ks.flowScheduler)
			ks.nodeToMachineID[node.ID] = machineNode.ResourceDesc.Uuid
			ks.machineToNodeID[machineNode.ResourceDesc.Uuid] = node.ID
		case <-done:
			finish = true
		default:
		}
	}
}
func addNewJob(jobMap *types.JobMap, scheduler flowscheduler.Scheduler) types.JobID {
	jobID := types.JobID(util.RandUint64())
	for jobMap.ContainsKey(jobID) {
		jobID = types.JobID(util.RandUint64())
	}
	jobUuid := strconv.FormatUint(uint64(jobID), 10)
	name := "Job " + jobUuid
	jobDesc := &pb.JobDescriptor{
		Uuid:  jobUuid,
		Name:  name,
		State: pb.JobDescriptor_Created,
	}
	jobMap.InsertIfNotPresent(jobID, jobDesc)
	scheduler.AddJob(jobDesc)
	return jobID
}
func addTaskToJob(jobID types.JobID, jobMap *types.JobMap, taskMap *types.TaskMap) types.TaskID {
	jobDesc := jobMap.FindPtrOrNull(jobID)
	jobUuid := strconv.FormatUint(uint64(jobID), 10)
	if jobDesc == nil {
		log.Panicf("No job for jobID:%v exists in the jobMap\n", jobID)
	}
	taskID := types.TaskID(util.RandUint64())
	for taskMap.ContainsKey(taskID) {
		taskID = types.TaskID(util.RandUint64())
	}
	name := "Task " + strconv.FormatUint(uint64(taskID), 10)
	task := &pb.TaskDescriptor{
		Uid:   uint64(taskID),
		Name:  name,
		State: pb.TaskDescriptor_Created,
		JobID: jobUuid,
	}
	taskMap.InsertIfNotPresent(taskID, task)
	if jobDesc.RootTask == nil {
		jobDesc.RootTask = task
	} else {
		jobDesc.RootTask.Spawned = append(jobDesc.RootTask.Spawned, task)
	}
	return taskID
}
func addMachine(nodeID string, tasksPerPu int, root *pb.ResourceTopologyNodeDescriptor, resourceMap *types.ResourceMap, scheduler flowscheduler.Scheduler) *pb.ResourceTopologyNodeDescriptor {
	machineNode := createMachineNode(tasksPerPu)
	root.Children = append(root.Children, machineNode)
	machineNode.ParentId = root.ResourceDesc.Uuid
	nodes := queue.NewFIFO()
	nodes.Push(machineNode)
	for !nodes.IsEmpty() {
		currNode := nodes.Pop().(*pb.ResourceTopologyNodeDescriptor)
		resourceStatus := &resourcestatus.ResourceStatus{
			Descriptor:   currNode.ResourceDesc,
			TopologyNode: currNode,
		}
		resourceMap.InsertIfNotPresent(util.MustResourceIDFromString(currNode.ResourceDesc.Uuid), resourceStatus)
		for _, childDesc := range currNode.GetChildren() {
			nodes.Push(childDesc)
		}
	}
	scheduler.RegisterResource(machineNode)
	return machineNode
}
func createMachineNode(tasksPerPu int) *pb.ResourceTopologyNodeDescriptor {
	totalCap := tasksPerPu
	machineNode := &pb.ResourceTopologyNodeDescriptor{
		ResourceDesc: createResourceDesc(pb.ResourceDescriptor_ResourceMachine, totalCap),
	}
	puNode := getNewPuRtnd(tasksPerPu)
	puNode.ParentId = machineNode.ResourceDesc.Uuid
	machineNode.Children = append(machineNode.Children, puNode)
	machineNode.ResourceDesc.ResourceCapacity = &pb.ResourceVector{
		CpuCores: float32(1),  
		RamCap:   1024,
	}
	return machineNode
}
func getNewPuRtnd(taskCap int) *pb.ResourceTopologyNodeDescriptor {
	return &pb.ResourceTopologyNodeDescriptor{
		ResourceDesc: createResourceDesc(pb.ResourceDescriptor_ResourcePu, taskCap),
	}
}
func createResourceDesc(resourceType pb.ResourceDescriptor_ResourceType, taskCap int) *pb.ResourceDescriptor {
	IDString := strconv.FormatUint(util.RandUint64(), 10)
	name := fmt.Sprintf("%s-%s", pb.ResourceDescriptor_ResourceType_name[int32(resourceType)], IDString)
	return &pb.ResourceDescriptor{
		Uuid:         IDString,
		FriendlyName: name,
		TaskCapacity: uint64(taskCap),
		Type:         resourceType,
		State:       pb.ResourceDescriptor_ResourceIdle,
		Schedulable: true,
	}
}