<commit_before>// SPDX-License-Identifier: Apache-2.0
// Copyright Authors of Cilium

package resource

import (
	"context"
	"sync"

	corev1 "k8s.io/api/core/v1"
	k8sRuntime "k8s.io/apimachinery/pkg/runtime"
	"k8s.io/client-go/tools/cache"
	"k8s.io/client-go/util/workqueue"

	"github.com/cilium/cilium/pkg/hive"
	"github.com/cilium/cilium/pkg/lock"
	"github.com/cilium/cilium/pkg/promise"
	"github.com/cilium/cilium/pkg/stream"
)

// Resource provides access to a Kubernetes resource through either
// a stream of events or a read-only store.
//
// Observing of the events can be done from a constructor as subscriber
// registration is non-blocking.
//
// Store() however should only be called from a start hook, or from a
// goroutine forked from the start hook as it blocks until the store
// has been synchronized.
//
// The subscriber can process the event synchronously with the
// Event[T].Handle function, or asynchronously by calling Event[T].Done
// once the event has been processed. On errors the object's key is requeued
// for later processing. Once maximum number of retries is reached the subscriber's
// event stream will be completed with the error from the last retry attempt.
//
// The resource is lazy, e.g. it will not start the informer until a call
// has been made to Observe() or Store().
type Resource[T k8sRuntime.Object] interface {
	stream.Observable[Event[T]]

	// Store retrieves the read-only store for the resource. Blocks until
	// the store has been synchronized or the context cancelled.
	// Returns a non-nil error if context is cancelled or the resource
	// has been stopped before store has synchronized.
	Store(context.Context) (Store[T], error)
}

// New creates a new Resource[T]. Use with hive.Provide:
//
//	var exampleCell = hive.Module(
//		"example",
//	 	cell.Provide(
//		 	// Provide `Resource[*slim_corev1.Pod]` to the hive:
//		 	func(lc hive.Lifecycle, c k8sClient.Clientset) resource.Resource[*slim_corev1.Pod] {
//				lw := utils.ListerWatcherFromTyped[*slim_corev1.PodList](
//					c.Slim().CoreV1().Pods(""),
//				)
//		 		return resource.New(lc, lw)
//		 	}
//			// Use the resource:
//			newExample,
//		}),
//		...
//	)
//	func newExample(pods resource.Resource[*slim_corev1.Pod]) *Example {
//		e := &Example{...}
//		pods.Observe(e.ctx, e.onPodUpdated, e.onPodsComplete)
//		return e
//	}
//	func (e *Example) onPodUpdated(key resource.Key, pod *slim_core.Pod) error {
//		// Process event ...
//	}
//	func (e *Example) onPodsComplete(err error) {
//		// Handle error ...
//	}
//
// See also pkg/k8s/resource/example/main.go for a runnable example.
func New[T k8sRuntime.Object](lc hive.Lifecycle, lw cache.ListerWatcher, opts ...Option) Resource[T] {
	r := &resource[T]{
		queues: make(map[uint64]*keyQueue),
		needed: make(chan struct{}, 1),
		opts:   defaultOptions(),
		lw:     lw,
	}
	for _, applyOpt := range opts {
		applyOpt(&r.opts)
	}
	r.ctx, r.cancel = context.WithCancel(context.Background())
	r.storeResolver, r.storePromise = promise.New[Store[T]]()
	lc.Append(r)
	return r
}

type resource[T k8sRuntime.Object] struct {
	mu     lock.RWMutex
	ctx    context.Context
	cancel context.CancelFunc
	wg     sync.WaitGroup

	needed chan struct{}

	queues map[uint64]*keyQueue
	subId  uint64

	lw   cache.ListerWatcher
	opts options

	storePromise  promise.Promise[Store[T]]
	storeResolver promise.Resolver[Store[T]]
}

var _ Resource[*corev1.Node] = &resource[*corev1.Node]{}

func (r *resource[T]) Store(ctx context.Context) (Store[T], error) {
	r.markNeeded()
	return r.storePromise.Await(ctx)
}

func (r *resource[T]) pushUpdate(key Key) {
	r.mu.RLock()
	for _, queue := range r.queues {
		queue.Add(&updateEntry{key})
	}
	r.mu.RUnlock()
}

func (r *resource[T]) pushDelete(lastState any) {
	key := NewKey(lastState)
	obj := lastState
	if d, ok := lastState.(cache.DeletedFinalStateUnknown); ok {
		obj = d.Obj
	}
	r.mu.RLock()
	for _, queue := range r.queues {
		queue.Add(&deleteEntry{key, obj})
	}
	r.mu.RUnlock()
}

func (r *resource[T]) Start(startCtx hive.HookContext) error {
	r.wg.Add(1)
	go r.startWhenNeeded()
	return nil
}

func (r *resource[T]) markNeeded() {
	select {
	case r.needed <- struct{}{}:
	default:
	}
}

func (r *resource[T]) startWhenNeeded() {
	defer r.wg.Done()

	// Wait until we're needed before starting the informer.
	select {
	case <-r.ctx.Done():
		return
	case <-r.needed:
	}

	// Short-circuit if we're being stopped.
	if r.ctx.Err() != nil {
		return
	}

	// Construct the informer and run it.
	var objType T
	handlerFuncs :=
		cache.ResourceEventHandlerFuncs{
			AddFunc:    func(obj any) { r.pushUpdate(NewKey(obj)) },
			UpdateFunc: func(old any, new any) { r.pushUpdate(NewKey(new)) },
			DeleteFunc: func(obj any) { r.pushDelete(obj) },
		}

	store, informer := cache.NewInformer(r.lw, objType, 0, handlerFuncs)

	r.wg.Add(1)
	go func() {
		defer r.wg.Done()
		informer.Run(r.ctx.Done())
	}()

	// Wait for cache to be synced before resolving the store
	if !cache.WaitForCacheSync(r.ctx.Done(), informer.HasSynced) {
		// The context is cancelled when stopping and all dependees of Resource[T] are
		// stopped before it, but to be safe, resolve the store with nil to catch
		// misbehaving dependees.
		r.storeResolver.Reject(r.ctx.Err())
	} else {
		r.storeResolver.Resolve(&typedStore[T]{store})
	}
}

func (r *resource[T]) Stop(stopCtx hive.HookContext) error {
	r.cancel()
	r.wg.Wait()
	return nil
}

func (r *resource[T]) Observe(subCtx context.Context, next func(Event[T]), complete func(error)) {
	r.markNeeded()

	subCtx, subCancel := context.WithCancel(subCtx)

	queue := &keyQueue{
		RateLimitingInterface: workqueue.NewRateLimitingQueue(r.opts.rateLimiter()),
		errorHandler:          r.opts.errorHandler,
	}
	r.mu.Lock()
	subId := r.subId
	r.subId++
	r.mu.Unlock()

	// Fork a goroutine to pop elements from the queue and pass them to the subscriber.
	r.wg.Add(1)
	go func() {
		defer r.wg.Done()

		// Make sure to call ShutDown() in the end. Calling ShutDownWithDrain is not
		// enough as DelayingQueue does not implement it, so without ShutDown() we'd
		// leak the (*delayingType).waitingLoop.
		defer queue.ShutDown()

		// Wait for the store so we can emit the initial items and the sync event.
		store, err := r.storePromise.Await(subCtx)
		if err != nil {
			complete(err)
			return
		}

		r.mu.Lock()
		r.queues[subId] = queue

		// Append the initial set of keys to the queue + sentinel for the sync event.
		// We go through the queue instead of directly emitting to avoid a race between
		// the queue add and listing keys, plus to have one error handling path for Event.Done().
		keyIter := store.IterKeys()
		for keyIter.Next() {
			queue.Add(&updateEntry{keyIter.Key()})
		}
		queue.Add(&syncEntry{})
		r.mu.Unlock()

		for {
			// Retrieve an item from the subscribers queue and then fetch the object
			// from the store.
			raw, shutdown := queue.Get()
			if shutdown {
				break
			}
			entry := raw.(queueEntry)
			baseEvent := baseEvent{
				func(err error) { queue.eventDone(entry, err) },
			}
			switch entry := entry.(type) {
			case *syncEntry:
				next(&SyncEvent[T]{baseEvent})
			case *deleteEntry:
				next(&DeleteEvent[T]{baseEvent, entry.key, entry.obj.(T)})
			case *updateEntry:
				obj, exists, err := store.GetByKey(entry.key)
				if err != nil {
					queue.setError(err)
					break
				}
				// Emit the update event if the item exists, if it doesn't, then
				// it has been deleted and a delete will follow soon.
				if exists {
					next(&UpdateEvent[T]{baseEvent, entry.key, obj})
				}
			}
		}
		r.mu.Lock()
		delete(r.queues, subId)
		r.mu.Unlock()
		complete(queue.getError())
	}()

	// Fork a goroutine to wait for either the subscriber cancelling or the resource
	// shutting down.
	r.wg.Add(1)
	go func() {
		defer r.wg.Done()
		select {
		case <-r.ctx.Done():
		case <-subCtx.Done():
		}
		subCancel()
		queue.ShutDownWithDrain()
	}()
}

// keyQueue wraps the workqueue to implement the error retry logic for a single subscriber,
// e.g. it implements the eventDone() method called by Event[T].Done().
type keyQueue struct {
	lock.Mutex
	workqueue.RateLimitingInterface

	errorHandler ErrorHandler

	err error
}

func (kq *keyQueue) getError() error {
	kq.Lock()
	defer kq.Unlock()
	return kq.err
}

func (kq *keyQueue) setError(err error) {
	kq.Lock()
	defer kq.Unlock()
	if kq.err == nil {
		kq.err = err
	}
}

func (kq *keyQueue) eventDone(entry queueEntry, err error) {
	// This is based on the example found in k8s.io/client-go/examples/workqueue/main.go.

	// Mark the object as done being processed. If it was marked dirty
	// during processing, it'll be processed again.
	defer kq.Done(entry)

	if err != nil {
		numRequeues := kq.NumRequeues(entry)

		var action ErrorAction
		switch entry := entry.(type) {
		case *syncEntry:
			action = ErrorActionStop
		case *updateEntry:
			action = kq.errorHandler(entry.key, numRequeues, err)
		case *deleteEntry:
			action = kq.errorHandler(entry.key, numRequeues, err)
		}

		switch action {
		case ErrorActionRetry:
			go kq.AddRateLimited(entry)
		case ErrorActionStop:
			kq.setError(err)
			kq.ShutDown()
		case ErrorActionIgnore:
			kq.Forget(entry)
		}
	} else {
		// As the object was processed successfully we can "forget" it.
		// This clears any rate limiter state associated with this object, so
		// it won't be throttled based on previous failure history.
		kq.Forget(entry)
	}
}

type queueEntry interface {
	isQueueEntry()
}

type syncEntry struct{}

func (syncEntry) isQueueEntry() {}

type updateEntry struct {
	key Key
}

func (updateEntry) isQueueEntry() {}

type deleteEntry struct {
	key Key
	obj any
}

func (deleteEntry) isQueueEntry() {}
<commit_msg>resource: Fix queue entry coalescing<commit_after>// SPDX-License-Identifier: Apache-2.0
// Copyright Authors of Cilium

package resource

import (
	"context"
	"sync"

	corev1 "k8s.io/api/core/v1"
	k8sRuntime "k8s.io/apimachinery/pkg/runtime"
	"k8s.io/client-go/tools/cache"
	"k8s.io/client-go/util/workqueue"

	"github.com/cilium/cilium/pkg/hive"
	"github.com/cilium/cilium/pkg/lock"
	"github.com/cilium/cilium/pkg/promise"
	"github.com/cilium/cilium/pkg/stream"
)

