<commit_before>/*
Copyright 2019 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 objectmeta

import (
	"testing"

	structuralschema "k8s.io/apiextensions-apiserver/pkg/apiserver/schema"
	"k8s.io/apimachinery/pkg/util/json"
	"k8s.io/apimachinery/pkg/util/validation/field"
)

func TestValidateEmbeddedResource(t *testing.T) {
	tests := []struct {
		name   string
		object map[string]interface{}
		errors []validationMatch
	}{
		{name: "empty", object: map[string]interface{}{}, errors: []validationMatch{
			required("apiVersion"),
			required("kind"),
		}},
		{name: "version and kind", object: map[string]interface{}{
			"apiVersion": "foo/v1",
			"kind":       "Foo",
		}},
		{name: "invalid kind", object: map[string]interface{}{
			"apiVersion": "foo/v1",
			"kind":       "foo.bar-com",
		}, errors: []validationMatch{
			invalid("kind"),
		}},
		{name: "no name", object: map[string]interface{}{
			"apiVersion": "foo/v1",
			"kind":       "Foo",
			"metadata": map[string]interface{}{
				"namespace": "kube-system",
			},
		}},
		{name: "no namespace", object: map[string]interface{}{
			"apiVersion": "foo/v1",
			"kind":       "Foo",
			"metadata": map[string]interface{}{
				"name": "foo",
			},
		}},
		{name: "invalid", object: map[string]interface{}{
			"apiVersion": "foo/v1",
			"kind":       "Foo",
			"metadata": map[string]interface{}{
				"name":      "..",
				"namespace": "$$$",
				"labels": map[string]interface{}{
					"#": "#",
				},
				"annotations": map[string]interface{}{
					"#": "#",
				},
			},
		}, errors: []validationMatch{
			invalid("metadata", "name"),
			invalid("metadata", "namespace"),
			invalid("metadata", "labels"),      // key
			invalid("metadata", "labels"),      // value
			invalid("metadata", "annotations"), // key
		}},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			schema := &structuralschema.Structural{Extensions: structuralschema.Extensions{XEmbeddedResource: true}}
			errs := validateEmbeddedResource(nil, tt.object, schema)
			seenErrs := make([]bool, len(errs))

			for _, expectedError := range tt.errors {
				found := false
				for i, err := range errs {
					if expectedError.matches(err) && !seenErrs[i] {
						found = true
						seenErrs[i] = true
						break
					}
				}

				if !found {
					t.Errorf("expected %v at %v, got %v", expectedError.errorType, expectedError.path.String(), errs)
				}
			}

			for i, seen := range seenErrs {
				if !seen {
					t.Errorf("unexpected error: %v", errs[i])
				}
			}
		})
	}
}

func TestValidate(t *testing.T) {
	tests := []struct {
		name        string
		object      string
		includeRoot bool
		errors      []validationMatch
	}{
		{name: "empty", object: `{}`, errors: []validationMatch{}},
		{name: "include root", object: `{}`, includeRoot: true, errors: []validationMatch{
			required("apiVersion"),
			required("kind"),
		}},
		{name: "embedded", object: `
{
  "embedded": {}
}`, errors: []validationMatch{
			required("embedded", "apiVersion"),
			required("embedded", "kind"),
		}},
		{name: "nested", object: `
{
  "nested": {
    "embedded": {}
  }
}`, errors: []validationMatch{
			required("nested", "apiVersion"),
			required("nested", "kind"),
			required("nested", "embedded", "apiVersion"),
			required("nested", "embedded", "kind"),
		}},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			schema := &structuralschema.Structural{
				Properties: map[string]structuralschema.Structural{
					"embedded": {Extensions: structuralschema.Extensions{XEmbeddedResource: true}},
					"nested": {
						Extensions: structuralschema.Extensions{XEmbeddedResource: true},
						Properties: map[string]structuralschema.Structural{
							"embedded": {Extensions: structuralschema.Extensions{XEmbeddedResource: true}},
						},
					},
				},
			}

			var obj map[string]interface{}
			if err := json.Unmarshal([]byte(tt.object), &obj); err != nil {
				t.Fatal(err)
			}

			errs := Validate(nil, obj, schema, tt.includeRoot)
			seenErrs := make([]bool, len(errs))

			for _, expectedError := range tt.errors {
				found := false
				for i, err := range errs {
					if expectedError.matches(err) && !seenErrs[i] {
						found = true
						seenErrs[i] = true
						break
					}
				}

				if !found {
					t.Errorf("expected %v at %v, got %v", expectedError.errorType, expectedError.path.String(), errs)
				}
			}

			for i, seen := range seenErrs {
				if !seen {
					t.Errorf("unexpected error: %v", errs[i])
				}
			}
		})
	}
}

type validationMatch struct {
	path      *field.Path
	errorType field.ErrorType
}

func required(path ...string) validationMatch {
	return validationMatch{path: field.NewPath(path[0], path[1:]...), errorType: field.ErrorTypeRequired}
}
func invalid(path ...string) validationMatch {
	return validationMatch{path: field.NewPath(path[0], path[1:]...), errorType: field.ErrorTypeInvalid}
}
func invalidIndex(index int, path ...string) validationMatch {
	return validationMatch{path: field.NewPath(path[0], path[1:]...).Index(index), errorType: field.ErrorTypeInvalid}
}
func unsupported(path ...string) validationMatch {
	return validationMatch{path: field.NewPath(path[0], path[1:]...), errorType: field.ErrorTypeNotSupported}
}
func immutable(path ...string) validationMatch {
	return validationMatch{path: field.NewPath(path[0], path[1:]...), errorType: field.ErrorTypeInvalid}
}
func forbidden(path ...string) validationMatch {
	return validationMatch{path: field.NewPath(path[0], path[1:]...), errorType: field.ErrorTypeForbidden}
}

func (v validationMatch) matches(err *field.Error) bool {
	return err.Type == v.errorType && err.Field == v.path.String()
}
<commit_msg>apiextensions: add items+additionalProperties test coverage for objectmeta.Validate<commit_after>/*
Copyright 2019 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 objectmeta

import (
	"testing"

	structuralschema "k8s.io/apiextensions-apiserver/pkg/apiserver/schema"
	"k8s.io/apimachinery/pkg/util/json"
	"k8s.io/apimachinery/pkg/util/validation/field"
)

func TestValidateEmbeddedResource(t *testing.T) {
	tests := []struct {
		name   string
		object map[string]interface{}
		errors []validationMatch
	}{
		{name: "empty", object: map[string]interface{}{}, errors: []validationMatch{
			required("apiVersion"),
			required("kind"),
		}},
		{name: "version and kind", object: map[string]interface{}{
			"apiVersion": "foo/v1",
			"kind":       "Foo",
		}},
		{name: "invalid kind", object: map[string]interface{}{
			"apiVersion": "foo/v1",
			"kind":       "foo.bar-com",
		}, errors: []validationMatch{
			invalid("kind"),
		}},
		{name: "no name", object: map[string]interface{}{
			"apiVersion": "foo/v1",
			"kind":       "Foo",
			"metadata": map[string]interface{}{
				"namespace": "kube-system",
			},
		}},
		{name: "no namespace", object: map[string]interface{}{
			"apiVersion": "foo/v1",
			"kind":       "Foo",
			"metadata": map[string]interface{}{
				"name": "foo",
			},
		}},
		{name: "invalid", object: map[string]interface{}{
			"apiVersion": "foo/v1",
			"kind":       "Foo",
			"metadata": map[string]interface{}{
				"name":      "..",
				"namespace": "$$$",
				"labels": map[string]interface{}{
					"#": "#",
				},
				"annotations": map[string]interface{}{
					"#": "#",
				},
			},
		}, errors: []validationMatch{
			invalid("metadata", "name"),
			invalid("metadata", "namespace"),
			invalid("metadata", "labels"),      // key
			invalid("metadata", "labels"),      // value
			invalid("metadata", "annotations"), // key
		}},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			schema := &structuralschema.Structural{Extensions: structuralschema.Extensions{XEmbeddedResource: true}}
			errs := validateEmbeddedResource(nil, tt.object, schema)
			seenErrs := make([]bool, len(errs))

			for _, expectedError := range tt.errors {
				found := false
				for i, err := range errs {
					if expectedError.matches(err) && !seenErrs[i] {
						found = true
						seenErrs[i] = true
						break
					}
				}

				if !found {
					t.Errorf("expected %v at %v, got %v", expectedError.errorType, expectedError.path.String(), errs)
				}
			}

			for i, seen := range seenErrs {
				if !seen {
					t.Errorf("unexpected error: %v", errs[i])
				}
			}
		})
	}
}

func TestValidate(t *testing.T) {
	tests := []struct {
		name        string
		object      string
		includeRoot bool
		errors      []validationMatch
	}{
		{name: "empty", object: `{}`, errors: []validationMatch{}},
		{name: "include root", object: `{}`, includeRoot: true, errors: []validationMatch{
			required("apiVersion"),
			required("kind"),
		}},
		{name: "embedded", object: `
{
  "embedded": {}
}`, errors: []validationMatch{
			required("embedded", "apiVersion"),
			required("embedded", "kind"),
		}},
		{name: "nested", object: `
{
  "nested": {
    "embedded": {}
  }
}`, errors: []validationMatch{
			required("nested", "apiVersion"),
			required("nested", "kind"),
			required("nested", "embedded", "apiVersion"),
			required("nested", "embedded", "kind"),
		}},
		{name: "items", object: `
{
  "items": [{}]
}`, errors: []validationMatch{
			required("items[0]", "apiVersion"),
			required("items[0]", "kind"),
		}},
		{name: "additionalProperties", object: `
{
  "additionalProperties": {"foo":{}}
}`, errors: []validationMatch{
			required("additionalProperties[foo]", "apiVersion"),
			required("additionalProperties[foo]", "kind"),
		}},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			schema := &structuralschema.Structural{
				Properties: map[string]structuralschema.Structural{
					"embedded": {Extensions: structuralschema.Extensions{XEmbeddedResource: true}},
					"nested": {
						Extensions: structuralschema.Extensions{XEmbeddedResource: true},
						Properties: map[string]structuralschema.Structural{
							"embedded": {Extensions: structuralschema.Extensions{XEmbeddedResource: true}},
						},
					},
					"items": {
						Items: &structuralschema.Structural{
							Extensions: structuralschema.Extensions{XEmbeddedResource: true},
						},
					},
					"additionalProperties": {
						Generic: structuralschema.Generic{
							AdditionalProperties: &structuralschema.StructuralOrBool{
								Structural: &structuralschema.Structural{
									Extensions: structuralschema.Extensions{XEmbeddedResource: true},
								},
							},
						},
					},
				},
			}

			var obj map[string]interface{}
			if err := json.Unmarshal([]byte(tt.object), &obj); err != nil {
				t.Fatal(err)
			}

			errs := Validate(nil, obj, schema, tt.includeRoot)
			seenErrs := make([]bool, len(errs))

			for _, expectedError := range tt.errors {
				found := false
				for i, err := range errs {
					if expectedError.matches(err) && !seenErrs[i] {
						found = true
						seenErrs[i] = true
						break
					}
				}

				if !found {
					t.Errorf("expected %v at %v, got %v", expectedError.errorType, expectedError.path.String(), errs)
				}
			}

			for i, seen := range seenErrs {
				if !seen {
					t.Errorf("unexpected error: %v", errs[i])
				}
			}
		})
	}
}

type validationMatch struct {
	path      *field.Path
	errorType field.ErrorType
}

func required(path ...string) validationMatch {
	return validationMatch{path: field.NewPath(path[0], path[1:]...), errorType: field.ErrorTypeRequired}
}
func invalid(path ...string) validationMatch {
	return validationMatch{path: field.NewPath(path[0], path[1:]...), errorType: field.ErrorTypeInvalid}
}
func invalidIndex(index int, path ...string) validationMatch {
	return validationMatch{path: field.NewPath(path[0], path[1:]...).Index(index), errorType: field.ErrorTypeInvalid}
}
func unsupported(path ...string) validationMatch {
	return validationMatch{path: field.NewPath(path[0], path[1:]...), errorType: field.ErrorTypeNotSupported}
}
func immutable(path ...string) validationMatch {
	return validationMatch{path: field.NewPath(path[0], path[1:]...), errorType: field.ErrorTypeInvalid}
}
func forbidden(path ...string) validationMatch {
	return validationMatch{path: field.NewPath(path[0], path[1:]...), errorType: field.ErrorTypeForbidden}
}

func (v validationMatch) matches(err *field.Error) bool {
	return err.Type == v.errorType && err.Field == v.path.String()
}
<|endoftext|>


<commit_before>package controller

import (
	"fmt"

	"github.com/golang/glog"

	apierrs "k8s.io/apimachinery/pkg/api/errors"
	utilruntime "k8s.io/apimachinery/pkg/util/runtime"
	"k8s.io/client-go/tools/cache"
	"k8s.io/client-go/util/flowcontrol"

	imageapi "github.com/openshift/origin/pkg/image/apis/image"
	imageclient "github.com/openshift/origin/pkg/image/generated/internalclientset/typed/image/internalversion"
	imageinformer "github.com/openshift/origin/pkg/image/generated/listers/image/internalversion"
)

type uniqueItem struct {
	uid             string
	resourceVersion string
}

type ScheduledImageStreamController struct {
	// boolean flag whether this controller is active
	enabled bool

	// image stream client
	client imageclient.ImageInterface

	// lister can list/get image streams from a shared informer's cache
	lister imageinformer.ImageStreamLister
	// listerSynced makes sure the is store is synced before reconciling streams
	listerSynced cache.InformerSynced

	// rateLimiter to be used when re-importing images
	rateLimiter flowcontrol.RateLimiter

	// scheduler for timely image re-imports
	scheduler *scheduler
}

// Importing is invoked when the controller decides to import a stream in order to push back
// the next schedule time.
func (s *ScheduledImageStreamController) Importing(stream *imageapi.ImageStream) {
	if !s.enabled {
		return
	}
	glog.V(5).Infof("DEBUG: stream %s was just imported", stream.Name)
	// Push the current key back to the end of the queue because it's just been imported
	key, _ := cache.MetaNamespaceKeyFunc(stream)
	s.scheduler.Delay(key)
}

// Run begins watching and syncing.
func (s *ScheduledImageStreamController) Run(stopCh <-chan struct{}) {
	defer utilruntime.HandleCrash()

	glog.Infof("Starting scheduled import controller")

	// Wait for the stream store to sync before starting any work in this controller.
	if !cache.WaitForCacheSync(stopCh, s.listerSynced) {
		utilruntime.HandleError(fmt.Errorf("timed out waiting for caches to sync"))
		return
	}

	go s.scheduler.RunUntil(stopCh)

	<-stopCh
	glog.Infof("Shutting down image stream controller")
}

func (s *ScheduledImageStreamController) addImageStream(obj interface{}) {
	stream := obj.(*imageapi.ImageStream)
	s.enqueueImageStream(stream)
}

func (s *ScheduledImageStreamController) updateImageStream(old, cur interface{}) {
	curStream, ok := cur.(*imageapi.ImageStream)
	if !ok {
		return
	}
	oldStream, ok := old.(*imageapi.ImageStream)
	if !ok {
		return
	}
	// we only compare resource version, since deeper inspection if a stream
	// needs to be re-imported happens in syncImageStream
	if curStream.ResourceVersion == oldStream.ResourceVersion {
		return
	}
	s.enqueueImageStream(curStream)
}

func (s *ScheduledImageStreamController) deleteImageStream(obj interface{}) {
	stream, isStream := obj.(*imageapi.ImageStream)
	if !isStream {
		tombstone, objIsTombstone := obj.(cache.DeletedFinalStateUnknown)
		if !objIsTombstone {
			return
		}
		var isImageStream bool
		stream, isImageStream = tombstone.Obj.(*imageapi.ImageStream)
		if !isImageStream {
			glog.V(2).Infof("tombstone contained unexpected object %#v", tombstone)
			return
		}
	}
	key, err := cache.MetaNamespaceKeyFunc(stream)
	if err != nil {
		glog.V(2).Infof("unable to get namespace key function for stream %s/%s: %v", stream.Namespace, stream.Name, err)
		return
	}
	s.scheduler.Remove(key, nil)
}

// enqueueImageStream ensures an image stream is checked for scheduling
func (s *ScheduledImageStreamController) enqueueImageStream(stream *imageapi.ImageStream) {
	if !s.enabled {
		return
	}
	if needsScheduling(stream) {
		key, err := cache.MetaNamespaceKeyFunc(stream)
		if err != nil {
			glog.V(2).Infof("unable to get namespace key function for stream %s/%s: %v", stream.Namespace, stream.Name, err)
			return
		}
		s.scheduler.Add(key, uniqueItem{uid: string(stream.UID), resourceVersion: stream.ResourceVersion})
	}
}

// syncTimed is invoked when a key is ready to be processed.
func (s *ScheduledImageStreamController) syncTimed(key, value interface{}) {
	if !s.enabled {
		s.scheduler.Remove(key, value)
		return
	}
	if s.rateLimiter != nil && !s.rateLimiter.TryAccept() {
		glog.V(5).Infof("DEBUG: check of %s exceeded rate limit, will retry later", key)
		return
	}
	namespace, name, err := cache.SplitMetaNamespaceKey(key.(string))
	if err != nil {
		glog.V(2).Infof("unable to split namespace key for key %q: %v", key, err)
		return
	}
	if err := s.syncTimedByName(namespace, name); err != nil {
		// the stream cannot be imported
		if err == ErrNotImportable {
			// value must match to be removed, so we avoid races against creation by ensuring that we only
			// remove the stream if the uid and resource version in the scheduler are exactly the same.
			s.scheduler.Remove(key, value)
			return
		}
		utilruntime.HandleError(err)
		return
	}
}

func (s *ScheduledImageStreamController) syncTimedByName(namespace, name string) error {
	sharedStream, err := s.lister.ImageStreams(namespace).Get(name)
	if err != nil {
		if apierrs.IsNotFound(err) {
			return ErrNotImportable
		}
		return err
	}
	if !needsScheduling(sharedStream) {
		return ErrNotImportable
	}

	stream := sharedStream.DeepCopy()
	resetScheduledTags(stream)

	glog.V(3).Infof("Scheduled import of stream %s/%s...", stream.Namespace, stream.Name)
	return handleImageStream(stream, s.client, nil)
}

// resetScheduledTags artificially increments the generation on the tags that should be imported.
func resetScheduledTags(stream *imageapi.ImageStream) {
	next := stream.Generation + 1
	for tag, tagRef := range stream.Spec.Tags {
		if tagImportable(tagRef) && tagRef.ImportPolicy.Scheduled {
			tagRef.Generation = &next
			stream.Spec.Tags[tag] = tagRef
		}
	}
}

// needsScheduling returns true if this image stream has any scheduled tags
func needsScheduling(stream *imageapi.ImageStream) bool {
	for _, tagRef := range stream.Spec.Tags {
		if tagImportable(tagRef) && tagRef.ImportPolicy.Scheduled {
			return true
		}
	}
	return false
}
<commit_msg>image: use deletion handling metanamespace key in import scheduler<commit_after>package controller

import (
	"fmt"

	"github.com/golang/glog"

	apierrs "k8s.io/apimachinery/pkg/api/errors"
	utilruntime "k8s.io/apimachinery/pkg/util/runtime"
	"k8s.io/client-go/tools/cache"
	"k8s.io/client-go/util/flowcontrol"

	imageapi "github.com/openshift/origin/pkg/image/apis/image"
	imageclient "github.com/openshift/origin/pkg/image/generated/internalclientset/typed/image/internalversion"
	imageinformer "github.com/openshift/origin/pkg/image/generated/listers/image/internalversion"
)

type uniqueItem struct {
	uid             string
	resourceVersion string
}

type ScheduledImageStreamController struct {
	// boolean flag whether this controller is active
	enabled bool

	// image stream client
	client imageclient.ImageInterface

	// lister can list/get image streams from a shared informer's cache
	lister imageinformer.ImageStreamLister
	// listerSynced makes sure the is store is synced before reconciling streams
	listerSynced cache.InformerSynced

	// rateLimiter to be used when re-importing images
	rateLimiter flowcontrol.RateLimiter

	// scheduler for timely image re-imports
	scheduler *scheduler
}

// Importing is invoked when the controller decides to import a stream in order to push back
// the next schedule time.
func (s *ScheduledImageStreamController) Importing(stream *imageapi.ImageStream) {
	if !s.enabled {
		return
	}
	glog.V(5).Infof("DEBUG: stream %s was just imported", stream.Name)
	// Push the current key back to the end of the queue because it's just been imported
	key, err := cache.DeletionHandlingMetaNamespaceKeyFunc(stream)
	if err != nil {
		utilruntime.HandleError(fmt.Errorf("failed to get the key for stream %s: %v", stream.Name, err))
		return
	}
	s.scheduler.Delay(key)
}

// Run begins watching and syncing.
func (s *ScheduledImageStreamController) Run(stopCh <-chan struct{}) {
	defer utilruntime.HandleCrash()

	glog.Infof("Starting scheduled import controller")

	// Wait for the stream store to sync before starting any work in this controller.
	if !cache.WaitForCacheSync(stopCh, s.listerSynced) {
		utilruntime.HandleError(fmt.Errorf("timed out waiting for caches to sync"))
		return
	}

	go s.scheduler.RunUntil(stopCh)

	<-stopCh
	glog.Infof("Shutting down image stream controller")
}

func (s *ScheduledImageStreamController) addImageStream(obj interface{}) {
	stream := obj.(*imageapi.ImageStream)
	s.enqueueImageStream(stream)
}

func (s *ScheduledImageStreamController) updateImageStream(old, cur interface{}) {
	curStream, ok := cur.(*imageapi.ImageStream)
	if !ok {
		return
	}
	oldStream, ok := old.(*imageapi.ImageStream)
	if !ok {
		return
	}
	// we only compare resource version, since deeper inspection if a stream
	// needs to be re-imported happens in syncImageStream
	if curStream.ResourceVersion == oldStream.ResourceVersion {
		return
	}
	s.enqueueImageStream(curStream)
}

func (s *ScheduledImageStreamController) deleteImageStream(obj interface{}) {
	key, err := cache.DeletionHandlingMetaNamespaceKeyFunc(obj)
	if err != nil {
		utilruntime.HandleError(fmt.Errorf("unable to get namespace key for %#v", obj))
		return
	}
	s.scheduler.Remove(key, nil)
}

// enqueueImageStream ensures an image stream is checked for scheduling
func (s *ScheduledImageStreamController) enqueueImageStream(stream *imageapi.ImageStream) {
	if !s.enabled {
		return
	}
	if needsScheduling(stream) {
		key, err := cache.DeletionHandlingMetaNamespaceKeyFunc(stream)
		if err != nil {
			glog.V(2).Infof("unable to get namespace key function for stream %s/%s: %v", stream.Namespace, stream.Name, err)
			return
		}
		s.scheduler.Add(key, uniqueItem{uid: string(stream.UID), resourceVersion: stream.ResourceVersion})
	}
}

// syncTimed is invoked when a key is ready to be processed.
func (s *ScheduledImageStreamController) syncTimed(key, value interface{}) {
	if !s.enabled {
		s.scheduler.Remove(key, value)
		return
	}
	if s.rateLimiter != nil && !s.rateLimiter.TryAccept() {
		glog.V(5).Infof("DEBUG: check of %s exceeded rate limit, will retry later", key)
		return
	}
	namespace, name, err := cache.SplitMetaNamespaceKey(key.(string))
	if err != nil {
		glog.V(2).Infof("unable to split namespace key for key %q: %v", key, err)
		return
	}
	if err := s.syncTimedByName(namespace, name); err != nil {
		// the stream cannot be imported
		if err == ErrNotImportable {
			// value must match to be removed, so we avoid races against creation by ensuring that we only
			// remove the stream if the uid and resource version in the scheduler are exactly the same.
			s.scheduler.Remove(key, value)
			return
		}
		utilruntime.HandleError(err)
		return
	}
}

func (s *ScheduledImageStreamController) syncTimedByName(namespace, name string) error {
	sharedStream, err := s.lister.ImageStreams(namespace).Get(name)
	if err != nil {
		if apierrs.IsNotFound(err) {
			return ErrNotImportable
		}
		return err
	}
	if !needsScheduling(sharedStream) {
		return ErrNotImportable
	}

	stream := sharedStream.DeepCopy()
	resetScheduledTags(stream)

	glog.V(3).Infof("Scheduled import of stream %s/%s...", stream.Namespace, stream.Name)
	return handleImageStream(stream, s.client, nil)
}

// resetScheduledTags artificially increments the generation on the tags that should be imported.
func resetScheduledTags(stream *imageapi.ImageStream) {
	next := stream.Generation + 1
	for tag, tagRef := range stream.Spec.Tags {
		if tagImportable(tagRef) && tagRef.ImportPolicy.Scheduled {
			tagRef.Generation = &next
			stream.Spec.Tags[tag] = tagRef
		}
	}
}

// needsScheduling returns true if this image stream has any scheduled tags
func needsScheduling(stream *imageapi.ImageStream) bool {
	for _, tagRef := range stream.Spec.Tags {
		if tagImportable(tagRef) && tagRef.ImportPolicy.Scheduled {
			return true
		}
	}
	return false
}
<|endoftext|>


<commit_before>package registry

import (
	"context"
	"testing"

	"github.com/stretchr/testify/assert"
)

func TestIntegrationBuildings(t *testing.T) {
	t.Run("get nonexistent building", func(t *testing.T) {
		assert := assert.New(t)
		building, err := testRegistrar.GetBuildingByID(context.Background(), "nonexistent")
		assert.NoError(err)
		assert.Nil(building)
	})

	var registeredBuilding *Building
	t.Run("register building", func(t *testing.T) {
		assert := assert.New(t)
		building := &Building{
			ID:   "something",
			Name: getULID().String(),
		}
		var err error
		registeredBuilding, err = testRegistrar.RegisterBuilding(context.Background(), building)
		assert.NoError(err)
		assert.NotEmpty(registeredBuilding.ID)
		assert.NotEqual(building.ID, registeredBuilding.ID, "registerbuliding should generate its own id")

	})

	t.Run("get existing building", func(t *testing.T) {
		assert := assert.New(t)
		building, err := testRegistrar.GetBuildingByID(context.Background(), registeredBuilding.ID)
		assert.NoError(err)
		assert.Equal(registeredBuilding, building)
	})

	t.Run("list buildings", func(t *testing.T) {
		assert := assert.New(t)
		buildings, err := testRegistrar.ListBuildings(context.Background())
		assert.NoError(err)
		assert.Len(buildings, 1)
		assert.Equal(registeredBuilding, buildings[0])
	})

	t.Run("deregister building", func(t *testing.T) {
		assert := assert.New(t)
		err := testRegistrar.DeregisterBuilding(context.Background(), registeredBuilding.ID)
		assert.NoError(err)
	})

	t.Run("get deregistered building", func(t *testing.T) {
		assert := assert.New(t)
		building, err := testRegistrar.GetBuildingByID(context.Background(), registeredBuilding.ID)
		assert.NoError(err)
		assert.Nil(building)
	})

	t.Run("deregister nonexistent building should not error", func(t *testing.T) {
		assert := assert.New(t)
		err := testRegistrar.DeregisterBuilding(context.Background(), "nonexistent")
		assert.NoError(err)
	})
}
<commit_msg>fix panic<commit_after>package registry

import (
	"context"
	"testing"

	"github.com/stretchr/testify/assert"
)

func TestIntegrationBuildings(t *testing.T) {
	t.Run("get nonexistent building", func(t *testing.T) {
		assert := assert.New(t)
		building, err := testRegistrar.GetBuildingByID(context.Background(), "nonexistent")
		assert.NoError(err)
		assert.Nil(building)
	})

	var registeredBuilding *Building
	t.Run("register building", func(t *testing.T) {
		assert := assert.New(t)
		building := &Building{
			ID:   "something",
			Name: getULID().String(),
		}
		var err error
		registeredBuilding, err = testRegistrar.RegisterBuilding(context.Background(), building)
		assert.NoError(err)
		assert.NotEmpty(registeredBuilding.ID)
		assert.NotEqual(building.ID, registeredBuilding.ID, "registerbuliding should generate its own id")

	})

	t.Run("get existing building", func(t *testing.T) {
		assert := assert.New(t)
		building, err := testRegistrar.GetBuildingByID(context.Background(), registeredBuilding.ID)
		assert.NoError(err)
		assert.Equal(registeredBuilding, building)
	})

	t.Run("list buildings", func(t *testing.T) {
		assert := assert.New(t)
		buildings, err := testRegistrar.ListBuildings(context.Background())
		if assert.NoError(err) {
			if assert.Len(buildings, 1) {
				assert.Equal(registeredBuilding, buildings[0])
			}
		}
	})

	t.Run("deregister building", func(t *testing.T) {
		assert := assert.New(t)
		err := testRegistrar.DeregisterBuilding(context.Background(), registeredBuilding.ID)
		assert.NoError(err)
	})

	t.Run("get deregistered building", func(t *testing.T) {
		assert := assert.New(t)
		building, err := testRegistrar.GetBuildingByID(context.Background(), registeredBuilding.ID)
		assert.NoError(err)
		assert.Nil(building)
	})

	t.Run("deregister nonexistent building should not error", func(t *testing.T) {
		assert := assert.New(t)
		err := testRegistrar.DeregisterBuilding(context.Background(), "nonexistent")
		assert.NoError(err)
	})
}
<|endoftext|>


<commit_before>/*
Copyright 2017 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 rest

import (
	"fmt"
	"time"

	"github.com/golang/glog"

	apierrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	utilruntime "k8s.io/apimachinery/pkg/util/runtime"
	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/apiserver/pkg/registry/generic"
	"k8s.io/apiserver/pkg/registry/rest"
	genericapiserver "k8s.io/apiserver/pkg/server"
	serverstorage "k8s.io/apiserver/pkg/server/storage"
	"k8s.io/kubernetes/pkg/api/legacyscheme"
	"k8s.io/kubernetes/pkg/apis/scheduling"
	schedulingapiv1alpha1 "k8s.io/kubernetes/pkg/apis/scheduling/v1alpha1"
	schedulingapiv1beta1 "k8s.io/kubernetes/pkg/apis/scheduling/v1beta1"
	schedulingclient "k8s.io/kubernetes/pkg/client/clientset_generated/internalclientset/typed/scheduling/internalversion"
	priorityclassstore "k8s.io/kubernetes/pkg/registry/scheduling/priorityclass/storage"
)

const PostStartHookName = "scheduling/bootstrap-system-priority-classes"

type RESTStorageProvider struct{}

var _ genericapiserver.PostStartHookProvider = RESTStorageProvider{}

func (p RESTStorageProvider) NewRESTStorage(apiResourceConfigSource serverstorage.APIResourceConfigSource, restOptionsGetter generic.RESTOptionsGetter) (genericapiserver.APIGroupInfo, bool) {
	apiGroupInfo := genericapiserver.NewDefaultAPIGroupInfo(scheduling.GroupName, legacyscheme.Scheme, legacyscheme.ParameterCodec, legacyscheme.Codecs)

	if apiResourceConfigSource.VersionEnabled(schedulingapiv1alpha1.SchemeGroupVersion) {
		apiGroupInfo.VersionedResourcesStorageMap[schedulingapiv1alpha1.SchemeGroupVersion.Version] = p.storage(apiResourceConfigSource, restOptionsGetter)
	}
	if apiResourceConfigSource.VersionEnabled(schedulingapiv1beta1.SchemeGroupVersion) {
		apiGroupInfo.VersionedResourcesStorageMap[schedulingapiv1beta1.SchemeGroupVersion.Version] = p.storage(apiResourceConfigSource, restOptionsGetter)
	}
	return apiGroupInfo, true
}

func (p RESTStorageProvider) storage(apiResourceConfigSource serverstorage.APIResourceConfigSource, restOptionsGetter generic.RESTOptionsGetter) map[string]rest.Storage {
	storage := map[string]rest.Storage{}
	// priorityclasses
	priorityClassStorage := priorityclassstore.NewREST(restOptionsGetter)
	storage["priorityclasses"] = priorityClassStorage

	return storage
}

func (p RESTStorageProvider) PostStartHook() (string, genericapiserver.PostStartHookFunc, error) {
	return PostStartHookName, AddSystemPriorityClasses(), nil
}

func AddSystemPriorityClasses() genericapiserver.PostStartHookFunc {
	return func(hookContext genericapiserver.PostStartHookContext) error {
		// Adding system priority classes is important. If they fail to add, many critical system
		// components may fail and cluster may break.
		err := wait.Poll(1*time.Second, 30*time.Second, func() (done bool, err error) {
			schedClientSet, err := schedulingclient.NewForConfig(hookContext.LoopbackClientConfig)
			if err != nil {
				utilruntime.HandleError(fmt.Errorf("unable to initialize client: %v", err))
				return false, nil
			}

			for _, pc := range scheduling.SystemPriorityClasses() {
				_, err := schedClientSet.PriorityClasses().Get(pc.Name, metav1.GetOptions{})
				if err != nil {
					if apierrors.IsNotFound(err) {
						_, err := schedClientSet.PriorityClasses().Create(pc)
						if err != nil {
							return false, err
						} else {
							glog.Infof("created PriorityClass %s with value %v", pc.Name, pc.Value)
						}
					} else {
						// Unable to get the priority class for reasons other than "not found".
						glog.Warningf("unable to get PriorityClass %v: %v. Retrying...", pc.Name, err)
						return false, err
					}
				}
			}
			glog.Infof("all system priority classes are created successfully or already exist.")
			return true, nil
		})
		// if we're never able to make it through initialization, kill the API server.
		if err != nil {
			return fmt.Errorf("unable to add default system priority classes: %v", err)
		}
		return nil
	}
}

func (p RESTStorageProvider) GroupName() string {
	return scheduling.GroupName
}
<commit_msg>fix apiserver crashed when priority classs already exists<commit_after>/*
Copyright 2017 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 rest

import (
	"fmt"
	"time"

	"github.com/golang/glog"

	apierrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	utilruntime "k8s.io/apimachinery/pkg/util/runtime"
	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/apiserver/pkg/registry/generic"
	"k8s.io/apiserver/pkg/registry/rest"
	genericapiserver "k8s.io/apiserver/pkg/server"
	serverstorage "k8s.io/apiserver/pkg/server/storage"
	"k8s.io/kubernetes/pkg/api/legacyscheme"
	"k8s.io/kubernetes/pkg/apis/scheduling"
	schedulingapiv1alpha1 "k8s.io/kubernetes/pkg/apis/scheduling/v1alpha1"
	schedulingapiv1beta1 "k8s.io/kubernetes/pkg/apis/scheduling/v1beta1"
	schedulingclient "k8s.io/kubernetes/pkg/client/clientset_generated/internalclientset/typed/scheduling/internalversion"
	priorityclassstore "k8s.io/kubernetes/pkg/registry/scheduling/priorityclass/storage"
)

const PostStartHookName = "scheduling/bootstrap-system-priority-classes"

type RESTStorageProvider struct{}

var _ genericapiserver.PostStartHookProvider = RESTStorageProvider{}

func (p RESTStorageProvider) NewRESTStorage(apiResourceConfigSource serverstorage.APIResourceConfigSource, restOptionsGetter generic.RESTOptionsGetter) (genericapiserver.APIGroupInfo, bool) {
	apiGroupInfo := genericapiserver.NewDefaultAPIGroupInfo(scheduling.GroupName, legacyscheme.Scheme, legacyscheme.ParameterCodec, legacyscheme.Codecs)

	if apiResourceConfigSource.VersionEnabled(schedulingapiv1alpha1.SchemeGroupVersion) {
		apiGroupInfo.VersionedResourcesStorageMap[schedulingapiv1alpha1.SchemeGroupVersion.Version] = p.storage(apiResourceConfigSource, restOptionsGetter)
	}
	if apiResourceConfigSource.VersionEnabled(schedulingapiv1beta1.SchemeGroupVersion) {
		apiGroupInfo.VersionedResourcesStorageMap[schedulingapiv1beta1.SchemeGroupVersion.Version] = p.storage(apiResourceConfigSource, restOptionsGetter)
	}
	return apiGroupInfo, true
}

func (p RESTStorageProvider) storage(apiResourceConfigSource serverstorage.APIResourceConfigSource, restOptionsGetter generic.RESTOptionsGetter) map[string]rest.Storage {
	storage := map[string]rest.Storage{}
	// priorityclasses
	priorityClassStorage := priorityclassstore.NewREST(restOptionsGetter)
	storage["priorityclasses"] = priorityClassStorage

	return storage
}

func (p RESTStorageProvider) PostStartHook() (string, genericapiserver.PostStartHookFunc, error) {
	return PostStartHookName, AddSystemPriorityClasses(), nil
}

func AddSystemPriorityClasses() genericapiserver.PostStartHookFunc {
	return func(hookContext genericapiserver.PostStartHookContext) error {
		// Adding system priority classes is important. If they fail to add, many critical system
		// components may fail and cluster may break.
		err := wait.Poll(1*time.Second, 30*time.Second, func() (done bool, err error) {
			schedClientSet, err := schedulingclient.NewForConfig(hookContext.LoopbackClientConfig)
			if err != nil {
				utilruntime.HandleError(fmt.Errorf("unable to initialize client: %v", err))
				return false, nil
			}

			for _, pc := range scheduling.SystemPriorityClasses() {
				_, err := schedClientSet.PriorityClasses().Get(pc.Name, metav1.GetOptions{})
				if err != nil {
					if apierrors.IsNotFound(err) {
						_, err := schedClientSet.PriorityClasses().Create(pc)
						if err != nil && !apierrors.IsAlreadyExists(err) {
							return false, err
						} else {
							glog.Infof("created PriorityClass %s with value %v", pc.Name, pc.Value)
						}
					} else {
						// Unable to get the priority class for reasons other than "not found".
						glog.Warningf("unable to get PriorityClass %v: %v. Retrying...", pc.Name, err)
						return false, err
					}
				}
			}
			glog.Infof("all system priority classes are created successfully or already exist.")
			return true, nil
		})
		// if we're never able to make it through initialization, kill the API server.
		if err != nil {
			return fmt.Errorf("unable to add default system priority classes: %v", err)
		}
		return nil
	}
}

func (p RESTStorageProvider) GroupName() string {
	return scheduling.GroupName
}
<|endoftext|>


<commit_before>package indexgateway

import (
	"context"
	"fmt"
	"net/http"
	"sync"
	"time"

	"github.com/go-kit/log"
	"github.com/go-kit/log/level"
	"github.com/grafana/dskit/kv"
	"github.com/grafana/dskit/ring"
	"github.com/grafana/dskit/services"
	"github.com/grafana/dskit/tenant"
	"github.com/pkg/errors"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/common/model"
	"github.com/prometheus/prometheus/model/labels"

	"github.com/grafana/loki/pkg/logproto"
	"github.com/grafana/loki/pkg/logql/syntax"
	"github.com/grafana/loki/pkg/storage/chunk"
	"github.com/grafana/loki/pkg/storage/chunk/fetcher"
	"github.com/grafana/loki/pkg/storage/stores/series/index"
	"github.com/grafana/loki/pkg/storage/stores/shipper/indexgateway/indexgatewaypb"
	"github.com/grafana/loki/pkg/storage/stores/shipper/util"
	util_log "github.com/grafana/loki/pkg/util/log"
)

const (
	maxIndexEntriesPerResponse     = 1000
	ringAutoForgetUnhealthyPeriods = 10
	ringNameForServer              = "index-gateway"
	ringNumTokens                  = 128
	ringCheckPeriod                = 3 * time.Second

	// RingIdentifier is used as a unique name to register the Index Gateway ring.
	RingIdentifier = "index-gateway"

	// RingKey is the name of the key used to register the different Index Gateway instances in the key-value store.
	RingKey = "index-gateway"
)

type IndexQuerier interface {
	GetChunkRefs(ctx context.Context, userID string, from, through model.Time, matchers ...*labels.Matcher) ([][]chunk.Chunk, []*fetcher.Fetcher, error)
	LabelValuesForMetricName(ctx context.Context, userID string, from, through model.Time, metricName string, labelName string, matchers ...*labels.Matcher) ([]string, error)
	LabelNamesForMetricName(ctx context.Context, userID string, from, through model.Time, metricName string) ([]string, error)
	Stop()
}

type IndexClient interface {
	QueryPages(ctx context.Context, queries []index.Query, callback index.QueryPagesCallback) error
	Stop()
}

type Gateway struct {
	services.Service

	indexQuerier IndexQuerier
	indexClient  IndexClient

	cfg Config
	log log.Logger

	shipper IndexQuerier

	subservices        *services.Manager
	subservicesWatcher *services.FailureWatcher

	ringLifecycler *ring.BasicLifecycler
	ring           *ring.Ring
}

// NewIndexGateway instantiates a new Index Gateway and start its services.
//
// In case it is configured to be in ring mode, a Basic Service wrapping the ring client is started.
// Otherwise, it starts an Idle Service that doesn't have lifecycle hooks.
func NewIndexGateway(cfg Config, log log.Logger, registerer prometheus.Registerer, indexQuerier IndexQuerier, indexClient IndexClient) (*Gateway, error) {
	g := &Gateway{
		indexQuerier: indexQuerier,
		indexClient:  indexClient,
		cfg:          cfg,
		log:          log,
	}

	if cfg.Mode == RingMode {
		ringStore, err := kv.NewClient(
			cfg.Ring.KVStore,
			ring.GetCodec(),
			kv.RegistererWithKVName(prometheus.WrapRegistererWithPrefix("loki_", registerer), "index-gateway"),
			log,
		)
		if err != nil {
			return nil, errors.Wrap(err, "create KV store client")
		}

		lifecyclerCfg, err := cfg.Ring.ToLifecyclerConfig(ringNumTokens, log)
		if err != nil {
			return nil, errors.Wrap(err, "invalid ring lifecycler config")
		}

		delegate := ring.BasicLifecyclerDelegate(g)
		delegate = ring.NewLeaveOnStoppingDelegate(delegate, log)
		delegate = ring.NewTokensPersistencyDelegate(cfg.Ring.TokensFilePath, ring.JOINING, delegate, log)
		delegate = ring.NewAutoForgetDelegate(ringAutoForgetUnhealthyPeriods*cfg.Ring.HeartbeatTimeout, delegate, log)

		g.ringLifecycler, err = ring.NewBasicLifecycler(lifecyclerCfg, ringNameForServer, RingKey, ringStore, delegate, log, registerer)
		if err != nil {
			return nil, errors.Wrap(err, "index gateway create ring lifecycler")
		}

		ringCfg := cfg.Ring.ToRingConfig(cfg.Ring.ReplicationFactor)
		g.ring, err = ring.NewWithStoreClientAndStrategy(ringCfg, ringNameForServer, RingKey, ringStore, ring.NewIgnoreUnhealthyInstancesReplicationStrategy(), prometheus.WrapRegistererWithPrefix("loki_", registerer), log)
		if err != nil {
			return nil, errors.Wrap(err, "index gateway create ring client")
		}

		svcs := []services.Service{g.ringLifecycler, g.ring}
		g.subservices, err = services.NewManager(svcs...)
		if err != nil {
			return nil, fmt.Errorf("new index gateway services manager: %w", err)
		}

		g.subservicesWatcher = services.NewFailureWatcher()
		g.subservicesWatcher.WatchManager(g.subservices)
		g.Service = services.NewBasicService(g.starting, g.running, g.stopping)
	} else {
		g.Service = services.NewIdleService(nil, func(failureCase error) error {
			g.indexQuerier.Stop()
			g.indexClient.Stop()
			return nil
		})
	}

	return g, nil
}

// starting implements the Lifecycler interface and is one of the lifecycle hooks.
//
// Only invoked if the Index Gateway is in ring mode.
func (g *Gateway) starting(ctx context.Context) (err error) {
	// In case this function will return error we want to unregister the instance
	// from the ring. We do it ensuring dependencies are gracefully stopped if they
	// were already started.
	defer func() {
		if err == nil || g.subservices == nil {
			return
		}

		if stopErr := services.StopManagerAndAwaitStopped(context.Background(), g.subservices); stopErr != nil {
			level.Error(util_log.Logger).Log("msg", "failed to gracefully stop index gateway dependencies", "err", stopErr)
		}
	}()

	if err := services.StartManagerAndAwaitHealthy(ctx, g.subservices); err != nil {
		return errors.Wrap(err, "unable to start index gateway subservices")
	}

	// The BasicLifecycler does not automatically move state to ACTIVE such that any additional work that
	// someone wants to do can be done before becoming ACTIVE. For the index gateway we don't currently
	// have any additional work so we can become ACTIVE right away.
	// Wait until the ring client detected this instance in the JOINING
	// state to make sure that when we'll run the initial sync we already
	// know the tokens assigned to this instance.
	level.Info(util_log.Logger).Log("msg", "waiting until index gateway is JOINING in the ring")
	if err := ring.WaitInstanceState(ctx, g.ring, g.ringLifecycler.GetInstanceID(), ring.JOINING); err != nil {
		return err
	}
	level.Info(util_log.Logger).Log("msg", "index gateway is JOINING in the ring")

	if err = g.ringLifecycler.ChangeState(ctx, ring.ACTIVE); err != nil {
		return errors.Wrapf(err, "switch instance to %s in the ring", ring.ACTIVE)
	}

	// Wait until the ring client detected this instance in the ACTIVE state to
	// make sure that when we'll run the loop it won't be detected as a ring
	// topology change.
	level.Info(util_log.Logger).Log("msg", "waiting until index gateway is ACTIVE in the ring")
	if err := ring.WaitInstanceState(ctx, g.ring, g.ringLifecycler.GetInstanceID(), ring.ACTIVE); err != nil {
		return err
	}
	level.Info(util_log.Logger).Log("msg", "index gateway is ACTIVE in the ring")

	return nil
}

// running implements the Lifecycler interface and is one of the lifecycle hooks.
//
// Only invoked if the Index Gateway is in ring mode.
func (g *Gateway) running(ctx context.Context) error {
	t := time.NewTicker(ringCheckPeriod)
	defer t.Stop()
	for {
		select {
		case <-ctx.Done():
			return nil
		case err := <-g.subservicesWatcher.Chan():
			return errors.Wrap(err, "running index gateway subservice failed")
		case <-t.C:
			continue
			// TODO: should we implement CAS check?
		}
	}
}

// stopping implements the Lifecycler interface and is one of the lifecycle hooks.
//
// Only invoked if the Index Gateway is in ring mode.
func (g *Gateway) stopping(_ error) error {
	level.Debug(util_log.Logger).Log("msg", "stopping index gateway")
	defer func() {
		g.indexQuerier.Stop()
		g.indexClient.Stop()
	}()
	return services.StopManagerAndAwaitStopped(context.Background(), g.subservices)
}

func (g *Gateway) QueryIndex(request *indexgatewaypb.QueryIndexRequest, server indexgatewaypb.IndexGateway_QueryIndexServer) error {
	var outerErr error
	var innerErr error

	queries := make([]index.Query, 0, len(request.Queries))
	for _, query := range request.Queries {
		queries = append(queries, index.Query{
			TableName:        query.TableName,
			HashValue:        query.HashValue,
			RangeValuePrefix: query.RangeValuePrefix,
			RangeValueStart:  query.RangeValueStart,
			ValueEqual:       query.ValueEqual,
		})
	}

	sendBatchMtx := sync.Mutex{}
	outerErr = g.indexClient.QueryPages(server.Context(), queries, func(query index.Query, batch index.ReadBatchResult) bool {
		innerErr = buildResponses(query, batch, func(response *indexgatewaypb.QueryIndexResponse) error {
			// do not send grpc responses concurrently. See https://github.com/grpc/grpc-go/blob/master/stream.go#L120-L123.
			sendBatchMtx.Lock()
			defer sendBatchMtx.Unlock()

			return server.Send(response)
		})

		if innerErr != nil {
			return false
		}

		return true
	})

	if innerErr != nil {
		return innerErr
	}

	return outerErr
}

func buildResponses(query index.Query, batch index.ReadBatchResult, callback func(*indexgatewaypb.QueryIndexResponse) error) error {
	itr := batch.Iterator()
	var resp []*indexgatewaypb.Row

	for itr.Next() {
		if len(resp) == maxIndexEntriesPerResponse {
			err := callback(&indexgatewaypb.QueryIndexResponse{
				QueryKey: util.QueryKey(query),
				Rows:     resp,
			})
			if err != nil {
				return err
			}
			resp = []*indexgatewaypb.Row{}
		}

		resp = append(resp, &indexgatewaypb.Row{
			RangeValue: itr.RangeValue(),
			Value:      itr.Value(),
		})
	}

	if len(resp) != 0 {
		err := callback(&indexgatewaypb.QueryIndexResponse{
			QueryKey: util.QueryKey(query),
			Rows:     resp,
		})
		if err != nil {
			return err
		}
	}

	return nil
}

func (g *Gateway) GetChunkRef(ctx context.Context, req *indexgatewaypb.GetChunkRefRequest) (*indexgatewaypb.GetChunkRefResponse, error) {
	instanceID, err := tenant.TenantID(ctx)
	if err != nil {
		return nil, err
	}
	matchers, err := syntax.ParseMatchers(req.Matchers)
	if err != nil {
		return nil, err
	}
	chunks, _, err := g.indexQuerier.GetChunkRefs(ctx, instanceID, req.From, req.Through, matchers...)
	if err != nil {
		return nil, err
	}
	result := &indexgatewaypb.GetChunkRefResponse{
		Refs: make([]*logproto.ChunkRef, 0, len(chunks)),
	}
	for _, cs := range chunks {
		for _, c := range cs {
			result.Refs = append(result.Refs, &c.ChunkRef)
		}
	}
	return result, nil
}

func (g *Gateway) LabelNamesForMetricName(ctx context.Context, req *indexgatewaypb.LabelNamesForMetricNameRequest) (*indexgatewaypb.LabelResponse, error) {
	instanceID, err := tenant.TenantID(ctx)
	if err != nil {
		return nil, err
	}
	names, err := g.indexQuerier.LabelNamesForMetricName(ctx, instanceID, req.From, req.Through, req.MetricName)
	if err != nil {
		return nil, err
	}
	return &indexgatewaypb.LabelResponse{
		Values: names,
	}, nil
}

func (g *Gateway) LabelValuesForMetricName(ctx context.Context, req *indexgatewaypb.LabelValuesForMetricNameRequest) (*indexgatewaypb.LabelResponse, error) {
	instanceID, err := tenant.TenantID(ctx)
	if err != nil {
		return nil, err
	}
	matchers, err := syntax.ParseMatchers(req.Matchers)
	if err != nil {
		return nil, err
	}
	names, err := g.indexQuerier.LabelValuesForMetricName(ctx, instanceID, req.From, req.Through, req.MetricName, req.LabelName, matchers...)
	if err != nil {
		return nil, err
	}
	return &indexgatewaypb.LabelResponse{
		Values: names,
	}, nil
}

// ServeHTTP serves the HTTP route /indexgateway/ring.
func (g *Gateway) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	if g.cfg.Mode == RingMode {
		g.ring.ServeHTTP(w, req)
	} else {
		w.Write([]byte("IndexGateway running with 'useIndexGatewayRing' disabled."))
	}
}
<commit_msg>fix index gateway chunk refs response (#5957)<commit_after>package indexgateway

import (
	"context"
	"fmt"
	"net/http"
	"sync"
	"time"

	"github.com/go-kit/log"
	"github.com/go-kit/log/level"
	"github.com/grafana/dskit/kv"
	"github.com/grafana/dskit/ring"
	"github.com/grafana/dskit/services"
	"github.com/grafana/dskit/tenant"
	"github.com/pkg/errors"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/common/model"
	"github.com/prometheus/prometheus/model/labels"

	"github.com/grafana/loki/pkg/logproto"
	"github.com/grafana/loki/pkg/logql/syntax"
	"github.com/grafana/loki/pkg/storage/chunk"
	"github.com/grafana/loki/pkg/storage/chunk/fetcher"
	"github.com/grafana/loki/pkg/storage/stores/series/index"
	"github.com/grafana/loki/pkg/storage/stores/shipper/indexgateway/indexgatewaypb"
	"github.com/grafana/loki/pkg/storage/stores/shipper/util"
	util_log "github.com/grafana/loki/pkg/util/log"
)

const (
	maxIndexEntriesPerResponse     = 1000
	ringAutoForgetUnhealthyPeriods = 10
	ringNameForServer              = "index-gateway"
	ringNumTokens                  = 128
	ringCheckPeriod                = 3 * time.Second

	// RingIdentifier is used as a unique name to register the Index Gateway ring.
	RingIdentifier = "index-gateway"

	// RingKey is the name of the key used to register the different Index Gateway instances in the key-value store.
	RingKey = "index-gateway"
)

type IndexQuerier interface {
	GetChunkRefs(ctx context.Context, userID string, from, through model.Time, matchers ...*labels.Matcher) ([][]chunk.Chunk, []*fetcher.Fetcher, error)
	LabelValuesForMetricName(ctx context.Context, userID string, from, through model.Time, metricName string, labelName string, matchers ...*labels.Matcher) ([]string, error)
	LabelNamesForMetricName(ctx context.Context, userID string, from, through model.Time, metricName string) ([]string, error)
	Stop()
}

type IndexClient interface {
	QueryPages(ctx context.Context, queries []index.Query, callback index.QueryPagesCallback) error
	Stop()
}

type Gateway struct {
	services.Service

	indexQuerier IndexQuerier
	indexClient  IndexClient

	cfg Config
	log log.Logger

	shipper IndexQuerier

	subservices        *services.Manager
	subservicesWatcher *services.FailureWatcher

	ringLifecycler *ring.BasicLifecycler
	ring           *ring.Ring
}

// NewIndexGateway instantiates a new Index Gateway and start its services.
//
// In case it is configured to be in ring mode, a Basic Service wrapping the ring client is started.
// Otherwise, it starts an Idle Service that doesn't have lifecycle hooks.
func NewIndexGateway(cfg Config, log log.Logger, registerer prometheus.Registerer, indexQuerier IndexQuerier, indexClient IndexClient) (*Gateway, error) {
	g := &Gateway{
		indexQuerier: indexQuerier,
		indexClient:  indexClient,
		cfg:          cfg,
		log:          log,
	}

	if cfg.Mode == RingMode {
		ringStore, err := kv.NewClient(
			cfg.Ring.KVStore,
			ring.GetCodec(),
			kv.RegistererWithKVName(prometheus.WrapRegistererWithPrefix("loki_", registerer), "index-gateway"),
			log,
		)
		if err != nil {
			return nil, errors.Wrap(err, "create KV store client")
		}

		lifecyclerCfg, err := cfg.Ring.ToLifecyclerConfig(ringNumTokens, log)
		if err != nil {
			return nil, errors.Wrap(err, "invalid ring lifecycler config")
		}

		delegate := ring.BasicLifecyclerDelegate(g)
		delegate = ring.NewLeaveOnStoppingDelegate(delegate, log)
		delegate = ring.NewTokensPersistencyDelegate(cfg.Ring.TokensFilePath, ring.JOINING, delegate, log)
		delegate = ring.NewAutoForgetDelegate(ringAutoForgetUnhealthyPeriods*cfg.Ring.HeartbeatTimeout, delegate, log)

		g.ringLifecycler, err = ring.NewBasicLifecycler(lifecyclerCfg, ringNameForServer, RingKey, ringStore, delegate, log, registerer)
		if err != nil {
			return nil, errors.Wrap(err, "index gateway create ring lifecycler")
		}

		ringCfg := cfg.Ring.ToRingConfig(cfg.Ring.ReplicationFactor)
		g.ring, err = ring.NewWithStoreClientAndStrategy(ringCfg, ringNameForServer, RingKey, ringStore, ring.NewIgnoreUnhealthyInstancesReplicationStrategy(), prometheus.WrapRegistererWithPrefix("loki_", registerer), log)
		if err != nil {
			return nil, errors.Wrap(err, "index gateway create ring client")
		}

		svcs := []services.Service{g.ringLifecycler, g.ring}
		g.subservices, err = services.NewManager(svcs...)
		if err != nil {
			return nil, fmt.Errorf("new index gateway services manager: %w", err)
		}

		g.subservicesWatcher = services.NewFailureWatcher()
		g.subservicesWatcher.WatchManager(g.subservices)
		g.Service = services.NewBasicService(g.starting, g.running, g.stopping)
	} else {
		g.Service = services.NewIdleService(nil, func(failureCase error) error {
			g.indexQuerier.Stop()
			g.indexClient.Stop()
			return nil
		})
	}

	return g, nil
}

// starting implements the Lifecycler interface and is one of the lifecycle hooks.
//
// Only invoked if the Index Gateway is in ring mode.
func (g *Gateway) starting(ctx context.Context) (err error) {
	// In case this function will return error we want to unregister the instance
	// from the ring. We do it ensuring dependencies are gracefully stopped if they
	// were already started.
	defer func() {
		if err == nil || g.subservices == nil {
			return
		}

		if stopErr := services.StopManagerAndAwaitStopped(context.Background(), g.subservices); stopErr != nil {
			level.Error(util_log.Logger).Log("msg", "failed to gracefully stop index gateway dependencies", "err", stopErr)
		}
	}()

	if err := services.StartManagerAndAwaitHealthy(ctx, g.subservices); err != nil {
		return errors.Wrap(err, "unable to start index gateway subservices")
	}

	// The BasicLifecycler does not automatically move state to ACTIVE such that any additional work that
	// someone wants to do can be done before becoming ACTIVE. For the index gateway we don't currently
	// have any additional work so we can become ACTIVE right away.
	// Wait until the ring client detected this instance in the JOINING
	// state to make sure that when we'll run the initial sync we already
	// know the tokens assigned to this instance.
	level.Info(util_log.Logger).Log("msg", "waiting until index gateway is JOINING in the ring")
	if err := ring.WaitInstanceState(ctx, g.ring, g.ringLifecycler.GetInstanceID(), ring.JOINING); err != nil {
		return err
	}
	level.Info(util_log.Logger).Log("msg", "index gateway is JOINING in the ring")

	if err = g.ringLifecycler.ChangeState(ctx, ring.ACTIVE); err != nil {
		return errors.Wrapf(err, "switch instance to %s in the ring", ring.ACTIVE)
	}

	// Wait until the ring client detected this instance in the ACTIVE state to
	// make sure that when we'll run the loop it won't be detected as a ring
	// topology change.
	level.Info(util_log.Logger).Log("msg", "waiting until index gateway is ACTIVE in the ring")
	if err := ring.WaitInstanceState(ctx, g.ring, g.ringLifecycler.GetInstanceID(), ring.ACTIVE); err != nil {
		return err
	}
	level.Info(util_log.Logger).Log("msg", "index gateway is ACTIVE in the ring")

	return nil
}

// running implements the Lifecycler interface and is one of the lifecycle hooks.
//
// Only invoked if the Index Gateway is in ring mode.
func (g *Gateway) running(ctx context.Context) error {
	t := time.NewTicker(ringCheckPeriod)
	defer t.Stop()
	for {
		select {
		case <-ctx.Done():
			return nil
		case err := <-g.subservicesWatcher.Chan():
			return errors.Wrap(err, "running index gateway subservice failed")
		case <-t.C:
			continue
			// TODO: should we implement CAS check?
		}
	}
}

// stopping implements the Lifecycler interface and is one of the lifecycle hooks.
//
// Only invoked if the Index Gateway is in ring mode.
func (g *Gateway) stopping(_ error) error {
	level.Debug(util_log.Logger).Log("msg", "stopping index gateway")
	defer func() {
		g.indexQuerier.Stop()
		g.indexClient.Stop()
	}()
	return services.StopManagerAndAwaitStopped(context.Background(), g.subservices)
}

func (g *Gateway) QueryIndex(request *indexgatewaypb.QueryIndexRequest, server indexgatewaypb.IndexGateway_QueryIndexServer) error {
	var outerErr error
	var innerErr error

	queries := make([]index.Query, 0, len(request.Queries))
	for _, query := range request.Queries {
		queries = append(queries, index.Query{
			TableName:        query.TableName,
			HashValue:        query.HashValue,
			RangeValuePrefix: query.RangeValuePrefix,
			RangeValueStart:  query.RangeValueStart,
			ValueEqual:       query.ValueEqual,
		})
	}

	sendBatchMtx := sync.Mutex{}
	outerErr = g.indexClient.QueryPages(server.Context(), queries, func(query index.Query, batch index.ReadBatchResult) bool {
		innerErr = buildResponses(query, batch, func(response *indexgatewaypb.QueryIndexResponse) error {
			// do not send grpc responses concurrently. See https://github.com/grpc/grpc-go/blob/master/stream.go#L120-L123.
			sendBatchMtx.Lock()
			defer sendBatchMtx.Unlock()

			return server.Send(response)
		})

		if innerErr != nil {
			return false
		}

		return true
	})

	if innerErr != nil {
		return innerErr
	}

	return outerErr
}

func buildResponses(query index.Query, batch index.ReadBatchResult, callback func(*indexgatewaypb.QueryIndexResponse) error) error {
	itr := batch.Iterator()
	var resp []*indexgatewaypb.Row

	for itr.Next() {
		if len(resp) == maxIndexEntriesPerResponse {
			err := callback(&indexgatewaypb.QueryIndexResponse{
				QueryKey: util.QueryKey(query),
				Rows:     resp,
			})
			if err != nil {
				return err
			}
			resp = []*indexgatewaypb.Row{}
		}

		resp = append(resp, &indexgatewaypb.Row{
			RangeValue: itr.RangeValue(),
			Value:      itr.Value(),
		})
	}

	if len(resp) != 0 {
		err := callback(&indexgatewaypb.QueryIndexResponse{
			QueryKey: util.QueryKey(query),
			Rows:     resp,
		})
		if err != nil {
			return err
		}
	}

	return nil
}

func (g *Gateway) GetChunkRef(ctx context.Context, req *indexgatewaypb.GetChunkRefRequest) (*indexgatewaypb.GetChunkRefResponse, error) {
	instanceID, err := tenant.TenantID(ctx)
	if err != nil {
		return nil, err
	}
	matchers, err := syntax.ParseMatchers(req.Matchers)
	if err != nil {
		return nil, err
	}
	chunks, _, err := g.indexQuerier.GetChunkRefs(ctx, instanceID, req.From, req.Through, matchers...)
	if err != nil {
		return nil, err
	}
	result := &indexgatewaypb.GetChunkRefResponse{
		Refs: make([]*logproto.ChunkRef, 0, len(chunks)),
	}
	for _, cs := range chunks {
		for i := range cs {
			result.Refs = append(result.Refs, &cs[i].ChunkRef)
		}
	}
	return result, nil
}

func (g *Gateway) LabelNamesForMetricName(ctx context.Context, req *indexgatewaypb.LabelNamesForMetricNameRequest) (*indexgatewaypb.LabelResponse, error) {
	instanceID, err := tenant.TenantID(ctx)
	if err != nil {
		return nil, err
	}
	names, err := g.indexQuerier.LabelNamesForMetricName(ctx, instanceID, req.From, req.Through, req.MetricName)
	if err != nil {
		return nil, err
	}
	return &indexgatewaypb.LabelResponse{
		Values: names,
	}, nil
}

func (g *Gateway) LabelValuesForMetricName(ctx context.Context, req *indexgatewaypb.LabelValuesForMetricNameRequest) (*indexgatewaypb.LabelResponse, error) {
	instanceID, err := tenant.TenantID(ctx)
	if err != nil {
		return nil, err
	}
	matchers, err := syntax.ParseMatchers(req.Matchers)
	if err != nil {
		return nil, err
	}
	names, err := g.indexQuerier.LabelValuesForMetricName(ctx, instanceID, req.From, req.Through, req.MetricName, req.LabelName, matchers...)
	if err != nil {
		return nil, err
	}
	return &indexgatewaypb.LabelResponse{
		Values: names,
	}, nil
}

// ServeHTTP serves the HTTP route /indexgateway/ring.
func (g *Gateway) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	if g.cfg.Mode == RingMode {
		g.ring.ServeHTTP(w, req)
	} else {
		w.Write([]byte("IndexGateway running with 'useIndexGatewayRing' disabled."))
	}
}
<|endoftext|>


<commit_before>// Copyright 2018 PingCAP, Inc.
//
// 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,
// See the License for the specific language governing permissions and
// limitations under the License.

package decoder

import (
	"sort"
	"time"

	"github.com/pingcap/parser/model"
	"github.com/pingcap/parser/mysql"
	"github.com/pingcap/tidb/expression"
	"github.com/pingcap/tidb/kv"
	"github.com/pingcap/tidb/sessionctx"
	"github.com/pingcap/tidb/table"
	"github.com/pingcap/tidb/table/tables"
	"github.com/pingcap/tidb/tablecodec"
	"github.com/pingcap/tidb/types"
	"github.com/pingcap/tidb/util/chunk"
	"github.com/pingcap/tidb/util/rowcodec"
)

// Column contains the info and generated expr of column.
type Column struct {
	Col     *table.Column
	GenExpr expression.Expression
}

// RowDecoder decodes a byte slice into datums and eval the generated column value.
type RowDecoder struct {
	tbl           table.Table
	mutRow        chunk.MutRow
	colMap        map[int64]Column
	colTypes      map[int64]*types.FieldType
	haveGenColumn bool
	defaultVals   []types.Datum
	cols          []*table.Column
	pkCols        []int64
}

// NewRowDecoder returns a new RowDecoder.
func NewRowDecoder(tbl table.Table, cols []*table.Column, decodeColMap map[int64]Column) *RowDecoder {
	tblInfo := tbl.Meta()
	colFieldMap := make(map[int64]*types.FieldType, len(decodeColMap))
	for id, col := range decodeColMap {
		colFieldMap[id] = &col.Col.ColumnInfo.FieldType
	}

	tps := make([]*types.FieldType, len(cols))
	for _, col := range cols {
		if col.ChangeStateInfo == nil {
			tps[col.Offset] = &col.FieldType
		} else {
			// Since changing column in the mutRow will be set with relative column's old value in the process of column-type-change,
			// we should set fieldType as the relative column does. Otherwise it may get a panic, take change json to int as an example,
			// setting json value to a int type column in mutRow will panic because it lacks of offset array.
			tps[col.Offset] = &cols[col.ChangeStateInfo.DependencyColumnOffset].FieldType
		}
	}
	var pkCols []int64
	switch {
	case tblInfo.IsCommonHandle:
		pkCols = tables.TryGetCommonPkColumnIds(tbl.Meta())
	case tblInfo.PKIsHandle:
		pkCols = []int64{tblInfo.GetPkColInfo().ID}
	default: // support decoding _tidb_rowid.
		pkCols = []int64{model.ExtraHandleID}
	}
	return &RowDecoder{
		tbl:         tbl,
		mutRow:      chunk.MutRowFromTypes(tps),
		colMap:      decodeColMap,
		colTypes:    colFieldMap,
		defaultVals: make([]types.Datum, len(cols)),
		cols:        cols,
		pkCols:      pkCols,
	}
}

// DecodeAndEvalRowWithMap decodes a byte slice into datums and evaluates the generated column value.
func (rd *RowDecoder) DecodeAndEvalRowWithMap(ctx sessionctx.Context, handle kv.Handle, b []byte, decodeLoc, sysLoc *time.Location, row map[int64]types.Datum) (map[int64]types.Datum, error) {
	var err error
	if rowcodec.IsNewFormat(b) {
		row, err = tablecodec.DecodeRowWithMapNew(b, rd.colTypes, decodeLoc, row)
	} else {
		row, err = tablecodec.DecodeRowWithMap(b, rd.colTypes, decodeLoc, row)
	}
	if err != nil {
		return nil, err
	}
	row, err = tablecodec.DecodeHandleToDatumMap(handle, rd.pkCols, rd.colTypes, decodeLoc, row)
	if err != nil {
		return nil, err
	}
	for _, dCol := range rd.colMap {
		colInfo := dCol.Col.ColumnInfo
		val, ok := row[colInfo.ID]
		if ok || dCol.GenExpr != nil {
			rd.mutRow.SetValue(colInfo.Offset, val.GetValue())
			continue
		}
		if dCol.Col.ChangeStateInfo != nil {
			val, _, err = tables.GetChangingColVal(ctx, rd.cols, dCol.Col, row, rd.defaultVals)
		} else {
			// Get the default value of the column in the generated column expression.
			val, err = tables.GetColDefaultValue(ctx, dCol.Col, rd.defaultVals)
		}
		if err != nil {
			return nil, err
		}
		rd.mutRow.SetValue(colInfo.Offset, val.GetValue())
	}
	keys := make([]int, 0)
	ids := make(map[int]int)
	for k, col := range rd.colMap {
		keys = append(keys, col.Col.Offset)
		ids[col.Col.Offset] = int(k)
	}
	sort.Ints(keys)
	for _, id := range keys {
		col := rd.colMap[int64(ids[id])]
		if col.GenExpr == nil {
			continue
		}
		// Eval the column value
		val, err := col.GenExpr.Eval(rd.mutRow.ToRow())
		if err != nil {
			return nil, err
		}
		val, err = table.CastValue(ctx, val, col.Col.ColumnInfo, false, true)
		if err != nil {
			return nil, err
		}

		if val.Kind() == types.KindMysqlTime && sysLoc != time.UTC {
			t := val.GetMysqlTime()
			if t.Type() == mysql.TypeTimestamp {
				err := t.ConvertTimeZone(sysLoc, time.UTC)
				if err != nil {
					return nil, err
				}
				val.SetMysqlTime(t)
			}
		}
		rd.mutRow.SetValue(col.Col.Offset, val.GetValue())

		row[int64(ids[id])] = val
	}
	return row, nil
}

// BuildFullDecodeColMap builds a map that contains [columnID -> struct{*table.Column, expression.Expression}] from all columns.
func BuildFullDecodeColMap(cols []*table.Column, schema *expression.Schema) map[int64]Column {
	decodeColMap := make(map[int64]Column, len(cols))
	for _, col := range cols {
		decodeColMap[col.ID] = Column{
			Col:     col,
			GenExpr: schema.Columns[col.Offset].VirtualExpr,
		}
	}
	return decodeColMap
}

// CurrentRowWithDefaultVal returns current decoding row with default column values set properly.
// Please make sure calling DecodeAndEvalRowWithMap first.
func (rd *RowDecoder) CurrentRowWithDefaultVal() chunk.Row {
	return rd.mutRow.ToRow()
}
<commit_msg>util: pre-alloc memory for decoding rows (#22508)<commit_after>// Copyright 2018 PingCAP, Inc.
//
// 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,
// See the License for the specific language governing permissions and
// limitations under the License.

package decoder

import (
	"sort"
	"time"

	"github.com/pingcap/parser/model"
	"github.com/pingcap/parser/mysql"
	"github.com/pingcap/tidb/expression"
	"github.com/pingcap/tidb/kv"
	"github.com/pingcap/tidb/sessionctx"
	"github.com/pingcap/tidb/table"
	"github.com/pingcap/tidb/table/tables"
	"github.com/pingcap/tidb/tablecodec"
	"github.com/pingcap/tidb/types"
	"github.com/pingcap/tidb/util/chunk"
	"github.com/pingcap/tidb/util/rowcodec"
)

// Column contains the info and generated expr of column.
type Column struct {
	Col     *table.Column
	GenExpr expression.Expression
}

// RowDecoder decodes a byte slice into datums and eval the generated column value.
type RowDecoder struct {
	tbl           table.Table
	mutRow        chunk.MutRow
	colMap        map[int64]Column
	colTypes      map[int64]*types.FieldType
	haveGenColumn bool
	defaultVals   []types.Datum
	cols          []*table.Column
	pkCols        []int64
}

// NewRowDecoder returns a new RowDecoder.
func NewRowDecoder(tbl table.Table, cols []*table.Column, decodeColMap map[int64]Column) *RowDecoder {
	tblInfo := tbl.Meta()
	colFieldMap := make(map[int64]*types.FieldType, len(decodeColMap))
	for id, col := range decodeColMap {
		colFieldMap[id] = &col.Col.ColumnInfo.FieldType
	}

	tps := make([]*types.FieldType, len(cols))
	for _, col := range cols {
		if col.ChangeStateInfo == nil {
			tps[col.Offset] = &col.FieldType
		} else {
			// Since changing column in the mutRow will be set with relative column's old value in the process of column-type-change,
			// we should set fieldType as the relative column does. Otherwise it may get a panic, take change json to int as an example,
			// setting json value to a int type column in mutRow will panic because it lacks of offset array.
			tps[col.Offset] = &cols[col.ChangeStateInfo.DependencyColumnOffset].FieldType
		}
	}
	var pkCols []int64
	switch {
	case tblInfo.IsCommonHandle:
		pkCols = tables.TryGetCommonPkColumnIds(tbl.Meta())
	case tblInfo.PKIsHandle:
		pkCols = []int64{tblInfo.GetPkColInfo().ID}
	default: // support decoding _tidb_rowid.
		pkCols = []int64{model.ExtraHandleID}
	}
	return &RowDecoder{
		tbl:         tbl,
		mutRow:      chunk.MutRowFromTypes(tps),
		colMap:      decodeColMap,
		colTypes:    colFieldMap,
		defaultVals: make([]types.Datum, len(cols)),
		cols:        cols,
		pkCols:      pkCols,
	}
}

// DecodeAndEvalRowWithMap decodes a byte slice into datums and evaluates the generated column value.
func (rd *RowDecoder) DecodeAndEvalRowWithMap(ctx sessionctx.Context, handle kv.Handle, b []byte, decodeLoc, sysLoc *time.Location, row map[int64]types.Datum) (map[int64]types.Datum, error) {
	var err error
	if rowcodec.IsNewFormat(b) {
		row, err = tablecodec.DecodeRowWithMapNew(b, rd.colTypes, decodeLoc, row)
	} else {
		row, err = tablecodec.DecodeRowWithMap(b, rd.colTypes, decodeLoc, row)
	}
	if err != nil {
		return nil, err
	}
	row, err = tablecodec.DecodeHandleToDatumMap(handle, rd.pkCols, rd.colTypes, decodeLoc, row)
	if err != nil {
		return nil, err
	}
	for _, dCol := range rd.colMap {
		colInfo := dCol.Col.ColumnInfo
		val, ok := row[colInfo.ID]
		if ok || dCol.GenExpr != nil {
			rd.mutRow.SetValue(colInfo.Offset, val.GetValue())
			continue
		}
		if dCol.Col.ChangeStateInfo != nil {
			val, _, err = tables.GetChangingColVal(ctx, rd.cols, dCol.Col, row, rd.defaultVals)
		} else {
			// Get the default value of the column in the generated column expression.
			val, err = tables.GetColDefaultValue(ctx, dCol.Col, rd.defaultVals)
		}
		if err != nil {
			return nil, err
		}
		rd.mutRow.SetValue(colInfo.Offset, val.GetValue())
	}
	keys := make([]int, 0, len(rd.colMap))
	ids := make(map[int]int, len(rd.colMap))
	for k, col := range rd.colMap {
		keys = append(keys, col.Col.Offset)
		ids[col.Col.Offset] = int(k)
	}
	sort.Ints(keys)
	for _, id := range keys {
		col := rd.colMap[int64(ids[id])]
		if col.GenExpr == nil {
			continue
		}
		// Eval the column value
		val, err := col.GenExpr.Eval(rd.mutRow.ToRow())
		if err != nil {
			return nil, err
		}
		val, err = table.CastValue(ctx, val, col.Col.ColumnInfo, false, true)
		if err != nil {
			return nil, err
		}

		if val.Kind() == types.KindMysqlTime && sysLoc != time.UTC {
			t := val.GetMysqlTime()
			if t.Type() == mysql.TypeTimestamp {
				err := t.ConvertTimeZone(sysLoc, time.UTC)
				if err != nil {
					return nil, err
				}
				val.SetMysqlTime(t)
			}
		}
		rd.mutRow.SetValue(col.Col.Offset, val.GetValue())

		row[int64(ids[id])] = val
	}
	return row, nil
}

// BuildFullDecodeColMap builds a map that contains [columnID -> struct{*table.Column, expression.Expression}] from all columns.
func BuildFullDecodeColMap(cols []*table.Column, schema *expression.Schema) map[int64]Column {
	decodeColMap := make(map[int64]Column, len(cols))
	for _, col := range cols {
		decodeColMap[col.ID] = Column{
			Col:     col,
			GenExpr: schema.Columns[col.Offset].VirtualExpr,
		}
	}
	return decodeColMap
}

// CurrentRowWithDefaultVal returns current decoding row with default column values set properly.
// Please make sure calling DecodeAndEvalRowWithMap first.
func (rd *RowDecoder) CurrentRowWithDefaultVal() chunk.Row {
	return rd.mutRow.ToRow()
}
<|endoftext|>


<commit_before>package docker

import (
	"bufio"
	"bytes"
	"context"
	"github.com/ViBiOh/docker-deploy/jsonHttp"
	"github.com/docker/docker/api/types"
	"github.com/docker/docker/api/types/container"
	"github.com/docker/docker/api/types/network"
	"github.com/docker/docker/api/types/strslice"
	"github.com/docker/docker/client"
	"gopkg.in/yaml.v2"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"regexp"
	"strings"
)

const host = `DOCKER_HOST`
const version = `DOCKER_VERSION`
const configurationFile = `./users`

var commaByte = []byte(`,`)
var splitLogs = regexp.MustCompile(`.{8}(.*?)\n`)

const cpuShare = 128
const memoryLimit = 134217728
const securityOptions = []string{`no-new-privileges`}

var jsonLogging = container.LogConfig{Type: `json-file`, Config: map[string]string{
	`max-size`: `50m`,
}}
var onFailureRestart = container.RestartPolicy{Name: `on-failure`, MaximumRetryCount: 5}

var containersRequest = regexp.MustCompile(`/containers/?$`)
var containerRequest = regexp.MustCompile(`/containers/([^/]+)/?$`)
var startRequest = regexp.MustCompile(`/containers/([^/]+)/start`)
var stopRequest = regexp.MustCompile(`/containers/([^/]+)/stop`)
var restartRequest = regexp.MustCompile(`/containers/([^/]+)/restart`)
var logRequest = regexp.MustCompile(`/containers/([^/]+)/logs`)

type results struct {
	Results interface{} `json:"results"`
}

type user struct {
	username string
	password string
}

type dockerCompose struct {
	Version  string
	Services map[string]struct {
		Image       string
		Command     string
		Environment map[string]string
		Labels      map[string]string
	}
}

var docker *client.Client
var users map[string]*user

func handleError(w http.ResponseWriter, err error) {
	log.Print(err)
	http.Error(w, err.Error(), http.StatusInternalServerError)
}

func readConfiguration(path string) map[string]*user {
	configFile, err := os.Open(path)
	defer configFile.Close()

	if err != nil {
		log.Print(err)
		return nil
	}

	users := make(map[string]*user)

	scanner := bufio.NewScanner(configFile)
	for scanner.Scan() {
		parts := bytes.Split(scanner.Bytes(), commaByte)
		user := user{string(parts[0]), string(parts[1])}

		users[strings.ToLower(user.username)] = &user
	}

	return users
}

func init() {
	users = readConfiguration(configurationFile)

	client, err := client.NewClient(os.Getenv(host), os.Getenv(version), nil, nil)
	if err != nil {
		log.Fatal(err)
	} else {
		docker = client
	}
}

func inspectContainer(w http.ResponseWriter, containerID []byte) {
	if container, err := docker.ContainerInspect(context.Background(), string(containerID)); err != nil {
		handleError(w, err)
	} else {
		jsonHttp.ResponseJSON(w, container)
	}
}

func startContainer(w http.ResponseWriter, containerID []byte) {
	if err := docker.ContainerStart(context.Background(), string(containerID), types.ContainerStartOptions{}); err != nil {
		handleError(w, err)
	} else {
		w.Write(nil)
	}
}

func stopContainer(w http.ResponseWriter, containerID []byte) {
	if err := docker.ContainerStop(context.Background(), string(containerID), nil); err != nil {
		handleError(w, err)
	} else {
		w.Write(nil)
	}
}

func restartContainer(w http.ResponseWriter, containerID []byte) {
	if err := docker.ContainerRestart(context.Background(), string(containerID), nil); err != nil {
		handleError(w, err)
	} else {
		w.Write(nil)
	}
}

func logContainer(w http.ResponseWriter, containerID []byte) {
	logs, err := docker.ContainerLogs(context.Background(), string(containerID), types.ContainerLogsOptions{ShowStdout: true, ShowStderr: true, Follow: false})
	if err != nil {
		handleError(w, err)
		return
	}

	defer logs.Close()

	if logLines, err := ioutil.ReadAll(logs); err != nil {
		handleError(w, err)
	} else {
		matches := splitLogs.FindAllSubmatch(logLines, -1)
		cleanLogs := make([]string, 0, len(matches))
		for _, match := range matches {
			cleanLogs = append(cleanLogs, string(match[1]))
		}

		jsonHttp.ResponseJSON(w, results{cleanLogs})
	}
}

func listContainers(w http.ResponseWriter) {
	if containers, err := docker.ContainerList(context.Background(), types.ContainerListOptions{All: true}); err != nil {
		handleError(w, err)
	} else {
		jsonHttp.ResponseJSON(w, results{containers})
	}
}

func readBody(body io.ReadCloser) ([]byte, error) {
	defer body.Close()
	return ioutil.ReadAll(body)
}

func runCompose(w http.ResponseWriter, name []byte, composeFile []byte) {
	compose := dockerCompose{}

	if err := yaml.Unmarshal(composeFile, &compose); err != nil {
		handleError(w, err)
		return
	}

	ids := make([]string, len(compose.Services))
	for serviceName, service := range compose.Services {
		environments := make([]string, len(service.Environment))
		for key, value := range service.Environment {
			environments = append(environments, key+`=`+value)
		}

		config := container.Config{
			Image:  service.Image,
			Labels: service.Labels,
			Env:    environments,
		}

		if service.Command != `` {
			config.Cmd = strslice.StrSlice([]string{service.Command})
		}

		id, err := docker.ContainerCreate(
			context.Background(),
			&config,
			&container.HostConfig{
				LogConfig:      jsonLogging,
				RestartPolicy:  onFailureRestart,
				ReadonlyRootfs: true,
				Resources: container.Resources{
					CPUShares: cpuShare,
					Memory:    memoryLimit,
				},
				SecurityOpt: securityOptions,
			},
			&network.NetworkingConfig{
				EndpointsConfig: map[string]*network.EndpointSettings{
					`traefik`: &network.EndpointSettings{},
				},
			},
			string(name)+`_`+serviceName,
		)

		if err != nil {
			handleError(w, err)
			return
		}

		ids = append(ids, id.ID)
	}

	jsonHttp.ResponseJSON(w, results{ids})
}

func isAuthenticated(r *http.Request) bool {
	username, password, ok := r.BasicAuth()

	if ok {
		user, ok := users[strings.ToLower(username)]

		if ok && user.password == password {
			return true
		}
	}

	return false
}

func unauthorized(w http.ResponseWriter) {
	http.Error(w, `Authentication required`, http.StatusUnauthorized)
}

// Handler for Hello request. Should be use with net/http
type Handler struct {
}

func (handler Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	w.Header().Add(`Access-Control-Allow-Origin`, `*`)
	w.Header().Add(`Access-Control-Allow-Headers`, `Content-Type, Authorization`)
	w.Header().Add(`Access-Control-Allow-Methods`, `GET, POST`)
	w.Header().Add(`X-Content-Type-Options`, `nosniff`)

	if r.Method == http.MethodOptions {
		w.Write(nil)
		return
	}

	urlPath := []byte(r.URL.Path)

	if containersRequest.Match(urlPath) && r.Method == http.MethodGet {
		listContainers(w)
	} else if isAuthenticated(r) {
		if containerRequest.Match(urlPath) && r.Method == http.MethodPost {
			if composeBody, err := readBody(r.Body); err != nil {
				handleError(w, err)
			} else {
				runCompose(w, containerRequest.FindSubmatch(urlPath)[1], composeBody)
			}
		} else if containerRequest.Match(urlPath) && r.Method == http.MethodGet {
			inspectContainer(w, containerRequest.FindSubmatch(urlPath)[1])
		} else if startRequest.Match(urlPath) && r.Method == http.MethodPost {
			startContainer(w, startRequest.FindSubmatch(urlPath)[1])
		} else if stopRequest.Match(urlPath) && r.Method == http.MethodPost {
			stopContainer(w, stopRequest.FindSubmatch(urlPath)[1])
		} else if restartRequest.Match(urlPath) && r.Method == http.MethodPost {
			restartContainer(w, restartRequest.FindSubmatch(urlPath)[1])
		} else if logRequest.Match(urlPath) && r.Method == http.MethodGet {
			logContainer(w, logRequest.FindSubmatch(urlPath)[1])
		}
	} else {
		unauthorized(w)
	}
}
<commit_msg>Removing read only root fs<commit_after>package docker

import (
	"bufio"
	"bytes"
	"context"
	"github.com/ViBiOh/docker-deploy/jsonHttp"
	"github.com/docker/docker/api/types"
	"github.com/docker/docker/api/types/container"
	"github.com/docker/docker/api/types/network"
	"github.com/docker/docker/api/types/strslice"
	"github.com/docker/docker/client"
	"gopkg.in/yaml.v2"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"regexp"
	"strings"
)

const host = `DOCKER_HOST`
const version = `DOCKER_VERSION`
const configurationFile = `./users`

var commaByte = []byte(`,`)
var splitLogs = regexp.MustCompile(`.{8}(.*?)\n`)

var hostConfig = container.HostConfig{
	LogConfig: container.LogConfig{Type: `json-file`, Config: map[string]string{
		`max-size`: `50m`,
	}},
	RestartPolicy: container.RestartPolicy{Name: `on-failure`, MaximumRetryCount: 5},
	Resources: container.Resources{
		CPUShares: 128,
		Memory:    134217728,
	},
	SecurityOpt: []string{`no-new-privileges`},
}
var networkConfig = network.NetworkingConfig{
	EndpointsConfig: map[string]*network.EndpointSettings{
		`traefik`: &network.EndpointSettings{},
	},
}

var containersRequest = regexp.MustCompile(`/containers/?$`)
var containerRequest = regexp.MustCompile(`/containers/([^/]+)/?$`)
var startRequest = regexp.MustCompile(`/containers/([^/]+)/start`)
var stopRequest = regexp.MustCompile(`/containers/([^/]+)/stop`)
var restartRequest = regexp.MustCompile(`/containers/([^/]+)/restart`)
var logRequest = regexp.MustCompile(`/containers/([^/]+)/logs`)

type results struct {
	Results interface{} `json:"results"`
}

type user struct {
	username string
	password string
}

type dockerCompose struct {
	Version  string
	Services map[string]struct {
		Image       string
		Command     string
		Environment map[string]string
		Labels      map[string]string
	}
}

var docker *client.Client
var users map[string]*user

func handleError(w http.ResponseWriter, err error) {
	log.Print(err)
	http.Error(w, err.Error(), http.StatusInternalServerError)
}

func readConfiguration(path string) map[string]*user {
	configFile, err := os.Open(path)
	defer configFile.Close()

	if err != nil {
		log.Print(err)
		return nil
	}

	users := make(map[string]*user)

	scanner := bufio.NewScanner(configFile)
	for scanner.Scan() {
		parts := bytes.Split(scanner.Bytes(), commaByte)
		user := user{string(parts[0]), string(parts[1])}

		users[strings.ToLower(user.username)] = &user
	}

	return users
}

func init() {
	users = readConfiguration(configurationFile)

	client, err := client.NewClient(os.Getenv(host), os.Getenv(version), nil, nil)
	if err != nil {
		log.Fatal(err)
	} else {
		docker = client
	}
}

func inspectContainer(w http.ResponseWriter, containerID []byte) {
	if container, err := docker.ContainerInspect(context.Background(), string(containerID)); err != nil {
		handleError(w, err)
	} else {
		jsonHttp.ResponseJSON(w, container)
	}
}

func startContainer(w http.ResponseWriter, containerID []byte) {
	if err := docker.ContainerStart(context.Background(), string(containerID), types.ContainerStartOptions{}); err != nil {
		handleError(w, err)
	} else {
		w.Write(nil)
	}
}

func stopContainer(w http.ResponseWriter, containerID []byte) {
	if err := docker.ContainerStop(context.Background(), string(containerID), nil); err != nil {
		handleError(w, err)
	} else {
		w.Write(nil)
	}
}

func restartContainer(w http.ResponseWriter, containerID []byte) {
	if err := docker.ContainerRestart(context.Background(), string(containerID), nil); err != nil {
		handleError(w, err)
	} else {
		w.Write(nil)
	}
}

func logContainer(w http.ResponseWriter, containerID []byte) {
	logs, err := docker.ContainerLogs(context.Background(), string(containerID), types.ContainerLogsOptions{ShowStdout: true, ShowStderr: true, Follow: false})
	if err != nil {
		handleError(w, err)
		return
	}

	defer logs.Close()

	if logLines, err := ioutil.ReadAll(logs); err != nil {
		handleError(w, err)
	} else {
		matches := splitLogs.FindAllSubmatch(logLines, -1)
		cleanLogs := make([]string, 0, len(matches))
		for _, match := range matches {
			cleanLogs = append(cleanLogs, string(match[1]))
		}

		jsonHttp.ResponseJSON(w, results{cleanLogs})
	}
}

func listContainers(w http.ResponseWriter) {
	if containers, err := docker.ContainerList(context.Background(), types.ContainerListOptions{All: true}); err != nil {
		handleError(w, err)
	} else {
		jsonHttp.ResponseJSON(w, results{containers})
	}
}

func readBody(body io.ReadCloser) ([]byte, error) {
	defer body.Close()
	return ioutil.ReadAll(body)
}

func runCompose(w http.ResponseWriter, name []byte, composeFile []byte) {
	compose := dockerCompose{}

	if err := yaml.Unmarshal(composeFile, &compose); err != nil {
		handleError(w, err)
		return
	}

	ids := make([]string, len(compose.Services))
	for serviceName, service := range compose.Services {
		environments := make([]string, len(service.Environment))
		for key, value := range service.Environment {
			environments = append(environments, key+`=`+value)
		}

		config := container.Config{
			Image:  service.Image,
			Labels: service.Labels,
			Env:    environments,
		}

		if service.Command != `` {
			config.Cmd = strslice.StrSlice([]string{service.Command})
		}

		id, err := docker.ContainerCreate(context.Background(), &config, &hostConfig, &networkConfig, string(name)+`_`+serviceName)

		if err != nil {
			handleError(w, err)
			return
		}

		ids = append(ids, id.ID)
	}

	jsonHttp.ResponseJSON(w, results{ids})
}

func isAuthenticated(r *http.Request) bool {
	username, password, ok := r.BasicAuth()

	if ok {
		user, ok := users[strings.ToLower(username)]

		if ok && user.password == password {
			return true
		}
	}

	return false
}

func unauthorized(w http.ResponseWriter) {
	http.Error(w, `Authentication required`, http.StatusUnauthorized)
}

// Handler for Hello request. Should be use with net/http
type Handler struct {
}

func (handler Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	w.Header().Add(`Access-Control-Allow-Origin`, `*`)
	w.Header().Add(`Access-Control-Allow-Headers`, `Content-Type, Authorization`)
	w.Header().Add(`Access-Control-Allow-Methods`, `GET, POST`)
	w.Header().Add(`X-Content-Type-Options`, `nosniff`)

	if r.Method == http.MethodOptions {
		w.Write(nil)
		return
	}

	urlPath := []byte(r.URL.Path)

	if containersRequest.Match(urlPath) && r.Method == http.MethodGet {
		listContainers(w)
	} else if isAuthenticated(r) {
		if containerRequest.Match(urlPath) && r.Method == http.MethodPost {
			if composeBody, err := readBody(r.Body); err != nil {
				handleError(w, err)
			} else {
				runCompose(w, containerRequest.FindSubmatch(urlPath)[1], composeBody)
			}
		} else if containerRequest.Match(urlPath) && r.Method == http.MethodGet {
			inspectContainer(w, containerRequest.FindSubmatch(urlPath)[1])
		} else if startRequest.Match(urlPath) && r.Method == http.MethodPost {
			startContainer(w, startRequest.FindSubmatch(urlPath)[1])
		} else if stopRequest.Match(urlPath) && r.Method == http.MethodPost {
			stopContainer(w, stopRequest.FindSubmatch(urlPath)[1])
		} else if restartRequest.Match(urlPath) && r.Method == http.MethodPost {
			restartContainer(w, restartRequest.FindSubmatch(urlPath)[1])
		} else if logRequest.Match(urlPath) && r.Method == http.MethodGet {
			logContainer(w, logRequest.FindSubmatch(urlPath)[1])
		}
	} else {
		unauthorized(w)
	}
}
<|endoftext|>


<commit_before>package alice

// Module is a marker interface for structs that defines how to initialize instances.
type Module interface {
	// IsModule indicates if this is a module.
	IsModule() bool
}

// BaseModule is an implementation of Module interface. It should be embeded into each module defined in the
// application.
//
// A typical module is defined as follows:
//
//	type ExampleModule struct {
//		alice.BaseModule
//		Foo Foo `alice:""`		// associated by type
//		Bar Bar `alice:"Bar"`	// associated by name
//		URL string				// not associated. Provided by creating the module.
//	}
//
//	func (m *ExampleModule) Baz() Baz {
//		return Baz{}
//	}
type BaseModule struct{}

// IsModule indicates it is a module.
func (b *BaseModule) IsModule() bool {
	return true
}
<commit_msg>Fix code format in documentation<commit_after>package alice

// Module is a marker interface for structs that defines how to initialize instances.
type Module interface {
	// IsModule indicates if this is a module.
	IsModule() bool
}

// BaseModule is an implementation of Module interface. It should be embeded into each module defined in the
// application.
//
// A typical module is defined as follows:
//
//	type ExampleModule struct {
//		alice.BaseModule
//		Foo Foo `alice:""`      // associated by type
//		Bar Bar `alice:"Bar"`   // associated by name
//		URL string              // not associated. Provided by creating the module.
//	}
//
//	func (m *ExampleModule) Baz() Baz {
//		return Baz{}
//	}
type BaseModule struct{}

// IsModule indicates it is a module.
func (b *BaseModule) IsModule() bool {
	return true
}
<|endoftext|>