k3s/test/e2e/autoscaling_utils.go

/*
Copyright 2015 The Kubernetes Authors All rights reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package e2e

import (
	"strconv"
	"time"

	"k8s.io/kubernetes/pkg/api"
	client "k8s.io/kubernetes/pkg/client/unversioned"
	"k8s.io/kubernetes/pkg/util"

	. "github.com/onsi/ginkgo"
)

const (
	consumptionTimeInSeconds = 30
	sleepTime                = 30 * time.Second
	requestSizeInMillicores  = 100
	requestSizeInMegabytes   = 100
	port                     = 80
	targetPort               = 8080
	timeoutRC                = 120 * time.Second
	image                    = "gcr.io/google_containers/resource_consumer:beta"
	rcIsNil                  = "ERROR: replicationController = nil"
)

/*
ResourceConsumer is a tool for testing. It helps create specified usage of CPU or memory (Warnig: memory not supported)
typical use case:
rc.ConsumeCPU(600)
// ... check your assumption here
rc.ConsumeCPU(300)
// ... check your assumption here
*/
type ResourceConsumer struct {
	name      string
	framework *Framework
	cpu       chan int
	mem       chan int
	stopCPU   chan int
	stopMem   chan int
}

/*
NewResourceConsumer creates new ResourceConsumer
initCPU argument is in millicores
initMemory argument is in megabytes
*/
func NewResourceConsumer(name string, replicas, initCPU, initMemory int, framework *Framework) *ResourceConsumer {
	runServiceAndRCForResourceConsumer(framework.Client, framework.Namespace.Name, name, replicas)
	rc := &ResourceConsumer{
		name:      name,
		framework: framework,
		cpu:       make(chan int),
		mem:       make(chan int),
		stopCPU:   make(chan int),
		stopMem:   make(chan int),
	}
	go rc.makeConsumeCPURequests()
	rc.ConsumeCPU(initCPU)
	go rc.makeConsumeMemRequests()
	rc.ConsumeMem(initMemory)
	return rc
}

// ConsumeCPU consumes given number of CPU
func (rc *ResourceConsumer) ConsumeCPU(millicores int) {
	rc.cpu <- millicores
}

// ConsumeMem consumes given number of Mem
func (rc *ResourceConsumer) ConsumeMem(megabytes int) {
	rc.mem <- megabytes
}

func (rc *ResourceConsumer) makeConsumeCPURequests() {
	var count int
	var rest int
	for {
		select {
		case millicores := <-rc.cpu:
			count = millicores / requestSizeInMillicores
			rest = millicores - count*requestSizeInMillicores
		case <-time.After(sleepTime):
			if count > 0 {
				rc.sendConsumeCPURequests(count, requestSizeInMillicores, consumptionTimeInSeconds)
			}
			if rest > 0 {
				go rc.sendOneConsumeCPURequest(rest, consumptionTimeInSeconds)
			}
		case <-rc.stopCPU:
			return
		}
	}
}

func (rc *ResourceConsumer) makeConsumeMemRequests() {
	var count int
	var rest int
	for {
		select {
		case megabytes := <-rc.mem:
			count = megabytes / requestSizeInMegabytes
			rest = megabytes - count*requestSizeInMegabytes
		case <-time.After(sleepTime):
			if count > 0 {
				rc.sendConsumeMemRequests(count, requestSizeInMegabytes, consumptionTimeInSeconds)
			}
			if rest > 0 {
				go rc.sendOneConsumeMemRequest(rest, consumptionTimeInSeconds)
			}
		case <-rc.stopMem:
			return
		}
	}
}

func (rc *ResourceConsumer) sendConsumeCPURequests(requests, millicores, durationSec int) {
	for i := 0; i < requests; i++ {
		go rc.sendOneConsumeCPURequest(millicores, durationSec)
	}
}

func (rc *ResourceConsumer) sendConsumeMemRequests(requests, megabytes, durationSec int) {
	for i := 0; i < requests; i++ {
		go rc.sendOneConsumeMemRequest(megabytes, durationSec)
	}
}

// sendOneConsumeCPURequest sends POST request for cpu consumption
func (rc *ResourceConsumer) sendOneConsumeCPURequest(millicores int, durationSec int) {
	defer GinkgoRecover()
	_, err := rc.framework.Client.Post().
		Prefix("proxy").
		Namespace(rc.framework.Namespace.Name).
		Resource("services").
		Name(rc.name).
		Suffix("ConsumeCPU").
		Param("millicores", strconv.Itoa(millicores)).
		Param("durationSec", strconv.Itoa(durationSec)).
		Do().
		Raw()
	expectNoError(err)
}

// sendOneConsumeMemRequest sends POST request for memory consumption
func (rc *ResourceConsumer) sendOneConsumeMemRequest(megabytes int, durationSec int) {
	defer GinkgoRecover()
	_, err := rc.framework.Client.Post().
		Prefix("proxy").
		Namespace(rc.framework.Namespace.Name).
		Resource("services").
		Name(rc.name).
		Suffix("ConsumeMem").
		Param("megabytes", strconv.Itoa(megabytes)).
		Param("durationSec", strconv.Itoa(durationSec)).
		Do().
		Raw()
	expectNoError(err)
}

func (rc *ResourceConsumer) GetReplicas() int {
	replicationController, err := rc.framework.Client.ReplicationControllers(rc.framework.Namespace.Name).Get(rc.name)
	expectNoError(err)
	if replicationController == nil {
		Failf(rcIsNil)
	}
	return replicationController.Status.Replicas
}

func (rc *ResourceConsumer) WaitForReplicas(desiredReplicas int) {
	timeout := 10 * time.Minute
	for start := time.Now(); time.Since(start) < timeout; time.Sleep(20 * time.Second) {
		if desiredReplicas == rc.GetReplicas() {
			Logf("Replication Controller current replicas number is equal to desired replicas number: %d", desiredReplicas)
			return
		} else {
			Logf("Replication Controller current replicas number %d waiting to be %d", rc.GetReplicas(), desiredReplicas)
		}
	}
	Failf("timeout waiting %v for pods size to be %d", timeout, desiredReplicas)
}

func (rc *ResourceConsumer) CleanUp() {
	rc.stopCPU <- 0
	rc.stopMem <- 0
	expectNoError(DeleteRC(rc.framework.Client, rc.framework.Namespace.Name, rc.name))
	expectNoError(rc.framework.Client.Services(rc.framework.Namespace.Name).Delete(rc.name))
	expectNoError(rc.framework.Client.Experimental().HorizontalPodAutoscalers(rc.framework.Namespace.Name).Delete(rc.name, api.NewDeleteOptions(0)))
}

func runServiceAndRCForResourceConsumer(c *client.Client, ns, name string, replicas int) {
	_, err := c.Services(ns).Create(&api.Service{
		ObjectMeta: api.ObjectMeta{
			Name: name,
		},
		Spec: api.ServiceSpec{
			Ports: []api.ServicePort{{
				Port:       port,
				TargetPort: util.NewIntOrStringFromInt(targetPort),
			}},

			Selector: map[string]string{
				"name": name,
			},
		},
	})
	expectNoError(err)
	config := RCConfig{
		Client:    c,
		Image:     image,
		Name:      name,
		Namespace: ns,
		Timeout:   timeoutRC,
		Replicas:  replicas,
	}
	expectNoError(RunRC(config))

}