// Resource provides access to a Kubernetes resource through either
// a stream of events or a read-only store.
//
// Observing of the events can be done from a constructor as subscriber
// registration is non-blocking.
//
// Store() however should only be called from a start hook, or from a
// goroutine forked from the start hook as it blocks until the store
// has been synchronized.
//
// The subscriber can process the event synchronously with the
// Event[T].Handle function, or asynchronously by calling Event[T].Done
// once the event has been processed. On errors the object's key is requeued
// for later processing. Once maximum number of retries is reached the subscriber's
// event stream will be completed with the error from the last retry attempt.
//
// The resource is lazy, e.g. it will not start the informer until a call
// has been made to Observe() or Store().
type Resource[T k8sRuntime.Object] interface {
	stream.Observable[Event[T]]

	// Store retrieves the read-only store for the resource. Blocks until
	// the store has been synchronized or the context cancelled.
	// Returns a non-nil error if context is cancelled or the resource
	// has been stopped before store has synchronized.
	Store(context.Context) (Store[T], error)
}

// New creates a new Resource[T]. Use with hive.Provide:
//
//	var exampleCell = hive.Module(
//		"example",
//	 	cell.Provide(
//		 	// Provide `Resource[*slim_corev1.Pod]` to the hive:
//		 	func(lc hive.Lifecycle, c k8sClient.Clientset) resource.Resource[*slim_corev1.Pod] {
//				lw := utils.ListerWatcherFromTyped[*slim_corev1.PodList](
//					c.Slim().CoreV1().Pods(""),
//				)
//		 		return resource.New(lc, lw)
//		 	}
//			// Use the resource:
//			newExample,
//		}),
//		...
//	)
//	func newExample(pods resource.Resource[*slim_corev1.Pod]) *Example {
//		e := &Example{...}
//		pods.Observe(e.ctx, e.onPodUpdated, e.onPodsComplete)
//		return e
//	}
//	func (e *Example) onPodUpdated(key resource.Key, pod *slim_core.Pod) error {
//		// Process event ...
//	}
//	func (e *Example) onPodsComplete(err error) {
//		// Handle error ...
//	}
//
// See also pkg/k8s/resource/example/main.go for a runnable example.
func New[T k8sRuntime.Object](lc hive.Lifecycle, lw cache.ListerWatcher, opts ...Option) Resource[T] {
	r := &resource[T]{
		queues: make(map[uint64]*keyQueue),
		needed: make(chan struct{}, 1),
		opts:   defaultOptions(),
		lw:     lw,
	}
	for _, applyOpt := range opts {
		applyOpt(&r.opts)
	}
	r.ctx, r.cancel = context.WithCancel(context.Background())
	r.storeResolver, r.storePromise = promise.New[Store[T]]()
	lc.Append(r)
	return r
}

type resource[T k8sRuntime.Object] struct {
	mu     lock.RWMutex
	ctx    context.Context
	cancel context.CancelFunc
	wg     sync.WaitGroup

	needed chan struct{}

	queues map[uint64]*keyQueue
	subId  uint64

	lw   cache.ListerWatcher
	opts options

	storePromise  promise.Promise[Store[T]]
	storeResolver promise.Resolver[Store[T]]
}

var _ Resource[*corev1.Node] = &resource[*corev1.Node]{}

func (r *resource[T]) Store(ctx context.Context) (Store[T], error) {
	r.markNeeded()
	return r.storePromise.Await(ctx)
}

func (r *resource[T]) pushUpdate(key Key) {
	r.mu.RLock()
	for _, queue := range r.queues {
		queue.AddUpdate(key)
	}
	r.mu.RUnlock()
}

func (r *resource[T]) pushDelete(lastState any) {
	key := NewKey(lastState)
	obj := lastState
	if d, ok := lastState.(cache.DeletedFinalStateUnknown); ok {
		obj = d.Obj
	}
	r.mu.RLock()
	for _, queue := range r.queues {
		queue.AddDelete(key, obj)
	}
	r.mu.RUnlock()
}

func (r *resource[T]) Start(startCtx hive.HookContext) error {
	r.wg.Add(1)
	go r.startWhenNeeded()
	return nil
}

func (r *resource[T]) markNeeded() {
	select {
	case r.needed <- struct{}{}:
	default:
	}
}

func (r *resource[T]) startWhenNeeded() {
	defer r.wg.Done()

	// Wait until we're needed before starting the informer.
	select {
	case <-r.ctx.Done():
		return
	case <-r.needed:
	}

	// Short-circuit if we're being stopped.
	if r.ctx.Err() != nil {
		return
	}

	// Construct the informer and run it.
	var objType T
	handlerFuncs :=
		cache.ResourceEventHandlerFuncs{
			AddFunc:    func(obj any) { r.pushUpdate(NewKey(obj)) },
			UpdateFunc: func(old any, new any) { r.pushUpdate(NewKey(new)) },
			DeleteFunc: func(obj any) { r.pushDelete(obj) },
		}

	store, informer := cache.NewInformer(r.lw, objType, 0, handlerFuncs)

	r.wg.Add(1)
	go func() {
		defer r.wg.Done()
		informer.Run(r.ctx.Done())
	}()

	// Wait for cache to be synced before resolving the store
	if !cache.WaitForCacheSync(r.ctx.Done(), informer.HasSynced) {
		// The context is cancelled when stopping and all dependees of Resource[T] are
		// stopped before it, but to be safe, resolve the store with nil to catch
		// misbehaving dependees.
		r.storeResolver.Reject(r.ctx.Err())
	} else {
		r.storeResolver.Resolve(&typedStore[T]{store})
	}
}

func (r *resource[T]) Stop(stopCtx hive.HookContext) error {
	r.cancel()
	r.wg.Wait()
	return nil
}

func (r *resource[T]) Observe(subCtx context.Context, next func(Event[T]), complete func(error)) {
	r.markNeeded()

	subCtx, subCancel := context.WithCancel(subCtx)

	queue := &keyQueue{
		RateLimitingInterface: workqueue.NewRateLimitingQueue(r.opts.rateLimiter()),
		errorHandler:          r.opts.errorHandler,
	}
	r.mu.Lock()
	subId := r.subId
	r.subId++
	r.mu.Unlock()

	// Fork a goroutine to pop elements from the queue and pass them to the subscriber.
	r.wg.Add(1)
	go func() {
		defer r.wg.Done()

		// Make sure to call ShutDown() in the end. Calling ShutDownWithDrain is not
		// enough as DelayingQueue does not implement it, so without ShutDown() we'd
		// leak the (*delayingType).waitingLoop.
		defer queue.ShutDown()

		// Wait for the store so we can emit the initial items and the sync event.
		store, err := r.storePromise.Await(subCtx)
		if err != nil {
			complete(err)
			return
		}

		r.mu.Lock()
		r.queues[subId] = queue

		// Append the initial set of keys to the queue + sentinel for the sync event.
		// We go through the queue instead of directly emitting to avoid a race between
		// the queue add and listing keys, plus to have one error handling path for Event.Done().
		keyIter := store.IterKeys()
		for keyIter.Next() {
			queue.AddUpdate(keyIter.Key())
		}
		queue.AddSync()
		r.mu.Unlock()

		for {
			// Retrieve an item from the subscribers queue and then fetch the object
			// from the store.
			raw, shutdown := queue.Get()
			if shutdown {
				break
			}
			entry := raw.(queueEntry)
			baseEvent := baseEvent{
				func(err error) { queue.eventDone(entry, err) },
			}
			switch entry := entry.(type) {
			case syncEntry:
				next(&SyncEvent[T]{baseEvent})
			case deleteEntry:
				next(&DeleteEvent[T]{baseEvent, entry.key, entry.obj.(T)})
			case updateEntry:
				obj, exists, err := store.GetByKey(entry.key)
				if err != nil {
					queue.setError(err)
					break
				}
				// Emit the update event if the item exists, if it doesn't, then
				// it has been deleted and a delete will follow soon.
				if exists {
					next(&UpdateEvent[T]{baseEvent, entry.key, obj})
				}
			}
		}
		r.mu.Lock()
		delete(r.queues, subId)
		r.mu.Unlock()
		complete(queue.getError())
	}()

	// Fork a goroutine to wait for either the subscriber cancelling or the resource
	// shutting down.
	r.wg.Add(1)
	go func() {
		defer r.wg.Done()
		select {
		case <-r.ctx.Done():
		case <-subCtx.Done():
		}
		subCancel()
		queue.ShutDownWithDrain()
	}()
}

// keyQueue wraps the workqueue to implement the error retry logic for a single subscriber,
// e.g. it implements the eventDone() method called by Event[T].Done().
type keyQueue struct {
	lock.Mutex
	err error

	workqueue.RateLimitingInterface
	errorHandler ErrorHandler
}

func (kq *keyQueue) AddSync() {
	kq.Add(syncEntry{})
}

func (kq *keyQueue) AddUpdate(key Key) {
	// The entries must be added by value and not by pointer in order for
	// them to be compared by value and not by pointer.
	kq.Add(updateEntry{key})
}

func (kq *keyQueue) AddDelete(key Key, obj any) {
	kq.Add(deleteEntry{key, obj})
}

func (kq *keyQueue) getError() error {
	kq.Lock()
	defer kq.Unlock()
	return kq.err
}

func (kq *keyQueue) setError(err error) {
	kq.Lock()
	defer kq.Unlock()
	if kq.err == nil {
		kq.err = err
	}
}

func (kq *keyQueue) eventDone(entry queueEntry, err error) {
	// This is based on the example found in k8s.io/client-go/examples/workqueue/main.go.

	// Mark the object as done being processed. If it was marked dirty
	// during processing, it'll be processed again.
	defer kq.Done(entry)

	if err != nil {
		numRequeues := kq.NumRequeues(entry)

		var action ErrorAction
		switch entry := entry.(type) {
		case syncEntry:
			action = ErrorActionStop
		case updateEntry:
			action = kq.errorHandler(entry.key, numRequeues, err)
		case deleteEntry:
			action = kq.errorHandler(entry.key, numRequeues, err)
		}

		switch action {
		case ErrorActionRetry:
			go kq.AddRateLimited(entry)
		case ErrorActionStop:
			kq.setError(err)
			kq.ShutDown()
		case ErrorActionIgnore:
			kq.Forget(entry)
		}
	} else {
		// As the object was processed successfully we can "forget" it.
		// This clears any rate limiter state associated with this object, so
		// it won't be throttled based on previous failure history.
		kq.Forget(entry)
	}
}

type queueEntry interface {
	isQueueEntry()
}

type syncEntry struct{}

func (syncEntry) isQueueEntry() {}

type updateEntry struct {
	key Key
}

func (updateEntry) isQueueEntry() {}

type deleteEntry struct {
	key Key
	obj any
}

func (deleteEntry) isQueueEntry() {}
<|endoftext|>


<commit_before>// Poseidon
// Copyright (c) The Poseidon 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
//
// THIS CODE IS PROVIDED ON AN *AS IS* BASIS, WITHOUT WARRANTIES OR
// CONDITIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING WITHOUT
// LIMITATION ANY IMPLIED WARRANTIES OR CONDITIONS OF TITLE, FITNESS FOR
// A PARTICULAR PURPOSE, MERCHANTABLITY OR NON-INFRINGEMENT.
//
// See the Apache Version 2.0 License for specific language governing
// permissions and limitations under the License.

package k8sclient

import (
	"fmt"
	"strconv"
	"sync"
	"time"

	"github.com/ICGog/poseidongo/pkg/firmament"
	"github.com/golang/glog"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	utilruntime "k8s.io/apimachinery/pkg/util/runtime"
	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/apimachinery/pkg/watch"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/pkg/api/v1"
	"k8s.io/client-go/tools/cache"
)

func NewNodeWatcher(client kubernetes.Interface, fc firmament.FirmamentSchedulerClient) *NodeWatcher {
	glog.Info("Starting NodeWatcher...")
	NodesCond = sync.NewCond(&sync.Mutex{})
	NodeToRTND = make(map[string]*firmament.ResourceTopologyNodeDescriptor)
	ResIDToNode = make(map[string]string)
	nodewatcher := &NodeWatcher{
		clientset: client,
		fc:        fc,
	}
	_, controller := cache.NewInformer(
		&cache.ListWatch{
			ListFunc: func(alo metav1.ListOptions) (runtime.Object, error) {
				return client.CoreV1().Nodes().List(alo)
			},
			WatchFunc: func(alo metav1.ListOptions) (watch.Interface, error) {
				return client.CoreV1().Nodes().Watch(alo)
			},
		},
		&v1.Node{},
		0,
		cache.ResourceEventHandlerFuncs{
			AddFunc: func(obj interface{}) {
				key, err := cache.MetaNamespaceKeyFunc(obj)
				if err != nil {
					glog.Errorf("AddFunc: error getting key %v", err)
				}
				nodewatcher.enqueueNodeAddition(key, obj)
			},
			UpdateFunc: func(old, new interface{}) {
				key, err := cache.MetaNamespaceKeyFunc(new)
				if err != nil {
					glog.Errorf("UpdateFunc: error getting key %v", err)
				}
				nodewatcher.enqueueNodeUpdate(key, old, new)
			},
			DeleteFunc: func(obj interface{}) {
				key, err := cache.MetaNamespaceKeyFunc(obj)
				if err != nil {
					glog.Errorf("DeleteFunc: error getting key %v", err)
				}
				nodewatcher.enqueueNodeDeletion(key, obj)
			},
		},
	)
	nodewatcher.controller = controller
	nodewatcher.nodeWorkQueue = NewKeyedQueue()
	return nodewatcher
}

func (this *NodeWatcher) enqueueNodeAddition(key, obj interface{}) {
	node := obj.(*v1.Node)
	if node.Spec.Unschedulable {
		glog.Info("enqueueNodeAddition: received an Unschedulable node", node.Name)
		return
	}
	isReady := false
	isOutOfDisk := false
	for _, cond := range node.Status.Conditions {
		switch cond.Type {
		case "OutOfDisk":
			isOutOfDisk = cond.Status == "True"
		case "Ready":
			isReady = cond.Status == "True"
		}
	}
	cpuCapQuantity := node.Status.Capacity["cpu"]
	cpuAllocQuantity := node.Status.Allocatable["cpu"]
	memCapQuantity := node.Status.Capacity["memory"]
	memCap, _ := memCapQuantity.AsInt64()
	memAllocQuantity := node.Status.Allocatable["memory"]
	memAlloc, _ := memAllocQuantity.AsInt64()
	addedNode := &Node{
		Hostname:         node.Name,
		Phase:            NodeAdded,
		IsReady:          isReady,
		IsOutOfDisk:      isOutOfDisk,
		CpuCapacity:      cpuCapQuantity.MilliValue(),
		CpuAllocatable:   cpuAllocQuantity.MilliValue(),
		MemCapacityKb:    memCap / bytesToKb,
		MemAllocatableKb: memAlloc / bytesToKb,
		Labels:           node.Labels,
		Annotations:      node.Annotations,
	}
	this.nodeWorkQueue.Add(key, addedNode)
	glog.Info("enqueueNodeAdition: Added node ", addedNode.Hostname)
}

func (this *NodeWatcher) enqueueNodeUpdate(key, oldObj, newObj interface{}) {
	// TODO(ionel): Implement!
}

func (this *NodeWatcher) enqueueNodeDeletion(key, obj interface{}) {
	node := obj.(*v1.Node)
	if node.Spec.Unschedulable {
		// Poseidon doesn't case about Unschedulable nodes.
		return
	}
	deletedNode := &Node{
		Hostname: node.Name,
		Phase:    NodeDeleted,
	}
	this.nodeWorkQueue.Add(key, deletedNode)
	glog.Info("enqueueNodeDeletion: Added node ", deletedNode.Hostname)
}

func (this *NodeWatcher) Run(stopCh <-chan struct{}, nWorkers int) {
	defer utilruntime.HandleCrash()

	// The workers can stop when we are done.
	defer this.nodeWorkQueue.ShutDown()
	defer glog.Info("Shutting down NodeWatcher")
	glog.Info("Geting node updates...")

	go this.controller.Run(stopCh)

	if !cache.WaitForCacheSync(stopCh, this.controller.HasSynced) {
		utilruntime.HandleError(fmt.Errorf("Timed out waiting for caches to sync"))
		return
	}

	glog.Info("Starting node watching workers")
	for i := 0; i < nWorkers; i++ {
		go wait.Until(this.nodeWorker, time.Second, stopCh)
	}

	<-stopCh
	glog.Info("Stopping node watcher")
}

func (this *NodeWatcher) nodeWorker() {
	for {
		func() {
			key, items, quit := this.nodeWorkQueue.Get()
			if quit {
				return
			}
			for _, item := range items {
				node := item.(*Node)
				switch node.Phase {
				case NodeAdded:
					NodesCond.L.Lock()
					rtnd := this.createResourceTopologyForNode(node)
					_, ok := NodeToRTND[node.Hostname]
					if ok {
						glog.Fatalf("Node %s already exists", node.Hostname)
					}
					NodeToRTND[node.Hostname] = rtnd
					ResIDToNode[rtnd.GetResourceDesc().GetUuid()] = node.Hostname
					NodesCond.L.Unlock()
					firmament.NodeAdded(this.fc, rtnd)
				case NodeDeleted:
					NodesCond.L.Lock()
					rtnd, ok := NodeToRTND[node.Hostname]
					NodesCond.L.Unlock()
					if !ok {
						glog.Fatalf("Node %s does not exist", node.Hostname)
					}
					resID := rtnd.GetResourceDesc().GetUuid()
					firmament.NodeRemoved(this.fc, &firmament.ResourceUID{ResourceUid: resID})
					NodesCond.L.Lock()
					delete(NodeToRTND, node.Hostname)
					delete(ResIDToNode, resID)
					NodesCond.L.Unlock()
				case NodeFailed:
					NodesCond.L.Lock()
					rtnd, ok := NodeToRTND[node.Hostname]
					NodesCond.L.Unlock()
					if !ok {
						glog.Fatalf("Node %s does not exist", node.Hostname)
					}
					resID := rtnd.GetResourceDesc().GetUuid()
					firmament.NodeFailed(this.fc, &firmament.ResourceUID{ResourceUid: resID})
					NodesCond.L.Lock()
					this.cleanResourceStateForNode(rtnd)
					delete(NodeToRTND, node.Hostname)
					delete(ResIDToNode, resID)
					NodesCond.L.Unlock()
				case NodeUpdated:
					// TODO(ionel): Handle update case.
				default:
					glog.Fatalf("Unexpected node %s phase %s", node.Hostname, node.Phase)
				}
			}
			defer this.nodeWorkQueue.Done(key)
		}()
	}
}

func (this *NodeWatcher) cleanResourceStateForNode(rtnd *firmament.ResourceTopologyNodeDescriptor) {
	delete(ResIDToNode, rtnd.GetResourceDesc().GetUuid())
	for _, childRTND := range rtnd.GetChildren() {
		this.cleanResourceStateForNode(childRTND)
	}
}

func (this *NodeWatcher) createResourceTopologyForNode(node *Node) *firmament.ResourceTopologyNodeDescriptor {
	resUuid := this.generateResourceID(node.Hostname)
	rtnd := &firmament.ResourceTopologyNodeDescriptor{
		ResourceDesc: &firmament.ResourceDescriptor{
			Uuid:         resUuid,
			Type:         firmament.ResourceDescriptor_RESOURCE_MACHINE,
			State:        firmament.ResourceDescriptor_RESOURCE_IDLE,
			FriendlyName: node.Hostname,
			ResourceCapacity: &firmament.ResourceVector{
				RamCap:   uint64(node.MemCapacityKb),
				CpuCores: float32(node.CpuCapacity),
			},
		},
	}
	ResIDToNode[resUuid] = node.Hostname
	// TODO(ionel) Add annotations.
	// Add labels.
	for label, value := range node.Labels {
		rtnd.ResourceDesc.Labels = append(rtnd.ResourceDesc.Labels,
			&firmament.Label{
				Key:   label,
				Value: value,
			})
	}
	// TODO(ionel): In the future, we want to get real node topology rather
	// than manually connecting PU RDs to the machine RD.
	for num_pu := int64(0); num_pu < node.CpuCapacity; num_pu++ {
		friendlyName := node.Hostname + "_pu" + strconv.FormatInt(num_pu, 10)
		puUuid := this.generateResourceID(friendlyName)
		puRtnd := &firmament.ResourceTopologyNodeDescriptor{
			ResourceDesc: &firmament.ResourceDescriptor{
				Uuid:         puUuid,
				Type:         firmament.ResourceDescriptor_RESOURCE_PU,
				State:        firmament.ResourceDescriptor_RESOURCE_IDLE,
				FriendlyName: friendlyName,
				Labels:       rtnd.ResourceDesc.Labels,
			},
			ParentId: resUuid,
		}
		rtnd.Children = append(rtnd.Children, puRtnd)
		ResIDToNode[puUuid] = node.Hostname
	}
	return rtnd
}

func (this *NodeWatcher) generateResourceID(seed string) string {
	return GenerateUUID(seed)
}
<commit_msg>Apply node updates.<commit_after>// Poseidon
// Copyright (c) The Poseidon 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
//
// THIS CODE IS PROVIDED ON AN *AS IS* BASIS, WITHOUT WARRANTIES OR
// CONDITIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING WITHOUT
// LIMITATION ANY IMPLIED WARRANTIES OR CONDITIONS OF TITLE, FITNESS FOR
// A PARTICULAR PURPOSE, MERCHANTABLITY OR NON-INFRINGEMENT.
//
// See the Apache Version 2.0 License for specific language governing
// permissions and limitations under the License.

package k8sclient

import (
	"fmt"
	"reflect"
	"strconv"
	"sync"
	"time"

	"github.com/ICGog/poseidongo/pkg/firmament"
	"github.com/golang/glog"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	utilruntime "k8s.io/apimachinery/pkg/util/runtime"
	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/apimachinery/pkg/watch"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/pkg/api/v1"
	"k8s.io/client-go/tools/cache"
)

func NewNodeWatcher(client kubernetes.Interface, fc firmament.FirmamentSchedulerClient) *NodeWatcher {
	glog.Info("Starting NodeWatcher...")
	NodesCond = sync.NewCond(&sync.Mutex{})
	NodeToRTND = make(map[string]*firmament.ResourceTopologyNodeDescriptor)
	ResIDToNode = make(map[string]string)
	nodewatcher := &NodeWatcher{
		clientset: client,
		fc:        fc,
	}
	_, controller := cache.NewInformer(
		&cache.ListWatch{
			ListFunc: func(alo metav1.ListOptions) (runtime.Object, error) {
				return client.CoreV1().Nodes().List(alo)
			},
			WatchFunc: func(alo metav1.ListOptions) (watch.Interface, error) {
				return client.CoreV1().Nodes().Watch(alo)
			},
		},
		&v1.Node{},
		0,
		cache.ResourceEventHandlerFuncs{
			AddFunc: func(obj interface{}) {
				key, err := cache.MetaNamespaceKeyFunc(obj)
				if err != nil {
					glog.Errorf("AddFunc: error getting key %v", err)
				}
				nodewatcher.enqueueNodeAddition(key, obj)
			},
			UpdateFunc: func(old, new interface{}) {
				key, err := cache.MetaNamespaceKeyFunc(new)
				if err != nil {
					glog.Errorf("UpdateFunc: error getting key %v", err)
				}
				nodewatcher.enqueueNodeUpdate(key, old, new)
			},
			DeleteFunc: func(obj interface{}) {
				key, err := cache.MetaNamespaceKeyFunc(obj)
				if err != nil {
					glog.Errorf("DeleteFunc: error getting key %v", err)
				}
				nodewatcher.enqueueNodeDeletion(key, obj)
			},
		},
	)
	nodewatcher.controller = controller
	nodewatcher.nodeWorkQueue = NewKeyedQueue()
	return nodewatcher
}

func (this *NodeWatcher) getReadyAndOutOfDiskConditions(node *v1.Node) (isReady bool, isOutOfDisk bool) {
	isReady = false
	isOutOfDisk = false
	for _, cond := range node.Status.Conditions {
		switch cond.Type {
		case "OutOfDisk":
			isOutOfDisk = cond.Status == "True"
		case "Ready":
			isReady = cond.Status == "True"
		}
	}
	return isReady, isOutOfDisk
}

func (this *NodeWatcher) parseNode(node *v1.Node, phase NodePhase) *Node {
	isReady, isOutOfDisk := this.getReadyAndOutOfDiskConditions(node)
	cpuCapQuantity := node.Status.Capacity["cpu"]
	cpuAllocQuantity := node.Status.Allocatable["cpu"]
	memCapQuantity := node.Status.Capacity["memory"]
	memCap, _ := memCapQuantity.AsInt64()
	memAllocQuantity := node.Status.Allocatable["memory"]
	memAlloc, _ := memAllocQuantity.AsInt64()
	return &Node{
		Hostname:         node.Name,
		Phase:            phase,
		IsReady:          isReady,
		IsOutOfDisk:      isOutOfDisk,
		CpuCapacity:      cpuCapQuantity.MilliValue(),
		CpuAllocatable:   cpuAllocQuantity.MilliValue(),
		MemCapacityKb:    memCap / bytesToKb,
		MemAllocatableKb: memAlloc / bytesToKb,
		Labels:           node.Labels,
		Annotations:      node.Annotations,
	}
}

func (this *NodeWatcher) enqueueNodeAddition(key, obj interface{}) {
	node := obj.(*v1.Node)
	if node.Spec.Unschedulable {
		glog.Info("enqueueNodeAddition: received an Unschedulable node", node.Name)
		return
	}
	addedNode := this.parseNode(node, NodeAdded)
	this.nodeWorkQueue.Add(key, addedNode)
	glog.Info("enqueueNodeAdition: Added node ", addedNode.Hostname)
}

func (this *NodeWatcher) enqueueNodeUpdate(key, oldObj, newObj interface{}) {
	oldNode := oldObj.(*v1.Node)
	newNode := newObj.(*v1.Node)
	if oldNode.Spec.Unschedulable != newNode.Spec.Unschedulable {
		if oldNode.Spec.Unschedulable {
			addedNode := this.parseNode(newNode, NodeAdded)
			this.nodeWorkQueue.Add(key, addedNode)
			glog.Info("enqueueNodeUpdate: Added node ", addedNode.Hostname)
			return
		} else {
			// Can not schedule pods on the node any more.
			deletedNode := this.parseNode(newNode, NodeDeleted)
			this.nodeWorkQueue.Add(key, deletedNode)
			glog.Info("enqueueNodeUpdate: Deleted node ", deletedNode.Hostname)
			return
		}
	}
	oldIsReady, oldIsOutOfDisk := this.getReadyAndOutOfDiskConditions(oldNode)
	newIsReady, newIsOutOfDisk := this.getReadyAndOutOfDiskConditions(newNode)

	if oldIsReady != newIsReady || oldIsOutOfDisk != newIsOutOfDisk {
		if newIsReady && !newIsOutOfDisk {
			addedNode := this.parseNode(newNode, NodeAdded)
			this.nodeWorkQueue.Add(key, addedNode)
			glog.Info("enqueueNodeUpdate: Added node ", addedNode.Hostname)
			return
		} else {
			failedNode := this.parseNode(newNode, NodeFailed)
			this.nodeWorkQueue.Add(key, failedNode)
			glog.Info("enqueueNodeUpdate: Failed node ", failedNode.Hostname)
			return
		}
	}
	nodeUpdated := false
	if !reflect.DeepEqual(oldNode.Labels, newNode.Labels) {
		nodeUpdated = true
	}
	if !reflect.DeepEqual(oldNode.Annotations, newNode.Annotations) {
		nodeUpdated = true
	}
	if nodeUpdated {
		updatedNode := this.parseNode(newNode, NodeUpdated)
		this.nodeWorkQueue.Add(key, updatedNode)
		glog.Info("enqueueNodeUpdate: Updated node ", updatedNode.Hostname)
	}
}

func (this *NodeWatcher) enqueueNodeDeletion(key, obj interface{}) {
	node := obj.(*v1.Node)
	if node.Spec.Unschedulable {
		// Poseidon doesn't case about Unschedulable nodes.
		return
	}
	deletedNode := &Node{
		Hostname: node.Name,
		Phase:    NodeDeleted,
	}
	this.nodeWorkQueue.Add(key, deletedNode)
	glog.Info("enqueueNodeDeletion: Added node ", deletedNode.Hostname)
}

func (this *NodeWatcher) Run(stopCh <-chan struct{}, nWorkers int) {
	defer utilruntime.HandleCrash()

	// The workers can stop when we are done.
	defer this.nodeWorkQueue.ShutDown()
	defer glog.Info("Shutting down NodeWatcher")
	glog.Info("Geting node updates...")

	go this.controller.Run(stopCh)

	if !cache.WaitForCacheSync(stopCh, this.controller.HasSynced) {
		utilruntime.HandleError(fmt.Errorf("Timed out waiting for caches to sync"))
		return
	}

	glog.Info("Starting node watching workers")
	for i := 0; i < nWorkers; i++ {
		go wait.Until(this.nodeWorker, time.Second, stopCh)
	}

	<-stopCh
	glog.Info("Stopping node watcher")
}

func (this *NodeWatcher) nodeWorker() {
	for {
		func() {
			key, items, quit := this.nodeWorkQueue.Get()
			if quit {
				return
			}
			for _, item := range items {
				node := item.(*Node)
				switch node.Phase {
				case NodeAdded:
					NodesCond.L.Lock()
					rtnd := this.createResourceTopologyForNode(node)
					_, ok := NodeToRTND[node.Hostname]
					if ok {
						glog.Fatalf("Node %s already exists", node.Hostname)
					}
					NodeToRTND[node.Hostname] = rtnd
					ResIDToNode[rtnd.GetResourceDesc().GetUuid()] = node.Hostname
					NodesCond.L.Unlock()
					firmament.NodeAdded(this.fc, rtnd)
				case NodeDeleted:
					NodesCond.L.Lock()
					rtnd, ok := NodeToRTND[node.Hostname]
					NodesCond.L.Unlock()
					if !ok {
						glog.Fatalf("Node %s does not exist", node.Hostname)
					}
					resID := rtnd.GetResourceDesc().GetUuid()
					firmament.NodeRemoved(this.fc, &firmament.ResourceUID{ResourceUid: resID})
					NodesCond.L.Lock()
					delete(NodeToRTND, node.Hostname)
					delete(ResIDToNode, resID)
					NodesCond.L.Unlock()
				case NodeFailed:
					NodesCond.L.Lock()
					rtnd, ok := NodeToRTND[node.Hostname]
					NodesCond.L.Unlock()
					if !ok {
						glog.Fatalf("Node %s does not exist", node.Hostname)
					}
					resID := rtnd.GetResourceDesc().GetUuid()
					firmament.NodeFailed(this.fc, &firmament.ResourceUID{ResourceUid: resID})
					NodesCond.L.Lock()
					this.cleanResourceStateForNode(rtnd)
					delete(NodeToRTND, node.Hostname)
					delete(ResIDToNode, resID)
					NodesCond.L.Unlock()
				case NodeUpdated:
					NodesCond.L.Lock()
					rtnd, ok := NodeToRTND[node.Hostname]
					NodesCond.L.Unlock()
					if !ok {
						glog.Fatalf("Node %s does not exist", node.Hostname)
					}
					// TODO(ionel): Handle update case.
				default:
					glog.Fatalf("Unexpected node %s phase %s", node.Hostname, node.Phase)
				}
			}
			defer this.nodeWorkQueue.Done(key)
		}()
	}
}

func (this *NodeWatcher) cleanResourceStateForNode(rtnd *firmament.ResourceTopologyNodeDescriptor) {
	delete(ResIDToNode, rtnd.GetResourceDesc().GetUuid())
	for _, childRTND := range rtnd.GetChildren() {
		this.cleanResourceStateForNode(childRTND)
	}
}

func (this *NodeWatcher) createResourceTopologyForNode(node *Node) *firmament.ResourceTopologyNodeDescriptor {
	resUuid := this.generateResourceID(node.Hostname)
	rtnd := &firmament.ResourceTopologyNodeDescriptor{
		ResourceDesc: &firmament.ResourceDescriptor{
			Uuid:         resUuid,
			Type:         firmament.ResourceDescriptor_RESOURCE_MACHINE,
			State:        firmament.ResourceDescriptor_RESOURCE_IDLE,
			FriendlyName: node.Hostname,
			ResourceCapacity: &firmament.ResourceVector{
				RamCap:   uint64(node.MemCapacityKb),
				CpuCores: float32(node.CpuCapacity),
			},
		},
	}
	ResIDToNode[resUuid] = node.Hostname
	// TODO(ionel) Add annotations.
	// Add labels.
	for label, value := range node.Labels {
		rtnd.ResourceDesc.Labels = append(rtnd.ResourceDesc.Labels,
			&firmament.Label{
				Key:   label,
				Value: value,
			})
	}
	// TODO(ionel): In the future, we want to get real node topology rather
	// than manually connecting PU RDs to the machine RD.
	for num_pu := int64(0); num_pu < node.CpuCapacity; num_pu++ {
		friendlyName := node.Hostname + "_pu" + strconv.FormatInt(num_pu, 10)
		puUuid := this.generateResourceID(friendlyName)
		puRtnd := &firmament.ResourceTopologyNodeDescriptor{
			ResourceDesc: &firmament.ResourceDescriptor{
				Uuid:         puUuid,
				Type:         firmament.ResourceDescriptor_RESOURCE_PU,
				State:        firmament.ResourceDescriptor_RESOURCE_IDLE,
				FriendlyName: friendlyName,
				Labels:       rtnd.ResourceDesc.Labels,
			},
			ParentId: resUuid,
		}
		rtnd.Children = append(rtnd.Children, puRtnd)
		ResIDToNode[puUuid] = node.Hostname
	}
	return rtnd
}

func (this *NodeWatcher) generateResourceID(seed string) string {
	return GenerateUUID(seed)
}
<|endoftext|>


<commit_before>/*
Copyright 2014 The Kubernetes Authors.

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 logs

import (
	"errors"
	"fmt"
	"io"
	"os"
	"time"

	"github.com/spf13/cobra"

	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/cli-runtime/pkg/genericclioptions"
	"k8s.io/client-go/rest"
	cmdutil "k8s.io/kubernetes/pkg/kubectl/cmd/util"
	"k8s.io/kubernetes/pkg/kubectl/polymorphichelpers"
	"k8s.io/kubernetes/pkg/kubectl/scheme"
	"k8s.io/kubernetes/pkg/kubectl/util"
	"k8s.io/kubernetes/pkg/kubectl/util/i18n"
	"k8s.io/kubernetes/pkg/kubectl/util/templates"
)

const (
	logsUsageStr = "logs [-f] [-p] (POD | TYPE/NAME) [-c CONTAINER]"
)

var (
	logsExample = templates.Examples(i18n.T(`
		# Return snapshot logs from pod nginx with only one container
		kubectl logs nginx

		# Return snapshot logs from pod nginx with multi containers
		kubectl logs nginx --all-containers=true

		# Return snapshot logs from all containers in pods defined by label app=nginx
		kubectl logs -lapp=nginx --all-containers=true

		# Return snapshot of previous terminated ruby container logs from pod web-1
		kubectl logs -p -c ruby web-1

		# Begin streaming the logs of the ruby container in pod web-1
		kubectl logs -f -c ruby web-1

		# Display only the most recent 20 lines of output in pod nginx
		kubectl logs --tail=20 nginx

		# Show all logs from pod nginx written in the last hour
		kubectl logs --since=1h nginx

		# Return snapshot logs from first container of a job named hello
		kubectl logs job/hello

		# Return snapshot logs from container nginx-1 of a deployment named nginx
		kubectl logs deployment/nginx -c nginx-1`))

	selectorTail    int64 = 10
	logsUsageErrStr       = fmt.Sprintf("expected '%s'.\nPOD or TYPE/NAME is a required argument for the logs command", logsUsageStr)
)

const (
	defaultPodLogsTimeout = 20 * time.Second
)

type LogsOptions struct {
	Namespace     string
	ResourceArg   string
	AllContainers bool
	Options       runtime.Object
	Resources     []string

	ConsumeRequestFn func(*rest.Request, io.Writer) error

	// PodLogOptions
	SinceTime    string
	SinceSeconds time.Duration
	Follow       bool
	Previous     bool
	Timestamps   bool
	LimitBytes   int64
	Tail         int64
	Container    string

	// whether or not a container name was given via --container
	ContainerNameSpecified bool
	Selector               string

	Object           runtime.Object
	GetPodTimeout    time.Duration
	RESTClientGetter genericclioptions.RESTClientGetter
	LogsForObject    polymorphichelpers.LogsForObjectFunc

	genericclioptions.IOStreams
}

func NewLogsOptions(streams genericclioptions.IOStreams, allContainers bool) *LogsOptions {
	return &LogsOptions{
		IOStreams:     streams,
		AllContainers: allContainers,
		Tail:          -1,
	}
}

// NewCmdLogs creates a new pod logs command
func NewCmdLogs(f cmdutil.Factory, streams genericclioptions.IOStreams) *cobra.Command {
	o := NewLogsOptions(streams, false)

	cmd := &cobra.Command{
		Use:                   logsUsageStr,
		DisableFlagsInUseLine: true,
		Short:                 i18n.T("Print the logs for a container in a pod"),
		Long:                  "Print the logs for a container in a pod or specified resource. If the pod has only one container, the container name is optional.",
		Example:               logsExample,
		PreRun: func(cmd *cobra.Command, args []string) {
			if len(os.Args) > 1 && os.Args[1] == "log" {
				fmt.Fprintf(o.ErrOut, "%s is DEPRECATED and will be removed in a future version. Use %s instead.\n", "log", "logs")
			}
		},
		Run: func(cmd *cobra.Command, args []string) {
			cmdutil.CheckErr(o.Complete(f, cmd, args))
			cmdutil.CheckErr(o.Validate())
			cmdutil.CheckErr(o.RunLogs())
		},
		Aliases: []string{"log"},
	}
	cmd.Flags().BoolVar(&o.AllContainers, "all-containers", o.AllContainers, "Get all containers's logs in the pod(s).")
	cmd.Flags().BoolVarP(&o.Follow, "follow", "f", o.Follow, "Specify if the logs should be streamed.")
	cmd.Flags().BoolVar(&o.Timestamps, "timestamps", o.Timestamps, "Include timestamps on each line in the log output")
	cmd.Flags().Int64Var(&o.LimitBytes, "limit-bytes", o.LimitBytes, "Maximum bytes of logs to return. Defaults to no limit.")
	cmd.Flags().BoolVarP(&o.Previous, "previous", "p", o.Previous, "If true, print the logs for the previous instance of the container in a pod if it exists.")
	cmd.Flags().Int64Var(&o.Tail, "tail", o.Tail, "Lines of recent log file to display. Defaults to -1 with no selector, showing all log lines otherwise 10, if a selector is provided.")
	cmd.Flags().StringVar(&o.SinceTime, "since-time", o.SinceTime, i18n.T("Only return logs after a specific date (RFC3339). Defaults to all logs. Only one of since-time / since may be used."))
	cmd.Flags().DurationVar(&o.SinceSeconds, "since", o.SinceSeconds, "Only return logs newer than a relative duration like 5s, 2m, or 3h. Defaults to all logs. Only one of since-time / since may be used.")
	cmd.Flags().StringVarP(&o.Container, "container", "c", o.Container, "Print the logs of this container")
	cmdutil.AddPodRunningTimeoutFlag(cmd, defaultPodLogsTimeout)
	cmd.Flags().StringVarP(&o.Selector, "selector", "l", o.Selector, "Selector (label query) to filter on.")
	return cmd
}

func (o *LogsOptions) ToLogOptions() (*corev1.PodLogOptions, error) {
	logOptions := &corev1.PodLogOptions{
		Container:  o.Container,
		Follow:     o.Follow,
		Previous:   o.Previous,
		Timestamps: o.Timestamps,
	}

	if len(o.SinceTime) > 0 {
		t, err := util.ParseRFC3339(o.SinceTime, metav1.Now)
		if err != nil {
			return nil, err
		}

		logOptions.SinceTime = &t
	}

	if o.LimitBytes != 0 {
		logOptions.LimitBytes = &o.LimitBytes
	}

	if o.SinceSeconds != 0 {
		// round up to the nearest second
		sec := int64(o.SinceSeconds.Round(time.Second).Seconds())
		logOptions.SinceSeconds = &sec
	}

	if len(o.Selector) > 0 && o.Tail == -1 {
		logOptions.TailLines = &selectorTail
	} else if o.Tail != -1 {
		logOptions.TailLines = &o.Tail
	}

	return logOptions, nil
}

func (o *LogsOptions) Complete(f cmdutil.Factory, cmd *cobra.Command, args []string) error {
	o.ContainerNameSpecified = cmd.Flag("container").Changed
	o.Resources = args

	switch len(args) {
	case 0:
		if len(o.Selector) == 0 {
			return cmdutil.UsageErrorf(cmd, "%s", logsUsageErrStr)
		}
	case 1:
		o.ResourceArg = args[0]
		if len(o.Selector) != 0 {
			return cmdutil.UsageErrorf(cmd, "only a selector (-l) or a POD name is allowed")
		}
	case 2:
		o.ResourceArg = args[0]
		o.Container = args[1]
	default:
		return cmdutil.UsageErrorf(cmd, "%s", logsUsageErrStr)
	}
	var err error
	o.Namespace, _, err = f.ToRawKubeConfigLoader().Namespace()
	if err != nil {
		return err
	}

	o.ConsumeRequestFn = DefaultConsumeRequest

	o.GetPodTimeout, err = cmdutil.GetPodRunningTimeoutFlag(cmd)
	if err != nil {
		return err
	}

	o.Options, err = o.ToLogOptions()
	if err != nil {
		return err
	}

	o.RESTClientGetter = f
	o.LogsForObject = polymorphichelpers.LogsForObjectFn

	if o.Object == nil {
		builder := f.NewBuilder().
			WithScheme(scheme.Scheme, scheme.Scheme.PrioritizedVersionsAllGroups()...).
			NamespaceParam(o.Namespace).DefaultNamespace().
			SingleResourceType()
		if o.ResourceArg != "" {
			builder.ResourceNames("pods", o.ResourceArg)
		}
		if o.Selector != "" {
			builder.ResourceTypes("pods").LabelSelectorParam(o.Selector)
		}
		infos, err := builder.Do().Infos()
		if err != nil {
			return err
		}
		if o.Selector == "" && len(infos) != 1 {
			return errors.New("expected a resource")
		}
		o.Object = infos[0].Object
	}

	return nil
}

func (o LogsOptions) Validate() error {
	if o.Follow && len(o.Selector) > 0 {
		return fmt.Errorf("only one of follow (-f) or selector (-l) is allowed")
	}

	if len(o.SinceTime) > 0 && o.SinceSeconds != 0 {
		return fmt.Errorf("at most one of `sinceTime` or `sinceSeconds` may be specified")
	}

	logsOptions, ok := o.Options.(*corev1.PodLogOptions)
	if !ok {
		return errors.New("unexpected logs options object")
	}
	if o.AllContainers && len(logsOptions.Container) > 0 {
		return fmt.Errorf("--all-containers=true should not be specified with container name %s", logsOptions.Container)
	}

	if o.ContainerNameSpecified && len(o.Resources) == 2 {
		return fmt.Errorf("only one of -c or an inline [CONTAINER] arg is allowed")
	}

	if o.LimitBytes < 0 {
		return fmt.Errorf("--limit-bytes must be greater than 0")
	}

	if logsOptions.SinceSeconds != nil && *logsOptions.SinceSeconds < int64(0) {
		return fmt.Errorf("--since must be greater than 0")
	}

	if logsOptions.TailLines != nil && *logsOptions.TailLines < 0 {
		return fmt.Errorf("TailLines must be greater than or equal to 0")
	}

	return nil
}

// RunLogs retrieves a pod log
func (o LogsOptions) RunLogs() error {
	requests, err := o.LogsForObject(o.RESTClientGetter, o.Object, o.Options, o.GetPodTimeout, o.AllContainers)
	if err != nil {
		return err
	}

	for _, request := range requests {
		if err := o.ConsumeRequestFn(request, o.Out); err != nil {
			return err
		}
	}

	return nil
}

func DefaultConsumeRequest(request *rest.Request, out io.Writer) error {
	readCloser, err := request.Stream()
	if err != nil {
		return err
	}
	defer readCloser.Close()

	_, err = io.Copy(out, readCloser)
	return err
}
<commit_msg>word<commit_after>/*
Copyright 2014 The Kubernetes Authors.

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 logs

import (
	"errors"
	"fmt"
	"io"
	"os"
	"time"

	"github.com/spf13/cobra"

	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/cli-runtime/pkg/genericclioptions"
	"k8s.io/client-go/rest"
	cmdutil "k8s.io/kubernetes/pkg/kubectl/cmd/util"
	"k8s.io/kubernetes/pkg/kubectl/polymorphichelpers"
	"k8s.io/kubernetes/pkg/kubectl/scheme"
	"k8s.io/kubernetes/pkg/kubectl/util"
	"k8s.io/kubernetes/pkg/kubectl/util/i18n"
	"k8s.io/kubernetes/pkg/kubectl/util/templates"
)

const (
	logsUsageStr = "logs [-f] [-p] (POD | TYPE/NAME) [-c CONTAINER]"
)

var (
	logsExample = templates.Examples(i18n.T(`
		# Return snapshot logs from pod nginx with only one container
		kubectl logs nginx

		# Return snapshot logs from pod nginx with multi containers
		kubectl logs nginx --all-containers=true

		# Return snapshot logs from all containers in pods defined by label app=nginx
		kubectl logs -lapp=nginx --all-containers=true

		# Return snapshot of previous terminated ruby container logs from pod web-1
		kubectl logs -p -c ruby web-1

		# Begin streaming the logs of the ruby container in pod web-1
		kubectl logs -f -c ruby web-1

		# Display only the most recent 20 lines of output in pod nginx
		kubectl logs --tail=20 nginx

		# Show all logs from pod nginx written in the last hour
		kubectl logs --since=1h nginx

		# Return snapshot logs from first container of a job named hello
		kubectl logs job/hello

		# Return snapshot logs from container nginx-1 of a deployment named nginx
		kubectl logs deployment/nginx -c nginx-1`))

	selectorTail    int64 = 10
	logsUsageErrStr       = fmt.Sprintf("expected '%s'.\nPOD or TYPE/NAME is a required argument for the logs command", logsUsageStr)
)

const (
	defaultPodLogsTimeout = 20 * time.Second
)

type LogsOptions struct {
	Namespace     string
	ResourceArg   string
	AllContainers bool
	Options       runtime.Object
	Resources     []string

	ConsumeRequestFn func(*rest.Request, io.Writer) error

	// PodLogOptions
	SinceTime    string
	SinceSeconds time.Duration
	Follow       bool
	Previous     bool
	Timestamps   bool
	LimitBytes   int64
	Tail         int64
	Container    string

	// whether or not a container name was given via --container
	ContainerNameSpecified bool
	Selector               string

	Object           runtime.Object
	GetPodTimeout    time.Duration
	RESTClientGetter genericclioptions.RESTClientGetter
	LogsForObject    polymorphichelpers.LogsForObjectFunc

	genericclioptions.IOStreams
}

func NewLogsOptions(streams genericclioptions.IOStreams, allContainers bool) *LogsOptions {
	return &LogsOptions{
		IOStreams:     streams,
		AllContainers: allContainers,
		Tail:          -1,
	}
}

// NewCmdLogs creates a new pod logs command
func NewCmdLogs(f cmdutil.Factory, streams genericclioptions.IOStreams) *cobra.Command {
	o := NewLogsOptions(streams, false)

	cmd := &cobra.Command{
		Use:                   logsUsageStr,
		DisableFlagsInUseLine: true,
		Short:                 i18n.T("Print the logs for a container in a pod"),
		Long:                  "Print the logs for a container in a pod or specified resource. If the pod has only one container, the container name is optional.",
		Example:               logsExample,
		PreRun: func(cmd *cobra.Command, args []string) {
			if len(os.Args) > 1 && os.Args[1] == "log" {
				fmt.Fprintf(o.ErrOut, "%s is DEPRECATED and will be removed in a future version. Use %s instead.\n", "log", "logs")
			}
		},
		Run: func(cmd *cobra.Command, args []string) {
			cmdutil.CheckErr(o.Complete(f, cmd, args))
			cmdutil.CheckErr(o.Validate())
			cmdutil.CheckErr(o.RunLogs())
		},
		Aliases: []string{"log"},
	}
	cmd.Flags().BoolVar(&o.AllContainers, "all-containers", o.AllContainers, "Get all containers' logs in the pod(s).")
	cmd.Flags().BoolVarP(&o.Follow, "follow", "f", o.Follow, "Specify if the logs should be streamed.")
	cmd.Flags().BoolVar(&o.Timestamps, "timestamps", o.Timestamps, "Include timestamps on each line in the log output")
	cmd.Flags().Int64Var(&o.LimitBytes, "limit-bytes", o.LimitBytes, "Maximum bytes of logs to return. Defaults to no limit.")
	cmd.Flags().BoolVarP(&o.Previous, "previous", "p", o.Previous, "If true, print the logs for the previous instance of the container in a pod if it exists.")
	cmd.Flags().Int64Var(&o.Tail, "tail", o.Tail, "Lines of recent log file to display. Defaults to -1 with no selector, showing all log lines otherwise 10, if a selector is provided.")
	cmd.Flags().StringVar(&o.SinceTime, "since-time", o.SinceTime, i18n.T("Only return logs after a specific date (RFC3339). Defaults to all logs. Only one of since-time / since may be used."))
	cmd.Flags().DurationVar(&o.SinceSeconds, "since", o.SinceSeconds, "Only return logs newer than a relative duration like 5s, 2m, or 3h. Defaults to all logs. Only one of since-time / since may be used.")
	cmd.Flags().StringVarP(&o.Container, "container", "c", o.Container, "Print the logs of this container")
	cmdutil.AddPodRunningTimeoutFlag(cmd, defaultPodLogsTimeout)
	cmd.Flags().StringVarP(&o.Selector, "selector", "l", o.Selector, "Selector (label query) to filter on.")
	return cmd
}

func (o *LogsOptions) ToLogOptions() (*corev1.PodLogOptions, error) {
	logOptions := &corev1.PodLogOptions{
		Container:  o.Container,
		Follow:     o.Follow,
		Previous:   o.Previous,
		Timestamps: o.Timestamps,
	}

	if len(o.SinceTime) > 0 {
		t, err := util.ParseRFC3339(o.SinceTime, metav1.Now)
		if err != nil {
			return nil, err
		}

		logOptions.SinceTime = &t
	}

	if o.LimitBytes != 0 {
		logOptions.LimitBytes = &o.LimitBytes
	}

	if o.SinceSeconds != 0 {
		// round up to the nearest second
		sec := int64(o.SinceSeconds.Round(time.Second).Seconds())
		logOptions.SinceSeconds = &sec
	}

	if len(o.Selector) > 0 && o.Tail == -1 {
		logOptions.TailLines = &selectorTail
	} else if o.Tail != -1 {
		logOptions.TailLines = &o.Tail
	}

	return logOptions, nil
}

func (o *LogsOptions) Complete(f cmdutil.Factory, cmd *cobra.Command, args []string) error {
	o.ContainerNameSpecified = cmd.Flag("container").Changed
	o.Resources = args

	switch len(args) {
	case 0:
		if len(o.Selector) == 0 {
			return cmdutil.UsageErrorf(cmd, "%s", logsUsageErrStr)
		}
	case 1:
		o.ResourceArg = args[0]
		if len(o.Selector) != 0 {
			return cmdutil.UsageErrorf(cmd, "only a selector (-l) or a POD name is allowed")
		}
	case 2:
		o.ResourceArg = args[0]
		o.Container = args[1]
	default:
		return cmdutil.UsageErrorf(cmd, "%s", logsUsageErrStr)
	}
	var err error
	o.Namespace, _, err = f.ToRawKubeConfigLoader().Namespace()
	if err != nil {
		return err
	}

	o.ConsumeRequestFn = DefaultConsumeRequest

	o.GetPodTimeout, err = cmdutil.GetPodRunningTimeoutFlag(cmd)
	if err != nil {
		return err
	}

	o.Options, err = o.ToLogOptions()
	if err != nil {
		return err
	}

	o.RESTClientGetter = f
	o.LogsForObject = polymorphichelpers.LogsForObjectFn

	if o.Object == nil {
		builder := f.NewBuilder().
			WithScheme(scheme.Scheme, scheme.Scheme.PrioritizedVersionsAllGroups()...).
			NamespaceParam(o.Namespace).DefaultNamespace().
			SingleResourceType()
		if o.ResourceArg != "" {
			builder.ResourceNames("pods", o.ResourceArg)
		}
		if o.Selector != "" {
			builder.ResourceTypes("pods").LabelSelectorParam(o.Selector)
		}
		infos, err := builder.Do().Infos()
		if err != nil {
			return err
		}
		if o.Selector == "" && len(infos) != 1 {
			return errors.New("expected a resource")
		}
		o.Object = infos[0].Object
	}

	return nil
}

func (o LogsOptions) Validate() error {
	if o.Follow && len(o.Selector) > 0 {
		return fmt.Errorf("only one of follow (-f) or selector (-l) is allowed")
	}

	if len(o.SinceTime) > 0 && o.SinceSeconds != 0 {
		return fmt.Errorf("at most one of `sinceTime` or `sinceSeconds` may be specified")
	}

	logsOptions, ok := o.Options.(*corev1.PodLogOptions)
	if !ok {
		return errors.New("unexpected logs options object")
	}
	if o.AllContainers && len(logsOptions.Container) > 0 {
		return fmt.Errorf("--all-containers=true should not be specified with container name %s", logsOptions.Container)
	}

	if o.ContainerNameSpecified && len(o.Resources) == 2 {
		return fmt.Errorf("only one of -c or an inline [CONTAINER] arg is allowed")
	}

	if o.LimitBytes < 0 {
		return fmt.Errorf("--limit-bytes must be greater than 0")
	}

	if logsOptions.SinceSeconds != nil && *logsOptions.SinceSeconds < int64(0) {
		return fmt.Errorf("--since must be greater than 0")
	}

	if logsOptions.TailLines != nil && *logsOptions.TailLines < 0 {
		return fmt.Errorf("TailLines must be greater than or equal to 0")
	}

	return nil
}

// RunLogs retrieves a pod log
func (o LogsOptions) RunLogs() error {
	requests, err := o.LogsForObject(o.RESTClientGetter, o.Object, o.Options, o.GetPodTimeout, o.AllContainers)
	if err != nil {
		return err
	}

	for _, request := range requests {
		if err := o.ConsumeRequestFn(request, o.Out); err != nil {
			return err
		}
	}

	return nil
}

func DefaultConsumeRequest(request *rest.Request, out io.Writer) error {
	readCloser, err := request.Stream()
	if err != nil {
		return err
	}
	defer readCloser.Close()

	_, err = io.Copy(out, readCloser)
	return err
}
<|endoftext|>


<commit_before>package operator

import (
	"context"
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"time"

	"github.com/coreos-inc/vault-operator/pkg/spec"
	"github.com/coreos-inc/vault-operator/pkg/util/k8sutil"

	"github.com/Sirupsen/logrus"
	vaultapi "github.com/hashicorp/vault/api"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/labels"
)

// monitorAndUpdateStaus monitors the vault service and replicas statuses, and
// updates the status resrouce in the vault CR item.
func (vs *Vaults) monitorAndUpdateStaus(ctx context.Context, vr *spec.Vault) {
	tlsConfig, err := vs.readClientTLSFromSecret(vr)
	if err != nil {
		panic(fmt.Errorf("failed to read TLS config for vault client: %v", err))
	}

	s := spec.VaultStatus{}

	for {
		select {
		case err := <-ctx.Done():
			logrus.Infof("stopped monitoring vault: %s (%v)", vr.GetName(), err)
			return
		case <-time.After(10 * time.Second):
		}

		vs.updateLocalVaultCRStatus(ctx, vr.GetName(), vr.GetNamespace(), &s, tlsConfig)

		err := vs.updateVaultCRStatus(ctx, vr.GetName(), vr.GetNamespace(), s)
		if err != nil {
			logrus.Errorf("failed updating the status for the vault service: %s (%v)", vr.GetName(), err)
		}
	}
}

// updateLocalVaultCRStatus updates local vault CR status by querying each vault pod's API.
func (vs *Vaults) updateLocalVaultCRStatus(ctx context.Context, name, namespace string, s *spec.VaultStatus, tlsConfig *vaultapi.TLSConfig) {
	sel := k8sutil.LabelsForVault(name)
	// TODO: handle upgrades when pods from two replicaset can co-exist :(
	opt := metav1.ListOptions{LabelSelector: labels.SelectorFromSet(sel).String()}
	pods, err := vs.kubecli.CoreV1().Pods(namespace).List(opt)
	if err != nil {
		logrus.Errorf("failed to update vault replica status: failed listing pods for the vault service (%s.%s): %v", name, namespace, err)
		return
	}

	var sealNodes []string
	// If it can't talk to any vault pod, we are not changing the state.
	inited := s.Initialized

	for _, p := range pods.Items {
		cfg := vaultapi.DefaultConfig()
		podURL := fmt.Sprintf("https://%s:8200", k8sutil.PodDNSName(p.Status.PodIP, namespace))
		cfg.Address = podURL
		cfg.ConfigureTLS(tlsConfig)
		vapi, err := vaultapi.NewClient(cfg)
		if err != nil {
			logrus.Errorf("failed to update vault replica status: failed creating client for the vault pod (%s/%s): %v", namespace, p.GetName(), err)
			continue
		}

		hr, err := vapi.Sys().Health()
		if err != nil {
			logrus.Errorf("failed to update vault replica status: failed requesting health info for the vault pod (%s/%s): %v", namespace, p.GetName(), err)
			continue
		}
		// is active node?
		// TODO: add to vaultutil?
		if hr.Initialized && !hr.Sealed && !hr.Standby {
			s.ActiveNode = podURL
		}
		if hr.Sealed {
			sealNodes = append(sealNodes, podURL)
		}
		if hr.Initialized {
			inited = true
		}
	}

	s.SealedNodes = sealNodes
	s.Initialized = inited
}

// updateVaultCRStatus updates the status field of the Vault CR.
func (vs *Vaults) updateVaultCRStatus(ctx context.Context, name, namespace string, status spec.VaultStatus) error {
	vault, err := vs.vaultsCRCli.Get(ctx, namespace, name)
	if err != nil {
		return err
	}
	vault.Status = status
	_, err = vs.vaultsCRCli.Update(ctx, vault)
	return err
}

func (vs *Vaults) readClientTLSFromSecret(vr *spec.Vault) (*vaultapi.TLSConfig, error) {
	secret, err := vs.kubecli.CoreV1().Secrets(vr.GetNamespace()).Get(vr.Spec.TLS.Static.ClientSecret, metav1.GetOptions{})
	if err != nil {
		return nil, fmt.Errorf("read client tls failed: failed to get secret (%s): %v", vr.Spec.TLS.Static.ClientSecret, err)
	}

	// Read the secret and write ca.crt to a temporary file
	caCertData := secret.Data[spec.CATLSCertName]
	if err := os.MkdirAll(k8sutil.VaultTLSAssetDir, 0700); err != nil {
		return nil, fmt.Errorf("read client tls failed: failed to make dir: %v", err)
	}
	caCertFile := path.Join(k8sutil.VaultTLSAssetDir, spec.CATLSCertName)
	err = ioutil.WriteFile(caCertFile, caCertData, 0600)
	if err != nil {
		return nil, fmt.Errorf("read client tls failed: write ca cert file failed: %v", err)
	}
	return &vaultapi.TLSConfig{CACert: caCertFile}, nil
}
<commit_msg>pkg/operator: add available nodes (#82)<commit_after>package operator

import (
	"context"
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"time"

	"github.com/coreos-inc/vault-operator/pkg/spec"
	"github.com/coreos-inc/vault-operator/pkg/util/k8sutil"

	"github.com/Sirupsen/logrus"
	vaultapi "github.com/hashicorp/vault/api"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/labels"
)

// monitorAndUpdateStaus monitors the vault service and replicas statuses, and
// updates the status resrouce in the vault CR item.
func (vs *Vaults) monitorAndUpdateStaus(ctx context.Context, vr *spec.Vault) {
	tlsConfig, err := vs.readClientTLSFromSecret(vr)
	if err != nil {
		panic(fmt.Errorf("failed to read TLS config for vault client: %v", err))
	}

	s := spec.VaultStatus{}

	for {
		select {
		case err := <-ctx.Done():
			logrus.Infof("stopped monitoring vault: %s (%v)", vr.GetName(), err)
			return
		case <-time.After(10 * time.Second):
		}

		vs.updateLocalVaultCRStatus(ctx, vr.GetName(), vr.GetNamespace(), &s, tlsConfig)

		err := vs.updateVaultCRStatus(ctx, vr.GetName(), vr.GetNamespace(), s)
		if err != nil {
			logrus.Errorf("failed updating the status for the vault service: %s (%v)", vr.GetName(), err)
		}
	}
}

// updateLocalVaultCRStatus updates local vault CR status by querying each vault pod's API.
func (vs *Vaults) updateLocalVaultCRStatus(ctx context.Context, name, namespace string, s *spec.VaultStatus, tlsConfig *vaultapi.TLSConfig) {
	sel := k8sutil.LabelsForVault(name)
	// TODO: handle upgrades when pods from two replicaset can co-exist :(
	opt := metav1.ListOptions{LabelSelector: labels.SelectorFromSet(sel).String()}
	pods, err := vs.kubecli.CoreV1().Pods(namespace).List(opt)
	if err != nil {
		logrus.Errorf("failed to update vault replica status: failed listing pods for the vault service (%s.%s): %v", name, namespace, err)
		return
	}

	var sealNodes []string
	var availableNodes []string
	// If it can't talk to any vault pod, we are not changing the state.
	inited := s.Initialized

	for _, p := range pods.Items {
		availableNodes = append(availableNodes, p.GetName())

		cfg := vaultapi.DefaultConfig()
		podURL := fmt.Sprintf("https://%s:8200", k8sutil.PodDNSName(p.Status.PodIP, namespace))
		cfg.Address = podURL
		cfg.ConfigureTLS(tlsConfig)
		vapi, err := vaultapi.NewClient(cfg)
		if err != nil {
			logrus.Errorf("failed to update vault replica status: failed creating client for the vault pod (%s/%s): %v", namespace, p.GetName(), err)
			continue
		}

		hr, err := vapi.Sys().Health()
		if err != nil {
			logrus.Errorf("failed to update vault replica status: failed requesting health info for the vault pod (%s/%s): %v", namespace, p.GetName(), err)
			continue
		}
		// is active node?
		// TODO: add to vaultutil?
		if hr.Initialized && !hr.Sealed && !hr.Standby {
			s.ActiveNode = p.GetName()
		}
		if hr.Sealed {
			sealNodes = append(sealNodes, p.GetName())
		}
		if hr.Initialized {
			inited = true
		}
	}

	s.AvailableNodes = availableNodes
	s.SealedNodes = sealNodes
	s.Initialized = inited
}

// updateVaultCRStatus updates the status field of the Vault CR.
func (vs *Vaults) updateVaultCRStatus(ctx context.Context, name, namespace string, status spec.VaultStatus) error {
	vault, err := vs.vaultsCRCli.Get(ctx, namespace, name)
	if err != nil {
		return err
	}
	vault.Status = status
	_, err = vs.vaultsCRCli.Update(ctx, vault)
	return err
}

func (vs *Vaults) readClientTLSFromSecret(vr *spec.Vault) (*vaultapi.TLSConfig, error) {
	secret, err := vs.kubecli.CoreV1().Secrets(vr.GetNamespace()).Get(vr.Spec.TLS.Static.ClientSecret, metav1.GetOptions{})
	if err != nil {
		return nil, fmt.Errorf("read client tls failed: failed to get secret (%s): %v", vr.Spec.TLS.Static.ClientSecret, err)
	}

	// Read the secret and write ca.crt to a temporary file
	caCertData := secret.Data[spec.CATLSCertName]
	if err := os.MkdirAll(k8sutil.VaultTLSAssetDir, 0700); err != nil {
		return nil, fmt.Errorf("read client tls failed: failed to make dir: %v", err)
	}
	caCertFile := path.Join(k8sutil.VaultTLSAssetDir, spec.CATLSCertName)
	err = ioutil.WriteFile(caCertFile, caCertData, 0600)
	if err != nil {
		return nil, fmt.Errorf("read client tls failed: write ca cert file failed: %v", err)
	}
	return &vaultapi.TLSConfig{CACert: caCertFile}, nil
}
<|endoftext|>


<commit_before>// Package passphrase is a utility function for managing passphrase
// for TUF and Notary keys.
package passphrase

import (
	"bufio"
	"errors"
	"fmt"
	"io"
	"os"
	"strings"

	"path/filepath"

	"github.com/docker/docker/pkg/term"
)

// Retriever is a callback function that should retrieve a passphrase
// for a given named key. If it should be treated as new passphrase (e.g. with
// confirmation), createNew will be true. Attempts is passed in so that implementers
// decide how many chances to give to a human, for example.
type Retriever func(keyName, alias string, createNew bool, attempts int) (passphrase string, giveup bool, err error)

const (
	idBytesToDisplay            = 7
	tufRootAlias                = "root"
	tufTargetsAlias             = "targets"
	tufSnapshotAlias            = "snapshot"
	tufRootKeyGenerationWarning = `You are about to create a new root signing key passphrase. This passphrase
will be used to protect the most sensitive key in your signing system. Please
choose a long, complex passphrase and be careful to keep the password and the
key file itself secure and backed up. It is highly recommended that you use a
password manager to generate the passphrase and keep it safe. There will be no
way to recover this key. You can find the key in your config directory.`
)

var (
	// ErrTooShort is returned if the passphrase entered for a new key is
	// below the minimum length
	ErrTooShort = errors.New("Passphrase too short")

	// ErrDontMatch is returned if the two entered passphrases don't match.
	// new key is below the minimum length
	ErrDontMatch = errors.New("The entered passphrases do not match")

	// ErrTooManyAttempts is returned if the maximum number of passphrase
	// entry attempts is reached.
	ErrTooManyAttempts = errors.New("Too many attempts")
)

// PromptRetriever returns a new Retriever which will provide a prompt on stdin
// and stdout to retrieve a passphrase. The passphrase will be cached such that
// subsequent prompts will produce the same passphrase.
func PromptRetriever() Retriever {
	return PromptRetrieverWithInOut(os.Stdin, os.Stdout, nil)
}

// PromptRetrieverWithInOut returns a new Retriever which will provide a
// prompt using the given in and out readers. The passphrase will be cached
// such that subsequent prompts will produce the same passphrase.
// aliasMap can be used to specify display names for TUF key aliases. If aliasMap
// is nil, a sensible default will be used.
func PromptRetrieverWithInOut(in io.Reader, out io.Writer, aliasMap map[string]string) Retriever {
	userEnteredTargetsSnapshotsPass := false
	targetsSnapshotsPass := ""
	userEnteredRootsPass := false
	rootsPass := ""

	return func(keyName string, alias string, createNew bool, numAttempts int) (string, bool, error) {
		if alias == tufRootAlias && createNew && numAttempts == 0 {
			fmt.Fprintln(out, tufRootKeyGenerationWarning)
		}
		if numAttempts > 0 {
			if !createNew {
				fmt.Fprintln(out, "Passphrase incorrect. Please retry.")
			}
		}

		// Figure out if we should display a different string for this alias
		displayAlias := alias
		if aliasMap != nil {
			if val, ok := aliasMap[alias]; ok {
				displayAlias = val
			}

		}

		// First, check if we have a password cached for this alias.
		if numAttempts == 0 {
			if userEnteredTargetsSnapshotsPass && (alias == tufSnapshotAlias || alias == tufTargetsAlias) {
				return targetsSnapshotsPass, false, nil
			}
			if userEnteredRootsPass && (alias == "root") {
				return rootsPass, false, nil
			}
		}

		if numAttempts > 3 && !createNew {
			return "", true, ErrTooManyAttempts
		}

		state, err := term.SaveState(0)
		if err != nil {
			return "", false, err
		}
		term.DisableEcho(0, state)
		defer term.RestoreTerminal(0, state)

		stdin := bufio.NewReader(in)

		indexOfLastSeparator := strings.LastIndex(keyName, string(filepath.Separator))
		if indexOfLastSeparator == -1 {
			indexOfLastSeparator = 0
		}

		if len(keyName) > indexOfLastSeparator+idBytesToDisplay {
			if indexOfLastSeparator > 0 {
				keyNamePrefix := keyName[:indexOfLastSeparator]
				keyNameID := keyName[indexOfLastSeparator+1 : indexOfLastSeparator+idBytesToDisplay+1]
				keyName = keyNamePrefix + " (" + keyNameID + ")"
			} else {
				keyName = keyName[indexOfLastSeparator : indexOfLastSeparator+idBytesToDisplay]
			}
		}

		if createNew {
			fmt.Fprintf(out, "Enter passphrase for new %s key with id %s: ", displayAlias, keyName)
		} else {
			fmt.Fprintf(out, "Enter key passphrase for %s key with id %s: ", displayAlias, keyName)
		}

		passphrase, err := stdin.ReadBytes('\n')
		fmt.Fprintln(out)
		if err != nil {
			return "", false, err
		}

		retPass := strings.TrimSpace(string(passphrase))

		if !createNew {
			if alias == tufSnapshotAlias || alias == tufTargetsAlias {
				userEnteredTargetsSnapshotsPass = true
				targetsSnapshotsPass = retPass
			}
			if alias == tufRootAlias {
				userEnteredRootsPass = true
				rootsPass = retPass
			}
			return retPass, false, nil
		}

		if len(retPass) < 8 {
			fmt.Fprintln(out, "Please use a password manager to generate and store a good random passphrase.")
			return "", false, ErrTooShort
		}

		fmt.Fprintf(out, "Repeat passphrase for new %s key with id %s: ", displayAlias, keyName)
		confirmation, err := stdin.ReadBytes('\n')
		fmt.Fprintln(out)
		if err != nil {
			return "", false, err
		}
		confirmationStr := strings.TrimSpace(string(confirmation))

		if retPass != confirmationStr {
			fmt.Fprintln(out, "Passphrases do not match. Please retry.")
			return "", false, ErrDontMatch
		}

		if alias == tufSnapshotAlias || alias == tufTargetsAlias {
			userEnteredTargetsSnapshotsPass = true
			targetsSnapshotsPass = retPass
		}
		if alias == tufRootAlias {
			userEnteredRootsPass = true
			rootsPass = retPass
		}

		return retPass, false, nil
	}
}
<commit_msg>Update passphrase error message<commit_after>// Package passphrase is a utility function for managing passphrase
// for TUF and Notary keys.
package passphrase

import (
	"bufio"
	"errors"
	"fmt"
	"io"
	"os"
	"strings"

	"path/filepath"

	"github.com/docker/docker/pkg/term"
)

// Retriever is a callback function that should retrieve a passphrase
// for a given named key. If it should be treated as new passphrase (e.g. with
// confirmation), createNew will be true. Attempts is passed in so that implementers
// decide how many chances to give to a human, for example.
type Retriever func(keyName, alias string, createNew bool, attempts int) (passphrase string, giveup bool, err error)

const (
	idBytesToDisplay            = 7
	tufRootAlias                = "root"
	tufTargetsAlias             = "targets"
	tufSnapshotAlias            = "snapshot"
	tufRootKeyGenerationWarning = `You are about to create a new root signing key passphrase. This passphrase
will be used to protect the most sensitive key in your signing system. Please
choose a long, complex passphrase and be careful to keep the password and the
key file itself secure and backed up. It is highly recommended that you use a
password manager to generate the passphrase and keep it safe. There will be no
way to recover this key. You can find the key in your config directory.`
)

var (
	// ErrTooShort is returned if the passphrase entered for a new key is
	// below the minimum length
	ErrTooShort = errors.New("Passphrase too short")

	// ErrDontMatch is returned if the two entered passphrases don't match.
	// new key is below the minimum length
	ErrDontMatch = errors.New("The entered passphrases do not match")

	// ErrTooManyAttempts is returned if the maximum number of passphrase
	// entry attempts is reached.
	ErrTooManyAttempts = errors.New("Too many attempts")
)

// PromptRetriever returns a new Retriever which will provide a prompt on stdin
// and stdout to retrieve a passphrase. The passphrase will be cached such that
// subsequent prompts will produce the same passphrase.
func PromptRetriever() Retriever {
	return PromptRetrieverWithInOut(os.Stdin, os.Stdout, nil)
}

// PromptRetrieverWithInOut returns a new Retriever which will provide a
// prompt using the given in and out readers. The passphrase will be cached
// such that subsequent prompts will produce the same passphrase.
// aliasMap can be used to specify display names for TUF key aliases. If aliasMap
// is nil, a sensible default will be used.
func PromptRetrieverWithInOut(in io.Reader, out io.Writer, aliasMap map[string]string) Retriever {
	userEnteredTargetsSnapshotsPass := false
	targetsSnapshotsPass := ""
	userEnteredRootsPass := false
	rootsPass := ""

	return func(keyName string, alias string, createNew bool, numAttempts int) (string, bool, error) {
		if alias == tufRootAlias && createNew && numAttempts == 0 {
			fmt.Fprintln(out, tufRootKeyGenerationWarning)
		}
		if numAttempts > 0 {
			if !createNew {
				fmt.Fprintln(out, "Passphrase incorrect. Please retry.")
			}
		}

		// Figure out if we should display a different string for this alias
		displayAlias := alias
		if aliasMap != nil {
			if val, ok := aliasMap[alias]; ok {
				displayAlias = val
			}

		}

		// First, check if we have a password cached for this alias.
		if numAttempts == 0 {
			if userEnteredTargetsSnapshotsPass && (alias == tufSnapshotAlias || alias == tufTargetsAlias) {
				return targetsSnapshotsPass, false, nil
			}
			if userEnteredRootsPass && (alias == "root") {
				return rootsPass, false, nil
			}
		}

		if numAttempts > 3 && !createNew {
			return "", true, ErrTooManyAttempts
		}

		state, err := term.SaveState(0)
		if err != nil {
			return "", false, err
		}
		term.DisableEcho(0, state)
		defer term.RestoreTerminal(0, state)

		stdin := bufio.NewReader(in)

		indexOfLastSeparator := strings.LastIndex(keyName, string(filepath.Separator))
		if indexOfLastSeparator == -1 {
			indexOfLastSeparator = 0
		}

		if len(keyName) > indexOfLastSeparator+idBytesToDisplay {
			if indexOfLastSeparator > 0 {
				keyNamePrefix := keyName[:indexOfLastSeparator]
				keyNameID := keyName[indexOfLastSeparator+1 : indexOfLastSeparator+idBytesToDisplay+1]
				keyName = keyNamePrefix + " (" + keyNameID + ")"
			} else {
				keyName = keyName[indexOfLastSeparator : indexOfLastSeparator+idBytesToDisplay]
			}
		}

		if createNew {
			fmt.Fprintf(out, "Enter passphrase for new %s key with id %s: ", displayAlias, keyName)
		} else {
			fmt.Fprintf(out, "Enter key passphrase for %s key with id %s: ", displayAlias, keyName)
		}

		passphrase, err := stdin.ReadBytes('\n')
		fmt.Fprintln(out)
		if err != nil {
			return "", false, err
		}

		retPass := strings.TrimSpace(string(passphrase))

		if !createNew {
			if alias == tufSnapshotAlias || alias == tufTargetsAlias {
				userEnteredTargetsSnapshotsPass = true
				targetsSnapshotsPass = retPass
			}
			if alias == tufRootAlias {
				userEnteredRootsPass = true
				rootsPass = retPass
			}
			return retPass, false, nil
		}

		if len(retPass) < 8 {
			fmt.Fprintln(out, "Passphrase is too short. Please use a password manager to generate and store a good random passphrase.")
			return "", false, ErrTooShort
		}

		fmt.Fprintf(out, "Repeat passphrase for new %s key with id %s: ", displayAlias, keyName)
		confirmation, err := stdin.ReadBytes('\n')
		fmt.Fprintln(out)
		if err != nil {
			return "", false, err
		}
		confirmationStr := strings.TrimSpace(string(confirmation))

		if retPass != confirmationStr {
			fmt.Fprintln(out, "Passphrases do not match. Please retry.")
			return "", false, ErrDontMatch
		}

		if alias == tufSnapshotAlias || alias == tufTargetsAlias {
			userEnteredTargetsSnapshotsPass = true
			targetsSnapshotsPass = retPass
		}
		if alias == tufRootAlias {
			userEnteredRootsPass = true
			rootsPass = retPass
		}

		return retPass, false, nil
	}
}
<|endoftext|>


<commit_before>package balancer

import (
	"errors"
	"math/rand"
)

type (
	// WeightBalancer balancer
	WeightBalancer struct{}
)

var (
	// ErrZeroWeight is used when there a zero value weight was given
	ErrZeroWeight = errors.New("invalid backend, weight 0 given")
)

// NewWeightBalancer creates a new instance of WeightBalancer
func NewWeightBalancer() *WeightBalancer {
	return &WeightBalancer{}
}

// Elect backend using weight strategy
func (b *WeightBalancer) Elect(hosts []*Target) (*Target, error) {
	if len(hosts) == 0 {
		return nil, ErrEmptyBackendList
	}

	totalWeight := 0
	for _, host := range hosts {
		totalWeight += host.Weight
	}

	if totalWeight <= 0 {
		return nil, ErrZeroWeight
	}

	r := rand.Intn(totalWeight)
	pos := 0

	for _, host := range hosts {
		pos += host.Weight
		if r >= pos {
			continue
		}
		return host, nil
	}

	return nil, ErrCannotElectBackend
}
<commit_msg>Simply return first host if there is only one possible upstream host<commit_after>package balancer

import (
	"errors"
	"math/rand"
)

type (
	// WeightBalancer balancer
	WeightBalancer struct{}
)

var (
	// ErrZeroWeight is used when there a zero value weight was given
	ErrZeroWeight = errors.New("invalid backend, weight 0 given")
)

// NewWeightBalancer creates a new instance of WeightBalancer
func NewWeightBalancer() *WeightBalancer {
	return &WeightBalancer{}
}

// Elect backend using weight strategy
func (b *WeightBalancer) Elect(hosts []*Target) (*Target, error) {
	if len(hosts) == 0 {
		return nil, ErrEmptyBackendList
	}

	if len(hosts) == 1 {
		return hosts[0], nil
	}

	totalWeight := 0
	for _, host := range hosts {
		totalWeight += host.Weight
	}

	if totalWeight <= 0 {
		return nil, ErrZeroWeight
	}

	r := rand.Intn(totalWeight)
	pos := 0

	for _, host := range hosts {
		pos += host.Weight
		if r >= pos {
			continue
		}
		return host, nil
	}

	return nil, ErrCannotElectBackend
}
<|endoftext|>