<commit_before>package self_test

import (
	"context"
	"fmt"
	"io/ioutil"
	mrand "math/rand"
	"net"
	"sync"
	"sync/atomic"
	"time"

	quic "github.com/lucas-clemente/quic-go"
	quicproxy "github.com/lucas-clemente/quic-go/integrationtests/tools/proxy"
	"github.com/lucas-clemente/quic-go/internal/protocol"
	"github.com/lucas-clemente/quic-go/internal/wire"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("0-RTT", func() {
	const rtt = 50 * time.Millisecond
	for _, v := range protocol.SupportedVersions {
		version := v

		Context(fmt.Sprintf("with QUIC version %s", version), func() {
			runTest := func(ln quic.Listener, proxyPort int, testdata []byte) {
				// dial the first session in order to receive a session ticket
				go func() {
					defer GinkgoRecover()
					_, err := ln.Accept(context.Background())
					Expect(err).ToNot(HaveOccurred())
				}()

				clientConf := getTLSClientConfig()
				gets := make(chan string, 100)
				puts := make(chan string, 100)
				clientConf.ClientSessionCache = newClientSessionCache(gets, puts)
				sess, err := quic.DialAddr(
					fmt.Sprintf("localhost:%d", proxyPort),
					clientConf,
					&quic.Config{Versions: []protocol.VersionNumber{version}},
				)
				Expect(err).ToNot(HaveOccurred())
				Eventually(puts).Should(Receive())
				// received the session ticket. We're done here.
				Expect(sess.Close()).To(Succeed())

				// now dial the second session, and use 0-RTT to send some data
				done := make(chan struct{})
				go func() {
					defer GinkgoRecover()
					sess, err := ln.Accept(context.Background())
					Expect(err).ToNot(HaveOccurred())
					str, err := sess.AcceptUniStream(context.Background())
					Expect(err).ToNot(HaveOccurred())
					data, err := ioutil.ReadAll(str)
					Expect(err).ToNot(HaveOccurred())
					Expect(data).To(Equal(testdata))
					close(done)
				}()

				sess, err = quic.DialAddrEarly(
					fmt.Sprintf("localhost:%d", proxyPort),
					clientConf,
					&quic.Config{Versions: []protocol.VersionNumber{version}},
				)
				Expect(err).ToNot(HaveOccurred())
				str, err := sess.OpenUniStream()
				Expect(err).ToNot(HaveOccurred())
				_, err = str.Write(testdata)
				Expect(err).ToNot(HaveOccurred())
				Expect(str.Close()).To(Succeed())
				Eventually(done).Should(BeClosed())
			}

			It("transfers 0-RTT data", func() {
				var num0RTTPackets uint32 // to be used as an atomic

				ln, err := quic.ListenAddr(
					"localhost:0",
					getTLSConfig(),
					&quic.Config{
						Versions:    []protocol.VersionNumber{version},
						AcceptToken: func(_ net.Addr, _ *quic.Token) bool { return true },
					},
				)
				Expect(err).ToNot(HaveOccurred())
				defer ln.Close()
				serverPort := ln.Addr().(*net.UDPAddr).Port

				proxy, err := quicproxy.NewQuicProxy("localhost:0", &quicproxy.Opts{
					RemoteAddr: fmt.Sprintf("localhost:%d", serverPort),
					DelayPacket: func(_ quicproxy.Direction, data []byte) time.Duration {
						hdr, _, _, err := wire.ParsePacket(data, 0)
						Expect(err).ToNot(HaveOccurred())
						if hdr.Type == protocol.PacketType0RTT {
							atomic.AddUint32(&num0RTTPackets, 1)
						}
						return rtt / 2
					},
				})
				Expect(err).ToNot(HaveOccurred())
				defer proxy.Close()

				runTest(ln, proxy.LocalPort(), PRData)

				num0RTT := atomic.LoadUint32(&num0RTTPackets)
				fmt.Fprintf(GinkgoWriter, "Sent %d 0-RTT packets.", num0RTT)
				Expect(num0RTT).ToNot(BeZero())
			})

			It("transfers 0-RTT data, when 0-RTT packets are lost", func() {
				var (
					num0RTTPackets uint32 // to be used as an atomic
					num0RTTDropped uint32
				)

				ln, err := quic.ListenAddr(
					"localhost:0",
					getTLSConfig(),
					&quic.Config{
						Versions:    []protocol.VersionNumber{version},
						AcceptToken: func(_ net.Addr, _ *quic.Token) bool { return true },
					},
				)
				Expect(err).ToNot(HaveOccurred())
				defer ln.Close()
				serverPort := ln.Addr().(*net.UDPAddr).Port

				proxy, err := quicproxy.NewQuicProxy("localhost:0", &quicproxy.Opts{
					RemoteAddr: fmt.Sprintf("localhost:%d", serverPort),
					DelayPacket: func(_ quicproxy.Direction, data []byte) time.Duration {
						hdr, _, _, err := wire.ParsePacket(data, 0)
						Expect(err).ToNot(HaveOccurred())
						if hdr.Type == protocol.PacketType0RTT {
							atomic.AddUint32(&num0RTTPackets, 1)
						}
						return rtt / 2
					},
					DropPacket: func(_ quicproxy.Direction, data []byte) bool {
						hdr, _, _, err := wire.ParsePacket(data, 0)
						Expect(err).ToNot(HaveOccurred())
						if hdr.Type == protocol.PacketType0RTT {
							// drop 25% of the 0-RTT packets
							drop := mrand.Intn(4) == 0
							if drop {
								atomic.AddUint32(&num0RTTDropped, 1)
							}
							return drop
						}
						return false
					},
				})
				Expect(err).ToNot(HaveOccurred())
				defer proxy.Close()

				runTest(ln, proxy.LocalPort(), PRData)

				num0RTT := atomic.LoadUint32(&num0RTTPackets)
				numDropped := atomic.LoadUint32(&num0RTTDropped)
				fmt.Fprintf(GinkgoWriter, "Sent %d 0-RTT packets. Dropped %d of those.", num0RTT, numDropped)
				Expect(numDropped).ToNot(BeZero())
				Expect(num0RTT).ToNot(BeZero())
			})

			It("retransmits all 0-RTT data when the server performs a Retry", func() {
				var mutex sync.Mutex
				var firstConnID, secondConnID protocol.ConnectionID
				var firstCounter, secondCounter int

				ln, err := quic.ListenAddr(
					"localhost:0",
					getTLSConfig(),
					&quic.Config{Versions: []protocol.VersionNumber{version}},
				)
				Expect(err).ToNot(HaveOccurred())
				defer ln.Close()
				serverPort := ln.Addr().(*net.UDPAddr).Port

				proxy, err := quicproxy.NewQuicProxy("localhost:0", &quicproxy.Opts{
					RemoteAddr: fmt.Sprintf("localhost:%d", serverPort),
					DelayPacket: func(_ quicproxy.Direction, data []byte) time.Duration {
						hdr, _, _, err := wire.ParsePacket(data, 0)
						Expect(err).ToNot(HaveOccurred())
						if hdr.Type == protocol.PacketType0RTT {
							connID := hdr.DestConnectionID
							mutex.Lock()
							defer mutex.Unlock()
							if firstConnID == nil {
								firstConnID = connID
								firstCounter++
							} else if firstConnID != nil && firstConnID.Equal(connID) {
								Expect(secondConnID).To(BeNil())
								firstCounter++
							} else if secondConnID == nil {
								secondConnID = connID
								secondCounter++
							} else if secondConnID != nil && secondConnID.Equal(connID) {
								secondCounter++
							} else {
								Fail("received 3 connection IDs on 0-RTT packets")
							}
						}
						return rtt / 2
					},
				})
				Expect(err).ToNot(HaveOccurred())
				defer proxy.Close()

				runTest(ln, proxy.LocalPort(), GeneratePRData(5*1100)) // ~5 packets

				mutex.Lock()
				defer mutex.Unlock()
				Expect(firstCounter).To(BeNumerically("~", 5, 1)) // the FIN bit might be sent extra
				Expect(secondCounter).To(Equal(firstCounter))
			})
		})
	}
})
<commit_msg>add a test that waits for handshake completion after sending 0-RTT data<commit_after>package self_test

import (
	"context"
	"crypto/tls"
	"fmt"
	"io/ioutil"
	mrand "math/rand"
	"net"
	"sync"
	"sync/atomic"
	"time"

	quic "github.com/lucas-clemente/quic-go"
	quicproxy "github.com/lucas-clemente/quic-go/integrationtests/tools/proxy"
	"github.com/lucas-clemente/quic-go/internal/protocol"
	"github.com/lucas-clemente/quic-go/internal/wire"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("0-RTT", func() {
	const rtt = 50 * time.Millisecond
	for _, v := range protocol.SupportedVersions {
		version := v

		Context(fmt.Sprintf("with QUIC version %s", version), func() {
			runCountingProxy := func(serverPort int) (*quicproxy.QuicProxy, *uint32) {
				var num0RTTPackets uint32 // to be used as an atomic
				proxy, err := quicproxy.NewQuicProxy("localhost:0", &quicproxy.Opts{
					RemoteAddr: fmt.Sprintf("localhost:%d", serverPort),
					DelayPacket: func(_ quicproxy.Direction, data []byte) time.Duration {
						hdr, _, _, err := wire.ParsePacket(data, 0)
						Expect(err).ToNot(HaveOccurred())
						if hdr.Type == protocol.PacketType0RTT {
							atomic.AddUint32(&num0RTTPackets, 1)
						}
						return rtt / 2
					},
				})
				Expect(err).ToNot(HaveOccurred())

				return proxy, &num0RTTPackets
			}

			dialAndReceiveSessionTicket := func(ln quic.Listener, proxyPort int) *tls.Config {
				// dial the first session in order to receive a session ticket
				go func() {
					defer GinkgoRecover()
					_, err := ln.Accept(context.Background())
					Expect(err).ToNot(HaveOccurred())
				}()

				clientConf := getTLSClientConfig()
				gets := make(chan string, 100)
				puts := make(chan string, 100)
				clientConf.ClientSessionCache = newClientSessionCache(gets, puts)
				sess, err := quic.DialAddr(
					fmt.Sprintf("localhost:%d", proxyPort),
					clientConf,
					&quic.Config{Versions: []protocol.VersionNumber{version}},
				)
				Expect(err).ToNot(HaveOccurred())
				Eventually(puts).Should(Receive())
				// received the session ticket. We're done here.
				Expect(sess.Close()).To(Succeed())
				return clientConf
			}

			transfer0RTTData := func(ln quic.Listener, proxyPort int, clientConf *tls.Config, testdata []byte) {
				// now dial the second session, and use 0-RTT to send some data
				done := make(chan struct{})
				go func() {
					defer GinkgoRecover()
					sess, err := ln.Accept(context.Background())
					Expect(err).ToNot(HaveOccurred())
					str, err := sess.AcceptUniStream(context.Background())
					Expect(err).ToNot(HaveOccurred())
					data, err := ioutil.ReadAll(str)
					Expect(err).ToNot(HaveOccurred())
					Expect(data).To(Equal(testdata))
					close(done)
				}()

				sess, err := quic.DialAddrEarly(
					fmt.Sprintf("localhost:%d", proxyPort),
					clientConf,
					&quic.Config{Versions: []protocol.VersionNumber{version}},
				)
				Expect(err).ToNot(HaveOccurred())
				str, err := sess.OpenUniStream()
				Expect(err).ToNot(HaveOccurred())
				_, err = str.Write(testdata)
				Expect(err).ToNot(HaveOccurred())
				Expect(str.Close()).To(Succeed())
				Eventually(done).Should(BeClosed())
			}

			It("transfers 0-RTT data", func() {
				ln, err := quic.ListenAddr(
					"localhost:0",
					getTLSConfig(),
					&quic.Config{
						Versions:    []protocol.VersionNumber{version},
						AcceptToken: func(_ net.Addr, _ *quic.Token) bool { return true },
					},
				)
				Expect(err).ToNot(HaveOccurred())
				defer ln.Close()

				proxy, num0RTTPackets := runCountingProxy(ln.Addr().(*net.UDPAddr).Port)
				defer proxy.Close()

				clientConf := dialAndReceiveSessionTicket(ln, proxy.LocalPort())
				transfer0RTTData(ln, proxy.LocalPort(), clientConf, PRData)

				num0RTT := atomic.LoadUint32(num0RTTPackets)
				fmt.Fprintf(GinkgoWriter, "Sent %d 0-RTT packets.", num0RTT)
				Expect(num0RTT).ToNot(BeZero())
			})

			// Test that data intended to be sent with 1-RTT protection is not sent in 0-RTT packets.
			It("waits until a session until the handshake is done", func() {
				ln, err := quic.ListenAddr(
					"localhost:0",
					getTLSConfig(),
					&quic.Config{
						Versions:    []protocol.VersionNumber{version},
						AcceptToken: func(_ net.Addr, _ *quic.Token) bool { return true },
					},
				)
				Expect(err).ToNot(HaveOccurred())
				defer ln.Close()

				proxy, num0RTTPackets := runCountingProxy(ln.Addr().(*net.UDPAddr).Port)
				defer proxy.Close()

				clientConf := dialAndReceiveSessionTicket(ln, proxy.LocalPort())

				zeroRTTData := GeneratePRData(2 * 1100) // 2 packets
				oneRTTData := PRData

				// now dial the second session, and use 0-RTT to send some data
				done := make(chan struct{})
				go func() {
					defer GinkgoRecover()
					sess, err := ln.Accept(context.Background())
					Expect(err).ToNot(HaveOccurred())
					str, err := sess.AcceptUniStream(context.Background())
					Expect(err).ToNot(HaveOccurred())
					data, err := ioutil.ReadAll(str)
					Expect(err).ToNot(HaveOccurred())
					Expect(data).To(Equal(zeroRTTData))
					str, err = sess.AcceptUniStream(context.Background())
					Expect(err).ToNot(HaveOccurred())
					data, err = ioutil.ReadAll(str)
					Expect(err).ToNot(HaveOccurred())
					Expect(data).To(Equal(oneRTTData))
					close(done)
				}()

				sess, err := quic.DialAddrEarly(
					fmt.Sprintf("localhost:%d", proxy.LocalPort()),
					clientConf,
					&quic.Config{Versions: []protocol.VersionNumber{version}},
				)
				Expect(err).ToNot(HaveOccurred())
				sent0RTT := make(chan struct{})
				go func() {
					defer GinkgoRecover()
					defer close(sent0RTT)
					str, err := sess.OpenUniStream()
					Expect(err).ToNot(HaveOccurred())
					_, err = str.Write(zeroRTTData)
					Expect(err).ToNot(HaveOccurred())
					Expect(str.Close()).To(Succeed())
				}()
				Eventually(sent0RTT).Should(BeClosed())

				// wait for the handshake to complete
				Eventually(sess.HandshakeComplete().Done()).Should(BeClosed())
				str, err := sess.OpenUniStream()
				Expect(err).ToNot(HaveOccurred())
				_, err = str.Write(PRData)
				Expect(err).ToNot(HaveOccurred())
				Expect(str.Close()).To(Succeed())

				Eventually(done).Should(BeClosed())

				num0RTT := atomic.LoadUint32(num0RTTPackets)
				fmt.Fprintf(GinkgoWriter, "Sent %d 0-RTT packets.", num0RTT)
				Expect(num0RTT).To(Or(BeEquivalentTo(2), BeEquivalentTo(3))) // the FIN might be sent in a separate packet
			})

			It("transfers 0-RTT data, when 0-RTT packets are lost", func() {
				var (
					num0RTTPackets uint32 // to be used as an atomic
					num0RTTDropped uint32
				)

				ln, err := quic.ListenAddr(
					"localhost:0",
					getTLSConfig(),
					&quic.Config{
						Versions:    []protocol.VersionNumber{version},
						AcceptToken: func(_ net.Addr, _ *quic.Token) bool { return true },
					},
				)
				Expect(err).ToNot(HaveOccurred())
				defer ln.Close()
				serverPort := ln.Addr().(*net.UDPAddr).Port

				proxy, err := quicproxy.NewQuicProxy("localhost:0", &quicproxy.Opts{
					RemoteAddr: fmt.Sprintf("localhost:%d", serverPort),
					DelayPacket: func(_ quicproxy.Direction, data []byte) time.Duration {
						hdr, _, _, err := wire.ParsePacket(data, 0)
						Expect(err).ToNot(HaveOccurred())
						if hdr.Type == protocol.PacketType0RTT {
							atomic.AddUint32(&num0RTTPackets, 1)
						}
						return rtt / 2
					},
					DropPacket: func(_ quicproxy.Direction, data []byte) bool {
						hdr, _, _, err := wire.ParsePacket(data, 0)
						Expect(err).ToNot(HaveOccurred())
						if hdr.Type == protocol.PacketType0RTT {
							// drop 25% of the 0-RTT packets
							drop := mrand.Intn(4) == 0
							if drop {
								atomic.AddUint32(&num0RTTDropped, 1)
							}
							return drop
						}
						return false
					},
				})
				Expect(err).ToNot(HaveOccurred())
				defer proxy.Close()

				clientConf := dialAndReceiveSessionTicket(ln, proxy.LocalPort())
				transfer0RTTData(ln, proxy.LocalPort(), clientConf, PRData)

				num0RTT := atomic.LoadUint32(&num0RTTPackets)
				numDropped := atomic.LoadUint32(&num0RTTDropped)
				fmt.Fprintf(GinkgoWriter, "Sent %d 0-RTT packets. Dropped %d of those.", num0RTT, numDropped)
				Expect(numDropped).ToNot(BeZero())
				Expect(num0RTT).ToNot(BeZero())
			})

			It("retransmits all 0-RTT data when the server performs a Retry", func() {
				var mutex sync.Mutex
				var firstConnID, secondConnID protocol.ConnectionID
				var firstCounter, secondCounter int

				ln, err := quic.ListenAddr(
					"localhost:0",
					getTLSConfig(),
					&quic.Config{Versions: []protocol.VersionNumber{version}},
				)
				Expect(err).ToNot(HaveOccurred())
				defer ln.Close()
				serverPort := ln.Addr().(*net.UDPAddr).Port

				proxy, err := quicproxy.NewQuicProxy("localhost:0", &quicproxy.Opts{
					RemoteAddr: fmt.Sprintf("localhost:%d", serverPort),
					DelayPacket: func(_ quicproxy.Direction, data []byte) time.Duration {
						hdr, _, _, err := wire.ParsePacket(data, 0)
						Expect(err).ToNot(HaveOccurred())
						if hdr.Type == protocol.PacketType0RTT {
							connID := hdr.DestConnectionID
							mutex.Lock()
							defer mutex.Unlock()
							if firstConnID == nil {
								firstConnID = connID
								firstCounter++
							} else if firstConnID != nil && firstConnID.Equal(connID) {
								Expect(secondConnID).To(BeNil())
								firstCounter++
							} else if secondConnID == nil {
								secondConnID = connID
								secondCounter++
							} else if secondConnID != nil && secondConnID.Equal(connID) {
								secondCounter++
							} else {
								Fail("received 3 connection IDs on 0-RTT packets")
							}
						}
						return rtt / 2
					},
				})
				Expect(err).ToNot(HaveOccurred())
				defer proxy.Close()

				clientConf := dialAndReceiveSessionTicket(ln, proxy.LocalPort())
				transfer0RTTData(ln, proxy.LocalPort(), clientConf, GeneratePRData(5*1100)) // ~5 packets

				mutex.Lock()
				defer mutex.Unlock()
				Expect(firstCounter).To(BeNumerically("~", 5, 1)) // the FIN bit might be sent extra
				Expect(secondCounter).To(Equal(firstCounter))
			})
		})
	}
})
<|endoftext|>


<commit_before>/*
Copyright 2020 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 endpointslice

import (
	"context"
	"fmt"
	"testing"
	"time"

	corev1 "k8s.io/api/core/v1"
	discovery "k8s.io/api/discovery/v1beta1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/client-go/informers"
	clientset "k8s.io/client-go/kubernetes"
	restclient "k8s.io/client-go/rest"
	"k8s.io/kubernetes/pkg/controller/endpoint"
	"k8s.io/kubernetes/pkg/controller/endpointslice"
	"k8s.io/kubernetes/pkg/controller/endpointslicemirroring"
	"k8s.io/kubernetes/test/integration/framework"
)

func TestEndpointSliceMirroring(t *testing.T) {
	masterConfig := framework.NewIntegrationTestMasterConfig()
	_, server, closeFn := framework.RunAMaster(masterConfig)
	defer closeFn()

	config := restclient.Config{Host: server.URL}
	client, err := clientset.NewForConfig(&config)
	if err != nil {
		t.Fatalf("Error creating clientset: %v", err)
	}

	resyncPeriod := 12 * time.Hour
	informers := informers.NewSharedInformerFactory(client, resyncPeriod)

	epController := endpoint.NewEndpointController(
		informers.Core().V1().Pods(),
		informers.Core().V1().Services(),
		informers.Core().V1().Endpoints(),
		client,
		1*time.Second)

	epsController := endpointslice.NewController(
		informers.Core().V1().Pods(),
		informers.Core().V1().Services(),
		informers.Core().V1().Nodes(),
		informers.Discovery().V1beta1().EndpointSlices(),
		int32(100),
		client,
		1*time.Second)

	epsmController := endpointslicemirroring.NewController(
		informers.Core().V1().Endpoints(),
		informers.Discovery().V1beta1().EndpointSlices(),
		informers.Core().V1().Services(),
		int32(100),
		client,
		1*time.Second)

	// Start informer and controllers
	stopCh := make(chan struct{})
	defer close(stopCh)
	informers.Start(stopCh)
	go epController.Run(5, stopCh)
	go epsController.Run(5, stopCh)
	go epsmController.Run(5, stopCh)

	testCases := []struct {
		testName                     string
		service                      *corev1.Service
		endpoints                    *corev1.Endpoints
		expectEndpointSlice          bool
		expectEndpointSliceManagedBy string
	}{{
		testName: "Service with selector",
		service: &corev1.Service{
			ObjectMeta: metav1.ObjectMeta{
				Name: "test-123",
			},
			Spec: corev1.ServiceSpec{
				Ports: []corev1.ServicePort{{
					Port: int32(80),
				}},
				Selector: map[string]string{
					"foo": "bar",
				},
			},
		},
		endpoints: &corev1.Endpoints{
			ObjectMeta: metav1.ObjectMeta{
				Name: "test-123",
			},
		},
		expectEndpointSlice:          true,
		expectEndpointSliceManagedBy: "endpointslice-controller.k8s.io",
	}, {
		testName: "Service without selector",
		service: &corev1.Service{
			ObjectMeta: metav1.ObjectMeta{
				Name: "test-123",
			},
			Spec: corev1.ServiceSpec{
				Ports: []corev1.ServicePort{{
					Port: int32(80),
				}},
			},
		},
		endpoints: &corev1.Endpoints{
			ObjectMeta: metav1.ObjectMeta{
				Name: "test-123",
			},
			Subsets: []corev1.EndpointSubset{{
				Ports: []corev1.EndpointPort{{
					Port: 80,
				}},
				Addresses: []corev1.EndpointAddress{{
					IP: "10.0.0.1",
				}},
			}},
		},
		expectEndpointSlice:          true,
		expectEndpointSliceManagedBy: "endpointslicemirroring-controller.k8s.io",
	}, {
		testName: "Service without Endpoints",
		service: &corev1.Service{
			ObjectMeta: metav1.ObjectMeta{
				Name: "test-123",
			},
			Spec: corev1.ServiceSpec{
				Ports: []corev1.ServicePort{{
					Port: int32(80),
				}},
				Selector: map[string]string{
					"foo": "bar",
				},
			},
		},
		endpoints:                    nil,
		expectEndpointSlice:          true,
		expectEndpointSliceManagedBy: "endpointslice-controller.k8s.io",
	}, {
		testName: "Endpoints without Service",
		service:  nil,
		endpoints: &corev1.Endpoints{
			ObjectMeta: metav1.ObjectMeta{
				Name: "test-123",
			},
		},
		expectEndpointSlice: false,
	}}

	for i, tc := range testCases {
		t.Run(tc.testName, func(t *testing.T) {
			ns := framework.CreateTestingNamespace(fmt.Sprintf("test-endpointslice-mirroring-%d", i), server, t)
			defer framework.DeleteTestingNamespace(ns, server, t)

			resourceName := ""
			if tc.service != nil {
				resourceName = tc.service.Name
				tc.service.Namespace = ns.Name
				_, err = client.CoreV1().Services(ns.Name).Create(context.TODO(), tc.service, metav1.CreateOptions{})
				if err != nil {
					t.Fatalf("Error creating service: %v", err)
				}
			}

			if tc.endpoints != nil {
				resourceName = tc.endpoints.Name
				tc.endpoints.Namespace = ns.Name
				_, err = client.CoreV1().Endpoints(ns.Name).Create(context.TODO(), tc.endpoints, metav1.CreateOptions{})
				if err != nil {
					t.Fatalf("Error creating endpoints: %v", err)
				}
			}

			err = wait.PollImmediate(1*time.Second, wait.ForeverTestTimeout, func() (bool, error) {
				lSelector := discovery.LabelServiceName + "=" + resourceName
				esList, err := client.DiscoveryV1beta1().EndpointSlices(ns.Name).List(context.TODO(), metav1.ListOptions{LabelSelector: lSelector})
				if err != nil {
					t.Logf("Error listing EndpointSlices: %v", err)
					return false, err
				}

				if tc.expectEndpointSlice {
					if len(esList.Items) == 0 {
						t.Logf("Waiting for EndpointSlice to be created")
						return false, nil
					}
					if len(esList.Items) > 1 {
						return false, fmt.Errorf("Only expected 1 EndpointSlice, got %d", len(esList.Items))
					}
					endpointSlice := esList.Items[0]
					if tc.expectEndpointSliceManagedBy != "" {
						if endpointSlice.Labels[discovery.LabelManagedBy] != tc.expectEndpointSliceManagedBy {
							return false, fmt.Errorf("Expected EndpointSlice to be managed by %s, got %s", tc.expectEndpointSliceManagedBy, endpointSlice.Labels[discovery.LabelManagedBy])
						}
					}
				} else if len(esList.Items) > 0 {
					t.Logf("Waiting for EndpointSlices to be removed, still %d", len(esList.Items))
					return false, nil
				}

				return true, nil
			})
			if err != nil {
				t.Fatalf("Timed out waiting for conditions: %v", err)
			}
		})
	}

}
<commit_msg>Fixing a flaky EndpointSliceMirroring integration test<commit_after>/*
Copyright 2020 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 endpointslice

import (
	"context"
	"fmt"
	"testing"
	"time"

	corev1 "k8s.io/api/core/v1"
	discovery "k8s.io/api/discovery/v1beta1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/client-go/informers"
	clientset "k8s.io/client-go/kubernetes"
	restclient "k8s.io/client-go/rest"
	"k8s.io/kubernetes/pkg/controller/endpoint"
	"k8s.io/kubernetes/pkg/controller/endpointslice"
	"k8s.io/kubernetes/pkg/controller/endpointslicemirroring"
	"k8s.io/kubernetes/test/integration/framework"
)

func TestEndpointSliceMirroring(t *testing.T) {
	masterConfig := framework.NewIntegrationTestMasterConfig()
	_, server, closeFn := framework.RunAMaster(masterConfig)
	defer closeFn()

	config := restclient.Config{Host: server.URL}
	client, err := clientset.NewForConfig(&config)
	if err != nil {
		t.Fatalf("Error creating clientset: %v", err)
	}

	resyncPeriod := 12 * time.Hour
	informers := informers.NewSharedInformerFactory(client, resyncPeriod)

	epController := endpoint.NewEndpointController(
		informers.Core().V1().Pods(),
		informers.Core().V1().Services(),
		informers.Core().V1().Endpoints(),
		client,
		1*time.Second)

	epsController := endpointslice.NewController(
		informers.Core().V1().Pods(),
		informers.Core().V1().Services(),
		informers.Core().V1().Nodes(),
		informers.Discovery().V1beta1().EndpointSlices(),
		int32(100),
		client,
		1*time.Second)

	epsmController := endpointslicemirroring.NewController(
		informers.Core().V1().Endpoints(),
		informers.Discovery().V1beta1().EndpointSlices(),
		informers.Core().V1().Services(),
		int32(100),
		client,
		1*time.Second)

	// Start informer and controllers
	stopCh := make(chan struct{})
	defer close(stopCh)
	informers.Start(stopCh)
	go epController.Run(5, stopCh)
	go epsController.Run(5, stopCh)
	go epsmController.Run(5, stopCh)

	testCases := []struct {
		testName                     string
		service                      *corev1.Service
		customEndpoints              *corev1.Endpoints
		expectEndpointSlice          bool
		expectEndpointSliceManagedBy string
	}{{
		testName: "Service with selector",
		service: &corev1.Service{
			ObjectMeta: metav1.ObjectMeta{
				Name: "test-123",
			},
			Spec: corev1.ServiceSpec{
				Ports: []corev1.ServicePort{{
					Port: int32(80),
				}},
				Selector: map[string]string{
					"foo": "bar",
				},
			},
		},
		expectEndpointSlice:          true,
		expectEndpointSliceManagedBy: "endpointslice-controller.k8s.io",
	}, {
		testName: "Service without selector",
		service: &corev1.Service{
			ObjectMeta: metav1.ObjectMeta{
				Name: "test-123",
			},
			Spec: corev1.ServiceSpec{
				Ports: []corev1.ServicePort{{
					Port: int32(80),
				}},
			},
		},
		customEndpoints: &corev1.Endpoints{
			ObjectMeta: metav1.ObjectMeta{
				Name: "test-123",
			},
			Subsets: []corev1.EndpointSubset{{
				Ports: []corev1.EndpointPort{{
					Port: 80,
				}},
				Addresses: []corev1.EndpointAddress{{
					IP: "10.0.0.1",
				}},
			}},
		},
		expectEndpointSlice:          true,
		expectEndpointSliceManagedBy: "endpointslicemirroring-controller.k8s.io",
	}, {
		testName: "Service without Endpoints",
		service: &corev1.Service{
			ObjectMeta: metav1.ObjectMeta{
				Name: "test-123",
			},
			Spec: corev1.ServiceSpec{
				Ports: []corev1.ServicePort{{
					Port: int32(80),
				}},
				Selector: map[string]string{
					"foo": "bar",
				},
			},
		},
		customEndpoints:              nil,
		expectEndpointSlice:          true,
		expectEndpointSliceManagedBy: "endpointslice-controller.k8s.io",
	}, {
		testName: "Endpoints without Service",
		service:  nil,
		customEndpoints: &corev1.Endpoints{
			ObjectMeta: metav1.ObjectMeta{
				Name: "test-123",
			},
		},
		expectEndpointSlice: false,
	}}

	for i, tc := range testCases {
		t.Run(tc.testName, func(t *testing.T) {
			ns := framework.CreateTestingNamespace(fmt.Sprintf("test-endpointslice-mirroring-%d", i), server, t)
			defer framework.DeleteTestingNamespace(ns, server, t)

			resourceName := ""
			if tc.service != nil {
				resourceName = tc.service.Name
				tc.service.Namespace = ns.Name
				_, err = client.CoreV1().Services(ns.Name).Create(context.TODO(), tc.service, metav1.CreateOptions{})
				if err != nil {
					t.Fatalf("Error creating service: %v", err)
				}
			}

			if tc.customEndpoints != nil {
				resourceName = tc.customEndpoints.Name
				tc.customEndpoints.Namespace = ns.Name
				_, err = client.CoreV1().Endpoints(ns.Name).Create(context.TODO(), tc.customEndpoints, metav1.CreateOptions{})
				if err != nil {
					t.Fatalf("Error creating endpoints: %v", err)
				}
			}

			err = wait.PollImmediate(1*time.Second, wait.ForeverTestTimeout, func() (bool, error) {
				lSelector := discovery.LabelServiceName + "=" + resourceName
				esList, err := client.DiscoveryV1beta1().EndpointSlices(ns.Name).List(context.TODO(), metav1.ListOptions{LabelSelector: lSelector})
				if err != nil {
					t.Logf("Error listing EndpointSlices: %v", err)
					return false, err
				}

				if tc.expectEndpointSlice {
					if len(esList.Items) == 0 {
						t.Logf("Waiting for EndpointSlice to be created")
						return false, nil
					}
					if len(esList.Items) > 1 {
						return false, fmt.Errorf("Only expected 1 EndpointSlice, got %d", len(esList.Items))
					}
					endpointSlice := esList.Items[0]
					if tc.expectEndpointSliceManagedBy != "" {
						if endpointSlice.Labels[discovery.LabelManagedBy] != tc.expectEndpointSliceManagedBy {
							return false, fmt.Errorf("Expected EndpointSlice to be managed by %s, got %s", tc.expectEndpointSliceManagedBy, endpointSlice.Labels[discovery.LabelManagedBy])
						}
					}
				} else if len(esList.Items) > 0 {
					t.Logf("Waiting for EndpointSlices to be removed, still %d", len(esList.Items))
					return false, nil
				}

				return true, nil
			})
			if err != nil {
				t.Fatalf("Timed out waiting for conditions: %v", err)
			}
		})
	}

}
<|endoftext|>


<commit_before>//  Copyright 2016 The goscope 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 main

import (
	"flag"
	"fmt"
	"log"
	"sort"
	"strings"

	"github.com/zagrodzki/goscope/dummy"
	"github.com/zagrodzki/goscope/scope"
	"github.com/zagrodzki/goscope/usb"
)

var (
	dev       = flag.String("device", "", "Device to use, autodetect if empty")
	list      = flag.Bool("list", false, "If set, only list available devices")
	rate      = flag.Int("rate", 1e6, "sampling rate (per second)")
	voltRange = flag.Float64("range", 5.0, "sensitivity of the device (applied to all channels)")
	chID      = flag.String("chan", "", "name of the channel to use. If not specified, use the first channel")
	period    = flag.Duration("period", 0, "how long period of samples to collect, run forever if set to 0")
	showHist  = flag.Bool("histogram", false, "If true, output histogram of samples, otherwise only the mode")
)

func must(e error) {
	if e != nil {
		log.Fatalf(e.Error())
	}
}

type system struct {
	name      string
	enumerate func() map[string]string
	open      func(string) (scope.Device, error)
}

var (
	systems = []system{
		{
			name:      "dummy",
			enumerate: dummy.Enumerate,
			open:      dummy.Open,
		},
		{
			name:      "usb",
			enumerate: usb.Enumerate,
			open:      usb.Open,
		},
	}
	systemsByName = make(map[string]int)
)

type orderedHist struct {
	s map[scope.Sample]int
	k []scope.Sample
}

func (o *orderedHist) Len() int {
	return len(o.k)
}
func (o *orderedHist) Swap(i, j int) {
	o.k[i], o.k[j] = o.k[j], o.k[i]
}
func (o *orderedHist) Less(i, j int) bool {
	// sorting in reverse order
	return o.s[o.k[i]] > o.s[o.k[j]]
}
func (o *orderedHist) sort() {
	if len(o.k) != len(o.s) {
		o.k = make(scope.Sample, len(o.s))
		for s := range o.s {
			o.k = append(o.k, s)
		}
	}
	sort.Sort(o)
}

func main() {
	flag.Parse()
	var all []string
	for idx, sys := range systems {
		systemsByName[sys.name] = idx
		for id := range sys.enumerate() {
			all = append(all, fmt.Sprintf("%s:%s", sys.name, id))
		}
	}
	if len(all) == 0 {
		log.Fatalf("Did not find any supported devices")
	}
	if *list {
		fmt.Println("Devices found:")
		for _, d := range all {
			fmt.Println(d)
		}
		return
	}
	id := all[0]
	if *dev != "" {
		for _, d := range all {
			if d == *dev {
				id = d
				break
			}
		}
		if id != *dev {
			log.Fatalf("Device %s not detected on the list. Available devices: %v", *dev, all)
		}
	} else if len(all) > 1 {
		log.Printf("Multiple devices found: %v", all)
		log.Printf("Using the first device (%s)", id)
	}
	parts := strings.SplitN(id, ":", 2)
	s := systemsByName[parts[0]]
	osc, err := systems[s].open(parts[1])
	if err != nil {
		log.Fatalf("Open: %+v", err)
	}
	fmt.Println(osc)
	channels := osc.Channels()
	ch := channels[0]
	if *chID != "" {
		for _, c := range channels {
			if c == scope.ChanID(*chID) {
				ch = c
			}
		}
		if ch != scope.ChanID(*chID) {
			log.Fatalf("Device %s does not have a channel %q. Available channels: %v", id, *chID, channels)
		}
	}
	must(osc.Channel(ch).SetVoltRange(scope.VoltRange(*voltRange)))
	data, stop, err := osc.StartSampling()
	if err != nil {
		log.Fatalf("ReadData: %+v", err)
	}
	defer stop()
	rate := osc.GetSampleRate()
	log.Printf("Sampling rate %s (interval %s)", rate, rate.Interval())
	i := int(scope.DurationFromNano(*period) / rate.Interval())
	log.Printf("Reading %d samples", i)
	for s := range data {
		hist := &orderedHist{
			s: make(map[scope.Sample]int),
		}
		for _, d := range s.Samples[ch] {
			hist.s[d]++
		}
		hist.sort()
		if *showHist {
			out := make([]string, len(hist.k))
			for k := range hist.k {
				out[k] = fmt.Sprintf("%f: %d", hist.k[k], hist.s[hist.k[k]])
			}
			fmt.Println(out)
		} else {
			fmt.Println(hist.k[0])
		}
		i -= len(s.Samples[ch])
		if *period != 0 && i <= 0 {
			break
		}
	}
}
<commit_msg>* typo, make a slice, not a struct.<commit_after>//  Copyright 2016 The goscope 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 main

import (
	"flag"
	"fmt"
	"log"
	"sort"
	"strings"

	"github.com/zagrodzki/goscope/dummy"
	"github.com/zagrodzki/goscope/scope"
	"github.com/zagrodzki/goscope/usb"
)

var (
	dev       = flag.String("device", "", "Device to use, autodetect if empty")
	list      = flag.Bool("list", false, "If set, only list available devices")
	rate      = flag.Int("rate", 1e6, "sampling rate (per second)")
	voltRange = flag.Float64("range", 5.0, "sensitivity of the device (applied to all channels)")
	chID      = flag.String("chan", "", "name of the channel to use. If not specified, use the first channel")
	period    = flag.Duration("period", 0, "how long period of samples to collect, run forever if set to 0")
	showHist  = flag.Bool("histogram", false, "If true, output histogram of samples, otherwise only the mode")
)

func must(e error) {
	if e != nil {
		log.Fatalf(e.Error())
	}
}

type system struct {
	name      string
	enumerate func() map[string]string
	open      func(string) (scope.Device, error)
}

var (
	systems = []system{
		{
			name:      "dummy",
			enumerate: dummy.Enumerate,
			open:      dummy.Open,
		},
		{
			name:      "usb",
			enumerate: usb.Enumerate,
			open:      usb.Open,
		},
	}
	systemsByName = make(map[string]int)
)

type orderedHist struct {
	s map[scope.Sample]int
	k []scope.Sample
}

func (o *orderedHist) Len() int {
	return len(o.k)
}
func (o *orderedHist) Swap(i, j int) {
	o.k[i], o.k[j] = o.k[j], o.k[i]
}
func (o *orderedHist) Less(i, j int) bool {
	// sorting in reverse order
	return o.s[o.k[i]] > o.s[o.k[j]]
}
func (o *orderedHist) sort() {
	if len(o.k) != len(o.s) {
		o.k = make([]scope.Sample, len(o.s))
		for s := range o.s {
			o.k = append(o.k, s)
		}
	}
	sort.Sort(o)
}

func main() {
	flag.Parse()
	var all []string
	for idx, sys := range systems {
		systemsByName[sys.name] = idx
		for id := range sys.enumerate() {
			all = append(all, fmt.Sprintf("%s:%s", sys.name, id))
		}
	}
	if len(all) == 0 {
		log.Fatalf("Did not find any supported devices")
	}
	if *list {
		fmt.Println("Devices found:")
		for _, d := range all {
			fmt.Println(d)
		}
		return
	}
	id := all[0]
	if *dev != "" {
		for _, d := range all {
			if d == *dev {
				id = d
				break
			}
		}
		if id != *dev {
			log.Fatalf("Device %s not detected on the list. Available devices: %v", *dev, all)
		}
	} else if len(all) > 1 {
		log.Printf("Multiple devices found: %v", all)
		log.Printf("Using the first device (%s)", id)
	}
	parts := strings.SplitN(id, ":", 2)
	s := systemsByName[parts[0]]
	osc, err := systems[s].open(parts[1])
	if err != nil {
		log.Fatalf("Open: %+v", err)
	}
	fmt.Println(osc)
	channels := osc.Channels()
	ch := channels[0]
	if *chID != "" {
		for _, c := range channels {
			if c == scope.ChanID(*chID) {
				ch = c
			}
		}
		if ch != scope.ChanID(*chID) {
			log.Fatalf("Device %s does not have a channel %q. Available channels: %v", id, *chID, channels)
		}
	}
	must(osc.Channel(ch).SetVoltRange(scope.VoltRange(*voltRange)))
	data, stop, err := osc.StartSampling()
	if err != nil {
		log.Fatalf("ReadData: %+v", err)
	}
	defer stop()
	rate := osc.GetSampleRate()
	log.Printf("Sampling rate %s (interval %s)", rate, rate.Interval())
	i := int(scope.DurationFromNano(*period) / rate.Interval())
	log.Printf("Reading %d samples", i)
	for s := range data {
		hist := &orderedHist{
			s: make(map[scope.Sample]int),
		}
		for _, d := range s.Samples[ch] {
			hist.s[d]++
		}
		hist.sort()
		if *showHist {
			out := make([]string, len(hist.k))
			for k := range hist.k {
				out[k] = fmt.Sprintf("%f: %d", hist.k[k], hist.s[hist.k[k]])
			}
			fmt.Println(out)
		} else {
			fmt.Println(hist.k[0])
		}
		i -= len(s.Samples[ch])
		if *period != 0 && i <= 0 {
			break
		}
	}
}
<|endoftext|>


<commit_before>package brokers

import (
	"encoding/json"
	"errors"
	"fmt"
	"strings"
	"time"

	"sync"

	"github.com/RichardKnop/machinery/v1/config"
	"github.com/RichardKnop/machinery/v1/log"
	"github.com/RichardKnop/machinery/v1/tasks"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/sqs"
	"github.com/aws/aws-sdk-go/service/sqs/sqsiface"
)

// AWSSQSBroker represents a AWS SQS broker
type AWSSQSBroker struct {
	Broker
	processingWG      sync.WaitGroup // use wait group to make sure task processing completes on interrupt signal
	receivingWG       sync.WaitGroup
	stopReceivingChan chan int
	sess              *session.Session
	service           sqsiface.SQSAPI
}

// NewAWSSQSBroker creates new Broker instance
func NewAWSSQSBroker(cnf *config.Config) Interface {
	b := &AWSSQSBroker{Broker: New(cnf)}
	b.sess = session.Must(session.NewSessionWithOptions(session.Options{
		SharedConfigState: session.SharedConfigEnable,
	}))
	b.service = sqs.New(b.sess)

	return b
}

// GetPendingTasks returns a slice of task.Signatures waiting in the queue
func (b *AWSSQSBroker) GetPendingTasks(queue string) ([]*tasks.Signature, error) {
	return nil, errors.New("Not implemented")
}

// StartConsuming enters a loop and waits for incoming messages
func (b *AWSSQSBroker) StartConsuming(consumerTag string, concurrency int, taskProcessor TaskProcessor) (bool, error) {
	b.startConsuming(consumerTag, taskProcessor)
	qURL := b.defaultQueueURL()
	deliveries := make(chan *sqs.ReceiveMessageOutput)

	b.stopReceivingChan = make(chan int)
	b.receivingWG.Add(1)

	go func() {
		defer b.receivingWG.Done()

		log.INFO.Print("[*] Waiting for messages. To exit press CTRL+C")

		for {
			select {
			// A way to stop this goroutine from b.StopConsuming
			case <-b.stopReceivingChan:
				return
			default:
				output, err := b.receiveMessage(qURL)
				if err != nil {
					log.ERROR.Printf("Queue consume error: %s", err)
					continue
				}
				if len(output.Messages) == 0 {
					continue
				}

				deliveries <- output
			}

			whetherContinue, err := b.continueReceivingMessages(qURL, deliveries)
			if err != nil {
				log.ERROR.Printf("Error when receiving messages. Error: %v", err)
			}
			if whetherContinue == false {
				return
			}
		}
	}()

	if err := b.consume(deliveries, concurrency, taskProcessor); err != nil {
		return b.retry, err
	}

	return b.retry, nil
}

// StopConsuming quits the loop
func (b *AWSSQSBroker) StopConsuming() {
	b.stopConsuming()

	b.stopReceiving()

	// Waiting for any tasks being processed to finish
	b.processingWG.Wait()

	// Waiting for the receiving goroutine to have stopped
	b.receivingWG.Wait()
}

// Publish places a new message on the default queue
func (b *AWSSQSBroker) Publish(signature *tasks.Signature) error {

	msg, err := json.Marshal(signature)
	if err != nil {
		return fmt.Errorf("JSON marshal error: %s", err)
	}

	// Check that signature.RoutingKey is set, if not switch to DefaultQueue
	b.AdjustRoutingKey(signature)

	MsgInput := &sqs.SendMessageInput{
		MessageBody: aws.String(string(msg)),
		QueueUrl:    aws.String(b.cnf.Broker + "/" + signature.RoutingKey),
	}

	// if this is a fifo queue, there needs to be some addtional parameters.
	if strings.HasSuffix(signature.RoutingKey, ".fifo") {
		// Use Machinery's signature Task UUID as SQS Message Group ID.
		MsgDedupID := signature.UUID
		MsgInput.MessageDeduplicationId = aws.String(MsgDedupID)

		// Use Machinery's signature Group UUID as SQS Message Group ID.
		MsgGroupID := signature.GroupUUID
		MsgInput.MessageGroupId = aws.String(MsgGroupID)
	}

	// Check the ETA signature field, if it is set and it is in the future,
	// and is not a fifo queue, set a delay in seconds for the task.
	if signature.ETA != nil && !strings.HasSuffix(signature.RoutingKey, ".fifo") {
		now := time.Now().UTC()

		if signature.ETA.After(now) {
			MsgInput.DelaySeconds = aws.Int64(signature.ETA.Unix() - now.Unix())
		}
	}

	result, err := b.service.SendMessage(MsgInput)

	if err != nil {
		log.INFO.Println("Error", err)
		return err

	}
	log.INFO.Println("Success", *result.MessageId)
	return nil

}

func (b *AWSSQSBroker) consume(deliveries <-chan *sqs.ReceiveMessageOutput, concurrency int, taskProcessor TaskProcessor) error {
	pool := make(chan struct{}, concurrency)

	// initialize worker pool with maxWorkers workers
	go func() {
		b.initializePool(pool, concurrency)
	}()

	errorsChan := make(chan error)

	for {
		err := b.consumeDeliveries(deliveries, concurrency, taskProcessor, pool, errorsChan)
		if err != nil {
			return err
		}
	}
}

func (b *AWSSQSBroker) consumeOne(delivery *sqs.ReceiveMessageOutput, taskProcessor TaskProcessor) error {
	sig := new(tasks.Signature)
	if len(delivery.Messages) == 0 {
		log.ERROR.Printf("received an empty message, the delivery was %v", delivery)
		return errors.New("received empty message, the delivery is " + delivery.GoString())
	}
	msg := delivery.Messages[0].Body
	if err := json.Unmarshal([]byte(*msg), sig); err != nil {
		log.ERROR.Printf("unmarshal error. the delivery is %v", delivery)
		return err
	}

	// Delete message after consuming successfully
	err := b.deleteOne(delivery)
	if err != nil {
		log.ERROR.Printf("error when deleting the delivery. the delivery is %v", delivery)
	}

	// If the task is not registered, we requeue it,
	// there might be different workers for processing specific tasks
	if !b.IsTaskRegistered(sig.Name) {
		err := b.Publish(sig)
		if err != nil {
			return err
		}
		log.INFO.Printf("requeue a task to default queue: %v", sig)
		return nil
	}
	return taskProcessor.Process(sig)
}

func (b *AWSSQSBroker) deleteOne(delivery *sqs.ReceiveMessageOutput) error {
	qURL := b.defaultQueueURL()
	_, err := b.service.DeleteMessage(&sqs.DeleteMessageInput{
		QueueUrl:      qURL,
		ReceiptHandle: delivery.Messages[0].ReceiptHandle,
	})

	if err != nil {
		return err
	}
	return nil
}

func (b *AWSSQSBroker) defaultQueueURL() *string {
	return aws.String(b.cnf.Broker + "/" + b.cnf.DefaultQueue)
}

func (b *AWSSQSBroker) receiveMessage(qURL *string) (*sqs.ReceiveMessageOutput, error) {
	result, err := b.service.ReceiveMessage(&sqs.ReceiveMessageInput{
		AttributeNames: []*string{
			aws.String(sqs.MessageSystemAttributeNameSentTimestamp),
		},
		MessageAttributeNames: []*string{
			aws.String(sqs.QueueAttributeNameAll),
		},
		QueueUrl:            qURL,
		MaxNumberOfMessages: aws.Int64(1),
		VisibilityTimeout:   aws.Int64(int64(b.cnf.ResultsExpireIn)), // 10 hours
		WaitTimeSeconds:     aws.Int64(0),
	})
	if err != nil {
		return nil, err
	}
	return result, err
}

func (b *AWSSQSBroker) initializePool(pool chan struct{}, concurrency int) {
	for i := 0; i < concurrency; i++ {
		pool <- struct{}{}
	}
}

func (b *AWSSQSBroker) consumeDeliveries(deliveries <-chan *sqs.ReceiveMessageOutput, concurrency int, taskProcessor TaskProcessor, pool chan struct{}, errorsChan chan error) error {
	select {
	case err := <-errorsChan:
		return err
	case d := <-deliveries:
		if concurrency > 0 {
			// get worker from pool (blocks until one is available)
			<-pool
		}

		b.processingWG.Add(1)

		// Consume the task inside a goroutine so multiple tasks
		// can be processed concurrently
		go func() {

			if err := b.consumeOne(d, taskProcessor); err != nil {
				errorsChan <- err
			}

			b.processingWG.Done()

			if concurrency > 0 {
				// give worker back to pool
				pool <- struct{}{}
			}
		}()
	case <-b.stopChan:
		return nil
	}
	return nil
}

func (b *AWSSQSBroker) continueReceivingMessages(qURL *string, deliveries chan *sqs.ReceiveMessageOutput) (bool, error) {
	select {
	// A way to stop this goroutine from b.StopConsuming
	case <-b.stopReceivingChan:
		return false, nil
	default:
		output, err := b.receiveMessage(qURL)
		if err != nil {
			return true, err
		}
		if len(output.Messages) == 0 {
			return true, nil
		}
		go func() { deliveries <- output }()
	}
	return true, nil
}

func (b *AWSSQSBroker) stopReceiving() {
	// Stop the receiving goroutine
	b.stopReceivingChan <- 1
}
<commit_msg>fix: resolve issue with sqs routing check<commit_after>package brokers

import (
	"encoding/json"
	"errors"
	"fmt"
	"strings"
	"time"

	"sync"

	"github.com/RichardKnop/machinery/v1/config"
	"github.com/RichardKnop/machinery/v1/log"
	"github.com/RichardKnop/machinery/v1/tasks"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/sqs"
	"github.com/aws/aws-sdk-go/service/sqs/sqsiface"
)

// AWSSQSBroker represents a AWS SQS broker
type AWSSQSBroker struct {
	Broker
	processingWG      sync.WaitGroup // use wait group to make sure task processing completes on interrupt signal
	receivingWG       sync.WaitGroup
	stopReceivingChan chan int
	sess              *session.Session
	service           sqsiface.SQSAPI
}

// NewAWSSQSBroker creates new Broker instance
func NewAWSSQSBroker(cnf *config.Config) Interface {
	b := &AWSSQSBroker{Broker: New(cnf)}
	b.sess = session.Must(session.NewSessionWithOptions(session.Options{
		SharedConfigState: session.SharedConfigEnable,
	}))
	b.service = sqs.New(b.sess)

	return b
}

// GetPendingTasks returns a slice of task.Signatures waiting in the queue
func (b *AWSSQSBroker) GetPendingTasks(queue string) ([]*tasks.Signature, error) {
	return nil, errors.New("Not implemented")
}

// StartConsuming enters a loop and waits for incoming messages
func (b *AWSSQSBroker) StartConsuming(consumerTag string, concurrency int, taskProcessor TaskProcessor) (bool, error) {
	b.startConsuming(consumerTag, taskProcessor)
	qURL := b.defaultQueueURL()
	deliveries := make(chan *sqs.ReceiveMessageOutput)

	b.stopReceivingChan = make(chan int)
	b.receivingWG.Add(1)

	go func() {
		defer b.receivingWG.Done()

		log.INFO.Print("[*] Waiting for messages. To exit press CTRL+C")

		for {
			select {
			// A way to stop this goroutine from b.StopConsuming
			case <-b.stopReceivingChan:
				return
			default:
				output, err := b.receiveMessage(qURL)
				if err != nil {
					log.ERROR.Printf("Queue consume error: %s", err)
					continue
				}
				if len(output.Messages) == 0 {
					continue
				}

				deliveries <- output
			}

			whetherContinue, err := b.continueReceivingMessages(qURL, deliveries)
			if err != nil {
				log.ERROR.Printf("Error when receiving messages. Error: %v", err)
			}
			if whetherContinue == false {
				return
			}
		}
	}()

	if err := b.consume(deliveries, concurrency, taskProcessor); err != nil {
		return b.retry, err
	}

	return b.retry, nil
}

// StopConsuming quits the loop
func (b *AWSSQSBroker) StopConsuming() {
	b.stopConsuming()

	b.stopReceiving()

	// Waiting for any tasks being processed to finish
	b.processingWG.Wait()

	// Waiting for the receiving goroutine to have stopped
	b.receivingWG.Wait()
}

// Publish places a new message on the default queue
func (b *AWSSQSBroker) Publish(signature *tasks.Signature) error {

	msg, err := json.Marshal(signature)
	if err != nil {
		return fmt.Errorf("JSON marshal error: %s", err)
	}

	// Check that signature.RoutingKey is set, if not switch to DefaultQueue
	AdjustRoutingKey(b, signature)

	MsgInput := &sqs.SendMessageInput{
		MessageBody: aws.String(string(msg)),
		QueueUrl:    aws.String(b.cnf.Broker + "/" + signature.RoutingKey),
	}

	// if this is a fifo queue, there needs to be some addtional parameters.
	if strings.HasSuffix(signature.RoutingKey, ".fifo") {
		// Use Machinery's signature Task UUID as SQS Message Group ID.
		MsgDedupID := signature.UUID
		MsgInput.MessageDeduplicationId = aws.String(MsgDedupID)

		// Use Machinery's signature Group UUID as SQS Message Group ID.
		MsgGroupID := signature.GroupUUID
		MsgInput.MessageGroupId = aws.String(MsgGroupID)
	}

	// Check the ETA signature field, if it is set and it is in the future,
	// and is not a fifo queue, set a delay in seconds for the task.
	if signature.ETA != nil && !strings.HasSuffix(signature.RoutingKey, ".fifo") {
		now := time.Now().UTC()

		if signature.ETA.After(now) {
			MsgInput.DelaySeconds = aws.Int64(signature.ETA.Unix() - now.Unix())
		}
	}

	result, err := b.service.SendMessage(MsgInput)

	if err != nil {
		log.INFO.Println("Error", err)
		return err

	}
	log.INFO.Println("Success", *result.MessageId)
	return nil

}

func (b *AWSSQSBroker) consume(deliveries <-chan *sqs.ReceiveMessageOutput, concurrency int, taskProcessor TaskProcessor) error {
	pool := make(chan struct{}, concurrency)

	// initialize worker pool with maxWorkers workers
	go func() {
		b.initializePool(pool, concurrency)
	}()

	errorsChan := make(chan error)

	for {
		err := b.consumeDeliveries(deliveries, concurrency, taskProcessor, pool, errorsChan)
		if err != nil {
			return err
		}
	}
}

func (b *AWSSQSBroker) consumeOne(delivery *sqs.ReceiveMessageOutput, taskProcessor TaskProcessor) error {
	sig := new(tasks.Signature)
	if len(delivery.Messages) == 0 {
		log.ERROR.Printf("received an empty message, the delivery was %v", delivery)
		return errors.New("received empty message, the delivery is " + delivery.GoString())
	}
	msg := delivery.Messages[0].Body
	if err := json.Unmarshal([]byte(*msg), sig); err != nil {
		log.ERROR.Printf("unmarshal error. the delivery is %v", delivery)
		return err
	}

	// Delete message after consuming successfully
	err := b.deleteOne(delivery)
	if err != nil {
		log.ERROR.Printf("error when deleting the delivery. the delivery is %v", delivery)
	}

	// If the task is not registered, we requeue it,
	// there might be different workers for processing specific tasks
	if !b.IsTaskRegistered(sig.Name) {
		err := b.Publish(sig)
		if err != nil {
			return err
		}
		log.INFO.Printf("requeue a task to default queue: %v", sig)
		return nil
	}
	return taskProcessor.Process(sig)
}

func (b *AWSSQSBroker) deleteOne(delivery *sqs.ReceiveMessageOutput) error {
	qURL := b.defaultQueueURL()
	_, err := b.service.DeleteMessage(&sqs.DeleteMessageInput{
		QueueUrl:      qURL,
		ReceiptHandle: delivery.Messages[0].ReceiptHandle,
	})

	if err != nil {
		return err
	}
	return nil
}

func (b *AWSSQSBroker) defaultQueueURL() *string {
	return aws.String(b.cnf.Broker + "/" + b.cnf.DefaultQueue)
}

func (b *AWSSQSBroker) receiveMessage(qURL *string) (*sqs.ReceiveMessageOutput, error) {
	result, err := b.service.ReceiveMessage(&sqs.ReceiveMessageInput{
		AttributeNames: []*string{
			aws.String(sqs.MessageSystemAttributeNameSentTimestamp),
		},
		MessageAttributeNames: []*string{
			aws.String(sqs.QueueAttributeNameAll),
		},
		QueueUrl:            qURL,
		MaxNumberOfMessages: aws.Int64(1),
		VisibilityTimeout:   aws.Int64(int64(b.cnf.ResultsExpireIn)), // 10 hours
		WaitTimeSeconds:     aws.Int64(0),
	})
	if err != nil {
		return nil, err
	}
	return result, err
}

func (b *AWSSQSBroker) initializePool(pool chan struct{}, concurrency int) {
	for i := 0; i < concurrency; i++ {
		pool <- struct{}{}
	}
}

func (b *AWSSQSBroker) consumeDeliveries(deliveries <-chan *sqs.ReceiveMessageOutput, concurrency int, taskProcessor TaskProcessor, pool chan struct{}, errorsChan chan error) error {
	select {
	case err := <-errorsChan:
		return err
	case d := <-deliveries:
		if concurrency > 0 {
			// get worker from pool (blocks until one is available)
			<-pool
		}

		b.processingWG.Add(1)

		// Consume the task inside a goroutine so multiple tasks
		// can be processed concurrently
		go func() {

			if err := b.consumeOne(d, taskProcessor); err != nil {
				errorsChan <- err
			}

			b.processingWG.Done()

			if concurrency > 0 {
				// give worker back to pool
				pool <- struct{}{}
			}
		}()
	case <-b.stopChan:
		return nil
	}
	return nil
}

func (b *AWSSQSBroker) continueReceivingMessages(qURL *string, deliveries chan *sqs.ReceiveMessageOutput) (bool, error) {
	select {
	// A way to stop this goroutine from b.StopConsuming
	case <-b.stopReceivingChan:
		return false, nil
	default:
		output, err := b.receiveMessage(qURL)
		if err != nil {
			return true, err
		}
		if len(output.Messages) == 0 {
			return true, nil
		}
		go func() { deliveries <- output }()
	}
	return true, nil
}

func (b *AWSSQSBroker) stopReceiving() {
	// Stop the receiving goroutine
	b.stopReceivingChan <- 1
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"strings"
)

const gravizoBegin string = "(http://g.gravizo.com/svg?"

//const gravizoEnd string = "enduml)"

var encoder = strings.NewReplacer(";", "%3B", " ", "%20", "\n", "%0A", "@", "%40",
	"(", "%28", ")", "%29", "*", "%2A", "\\", "%5C")
var decoder = strings.NewReplacer("%3B", ";", "%20", " ", "%0A", "\n", "%40", "@",
	"%2A", "*", "%5C", "\\")

func check(err error) {
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
}

func findMatchingClose(text string, opening rune, closing rune) int {
	openingCount := 1
	for i, ch := range text {
		if ch == opening {
			openingCount++
		} else if ch == closing {
			openingCount--
			if openingCount == 0 {
				return i
			}
		}
	}
	return 0
}

func convert(filename string, replacer *strings.Replacer, backup bool) {
	buffer, err := ioutil.ReadFile(filename)
	check(err)

	if backup {
		err = ioutil.WriteFile(fmt.Sprint(filename+".bak"), buffer, 0644)
		check(err)
	}

	text := string(buffer)

	for offset, slice := range strings.Split(text, gravizoBegin) {
		if offset == 0 {
			continue
		}
		closeOffset := findMatchingClose(slice, '(', ')')
		if closeOffset > 0 {
			gravizoText := slice[:closeOffset]
			if len(gravizoText) > 0 {
				convertedText := replacer.Replace(gravizoText)
				text = strings.Replace(text, gravizoText, convertedText, -1)
			}
		}
	}

	err = ioutil.WriteFile(filename, []byte(text), 0644)
	check(err)
}

func main() {
	encode := flag.String("e", "", "Encode the given GitHub Markdown file")
	decode := flag.String("d", "", "Decode the given GitHub Markdown file")
	backup := flag.Bool("b", true, "Backup GitHub Markdown file before encode/decode")

	flag.Parse()

	if len(*encode) > 0 {
		convert(*encode, encoder, *backup)
	} else if len(*decode) > 0 {
		convert(*decode, decoder, *backup)
	}
}
<commit_msg>remove commented code<commit_after>package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"strings"
)

const gravizoBegin string = "(http://g.gravizo.com/svg?"

var encoder = strings.NewReplacer(";", "%3B", " ", "%20", "\n", "%0A", "@", "%40",
	"(", "%28", ")", "%29", "*", "%2A", "\\", "%5C")
var decoder = strings.NewReplacer("%3B", ";", "%20", " ", "%0A", "\n", "%40", "@",
	"%2A", "*", "%5C", "\\")

func check(err error) {
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
}

func findMatchingClose(text string, opening rune, closing rune) int {
	openingCount := 1
	for i, ch := range text {
		if ch == opening {
			openingCount++
		} else if ch == closing {
			openingCount--
			if openingCount == 0 {
				return i
			}
		}
	}
	return 0
}

func convert(filename string, replacer *strings.Replacer, backup bool) {
	buffer, err := ioutil.ReadFile(filename)
	check(err)

	if backup {
		err = ioutil.WriteFile(fmt.Sprint(filename+".bak"), buffer, 0644)
		check(err)
	}

	text := string(buffer)

	for offset, slice := range strings.Split(text, gravizoBegin) {
		if offset == 0 {
			continue
		}
		closeOffset := findMatchingClose(slice, '(', ')')
		if closeOffset > 0 {
			gravizoText := slice[:closeOffset]
			if len(gravizoText) > 0 {
				convertedText := replacer.Replace(gravizoText)
				text = strings.Replace(text, gravizoText, convertedText, -1)
			}
		}
	}

	err = ioutil.WriteFile(filename, []byte(text), 0644)
	check(err)
}

func main() {
	encode := flag.String("e", "", "Encode the given GitHub Markdown file")
	decode := flag.String("d", "", "Decode the given GitHub Markdown file")
	backup := flag.Bool("b", true, "Backup GitHub Markdown file before encode/decode")

	flag.Parse()

	if len(*encode) > 0 {
		convert(*encode, encoder, *backup)
	} else if len(*decode) > 0 {
		convert(*decode, decoder, *backup)
	}
}
<|endoftext|>


<commit_before>/*
Copyright The Helm 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 registry // import "helm.sh/helm/v3/internal/experimental/registry"

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"strings"

	"github.com/sirupsen/logrus"
	orascontext "oras.land/oras-go/pkg/context"
	"oras.land/oras-go/pkg/registry"

	"helm.sh/helm/v3/pkg/chart"
	"helm.sh/helm/v3/pkg/chart/loader"
)

// IsOCI determines whether or not a URL is to be treated as an OCI URL
func IsOCI(url string) bool {
	return strings.HasPrefix(url, fmt.Sprintf("%s://", OCIScheme))
}

// extractChartMeta is used to extract a chart metadata from a byte array
func extractChartMeta(chartData []byte) (*chart.Metadata, error) {
	ch, err := loader.LoadArchive(bytes.NewReader(chartData))
	if err != nil {
		return nil, err
	}
	return ch.Metadata, nil
}

// ctx retrieves a fresh context.
// disable verbose logging coming from ORAS (unless debug is enabled)
func ctx(out io.Writer, debug bool) context.Context {
	if !debug {
		return orascontext.Background()
	}
	ctx := orascontext.WithLoggerFromWriter(context.Background(), out)
	orascontext.GetLogger(ctx).Logger.SetLevel(logrus.DebugLevel)
	return ctx
}

// parseReference will parse and validate the reference, and clean tags when applicable
// 	tags are only cleaned when plus (+) signs are present, and are converted to underscores (_) before pushing
// 	See https://github.com/helm/helm/issues/10166
func parseReference(raw string) (registry.Reference, error) {
	// The sole possible reference modification is replacing plus (+) signs present in tags with underscores (_)
	// 	To do this properly, we first need to identify a tag, and then pass it on to the reference parser
	// 	NOTE: Passing immediately to the reference parser will fail since (+) signs are an invalid tag character, and
	// 		  simply replacing all plus (+) occurances could invalidate other portions of the URI
	parts := strings.Split(raw, ":")
	if len(parts) > 1 && !strings.Contains(parts[len(parts)-1], "/") {
		tag := parts[len(parts)-1]

		if tag != "" {
			// Replace any plus (+) signs with known underscore (_) conversion
			newTag := strings.ReplaceAll(tag, "+", "_")
			raw = strings.ReplaceAll(raw, tag, newTag)
		}
	}

	return registry.ParseReference(raw)
}
<commit_msg>Fix typo. Thanks buildbot. Also comments at 80 chars<commit_after>/*
Copyright The Helm 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 registry // import "helm.sh/helm/v3/internal/experimental/registry"

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"strings"

	"github.com/sirupsen/logrus"
	orascontext "oras.land/oras-go/pkg/context"
	"oras.land/oras-go/pkg/registry"

	"helm.sh/helm/v3/pkg/chart"
	"helm.sh/helm/v3/pkg/chart/loader"
)

// IsOCI determines whether or not a URL is to be treated as an OCI URL
func IsOCI(url string) bool {
	return strings.HasPrefix(url, fmt.Sprintf("%s://", OCIScheme))
}

// extractChartMeta is used to extract a chart metadata from a byte array
func extractChartMeta(chartData []byte) (*chart.Metadata, error) {
	ch, err := loader.LoadArchive(bytes.NewReader(chartData))
	if err != nil {
		return nil, err
	}
	return ch.Metadata, nil
}

// ctx retrieves a fresh context.
// disable verbose logging coming from ORAS (unless debug is enabled)
func ctx(out io.Writer, debug bool) context.Context {
	if !debug {
		return orascontext.Background()
	}
	ctx := orascontext.WithLoggerFromWriter(context.Background(), out)
	orascontext.GetLogger(ctx).Logger.SetLevel(logrus.DebugLevel)
	return ctx
}

// parseReference will parse and validate the reference, and clean tags when
// applicable tags are only cleaned when plus (+) signs are present, and are
// converted to underscores (_) before pushing
// See https://github.com/helm/helm/issues/10166
func parseReference(raw string) (registry.Reference, error) {
	// The sole possible reference modification is replacing plus (+) signs
	// present in tags with underscores (_). To do this properly, we first
	// need to identify a tag, and then pass it on to the reference parser
	// NOTE: Passing immediately to the reference parser will fail since (+)
	// signs are an invalid tag character, and simply replacing all plus (+)
	// occurrences could invalidate other portions of the URI
	parts := strings.Split(raw, ":")
	if len(parts) > 1 && !strings.Contains(parts[len(parts)-1], "/") {
		tag := parts[len(parts)-1]

		if tag != "" {
			// Replace any plus (+) signs with known underscore (_) conversion
			newTag := strings.ReplaceAll(tag, "+", "_")
			raw = strings.ReplaceAll(raw, tag, newTag)
		}
	}

	return registry.ParseReference(raw)
}
<|endoftext|>


<commit_before>package builtin_test

import (
	"os/exec"
	"strconv"
	"syscall"
	"testing"
)

func TestLinuxExit(t *testing.T) {
	type exitDesc struct {
		script string
		status string
		result string
		fail   bool
	}

	tests := map[string]exitDesc{
		"success": {
			script: "./testdata/exit.sh",
			status: "0",
			result: "0",
		},
		"failure": {
			script: "./testdata/exit.sh",
			status: "1",
			result: "1",
		},
		"maxStatus": {
			script: "./testdata/exit.sh",
			status: "255",
			result: "255",
		},
		"statusIsUnsigned": {
			script: "./testdata/exit.sh",
			status: "-1",
			result: "255",
		},
		"statusOverflow": {
			script: "./testdata/exit.sh",
			status: "666",
			result: "154", // Why ? For the glory of satan of course :-)
		},
	}

	//WHY: Not sure this is a great idea, but we need to exec with the
	//code under test nash, not the one installed on the system.
	//Can't circumvent the need for Exec here.
	//Other tests can just run nash inside their own process.
	projectnash := "../../../cmd/nash/nash"

	for name, desc := range tests {
		t.Run(name, func(t *testing.T) {
			cmd := exec.Command(projectnash, desc.script, desc.status)
			err := cmd.Run()
			if err == nil {
				if desc.status == "0" {
					return
				}
				t.Fatalf("expected error for status: %s", desc.status)

			}
			if exiterr, ok := err.(*exec.ExitError); ok {
				if status, ok := exiterr.Sys().(syscall.WaitStatus); ok {
					expectedStatus, err := strconv.Atoi(desc.result)
					if err != nil {
						t.Fatalf("error[%s] converting[%s]", err, desc.status)
					}
					got := status.ExitStatus()
					if expectedStatus != got {
						t.Fatalf("expected[%d] got[%d]", expectedStatus, got)
					}
				}
			} else {
				t.Fatalf("unexpected error: %v", err)
			}
		})
	}
}
<commit_msg>Make tests more robust<commit_after>package builtin_test

import (
	"os/exec"
	"strconv"
	"syscall"
	"testing"
)

func TestLinuxExit(t *testing.T) {
	type exitDesc struct {
		script string
		status string
		result string
		fail   bool
	}

	tests := map[string]exitDesc{
		"success": {
			script: "./testdata/exit.sh",
			status: "0",
			result: "0",
		},
		"failure": {
			script: "./testdata/exit.sh",
			status: "1",
			result: "1",
		},
		"maxStatus": {
			script: "./testdata/exit.sh",
			status: "255",
			result: "255",
		},
		"statusIsUnsigned": {
			script: "./testdata/exit.sh",
			status: "-1",
			result: "255",
		},
		"statusOverflow": {
			script: "./testdata/exit.sh",
			status: "666",
			result: "154", // Why ? For the glory of satan of course :-)
		},
	}

	//WHY: Not sure this is a great idea, but we need to exec with the
	//code under test nash, not the one installed on the system.
	//Can't circumvent the need for Exec here.
	//Other tests can just run nash inside their own process.
	projectnash := "../../../cmd/nash/nash"

	for name, desc := range tests {
		t.Run(name, func(t *testing.T) {
			cmd := exec.Command(projectnash, desc.script, desc.status)
			err := cmd.Run()
			if err == nil {
				if desc.status == "0" {
					return
				}
				t.Fatalf("expected error for status: %s", desc.status)

			}
			if exiterr, ok := err.(*exec.ExitError); ok {
				if status, ok := exiterr.Sys().(syscall.WaitStatus); ok {
					expectedStatus, err := strconv.Atoi(desc.result)
					if err != nil {
						t.Fatalf("error[%s] converting[%s]", err, desc.status)
					}
					got := status.ExitStatus()
					if expectedStatus != got {
						t.Fatalf("expected[%d] got[%d]", expectedStatus, got)
					}
				} else {
					t.Fatal("unable to extract status code from exec")
				}
			} else {
				t.Fatalf("unexpected error: %v", err)
			}
		})
	}
}
<|endoftext|>


<commit_before>// clock counts down to or up from a target time.
package main

import (
	"fmt"
	"os"
	"time"
)

func main() {
	const (
		indent          = "\t"
		highlight_start = "\x1b[1;36m"
		highlight_end   = "\x1b[0m"
	)

	fmt.Print(indent, highlight_start, "Just Go", highlight_end, "\n")
	target := time.Date(2016, 1, 1, 0, 0, 0, 0, time.UTC)
	fmt.Print(indent, target.Format(time.UnixDate), "\n")

	// Exit on Enter key
	go func() {
		buf := make([]byte, 1)
		_, _ = os.Stdin.Read(buf)
		os.Exit(0)
	}()

	var previous time.Time

	for {
		now := time.Now()
		now = now.Add(time.Duration(-now.Nanosecond())) // truncate to second

		if now != previous {
			previous = now
			remaining := target.Sub(now)

			var sign string
			if remaining > 0 {
				sign = "-" // countdown is "T minus..."
			} else {
				sign = "+" // count up is "T plus..."
				remaining = -remaining
			}

			var days int
			if remaining >= 24*time.Hour {
				days = int(remaining / (24 * time.Hour))
				remaining = remaining % (24 * time.Hour)
			}

			fmt.Print(indent, now.Format(time.UnixDate), "  ", sign)
			if days > 0 {
				fmt.Print(days, "d")
			}
			fmt.Print(remaining, "          \r")
		}
		time.Sleep(50 * time.Millisecond)
	}
}
<commit_msg>slight simplification<commit_after>// clock counts down to or up from a target time.
package main

import (
	"fmt"
	"os"
	"time"
)

func main() {
	target := time.Date(2016, 1, 1, 0, 0, 0, 0, time.UTC)

	// Exit on Enter key
	go func() {
		buf := make([]byte, 1)
		_, _ = os.Stdin.Read(buf)
		os.Exit(0)
	}()

	const (
		indent          = "\t"
		highlight_start = "\x1b[1;36m"
		highlight_end   = "\x1b[0m"
	)
	fmt.Print(indent, highlight_start, "Just Go", highlight_end, "\n")

	var previous time.Time

	for {
		now := time.Now()
		now = now.Add(time.Duration(-now.Nanosecond())) // truncate to second

		if now != previous {
			previous = now
			remaining := target.Sub(now)

			var sign string
			if remaining > 0 {
				sign = "-" // countdown is "T minus..."
			} else {
				sign = "+" // count up is "T plus..."
				remaining = -remaining
			}

			var days int
			if remaining >= 24*time.Hour {
				days = int(remaining / (24 * time.Hour))
				remaining = remaining % (24 * time.Hour)
			}

			fmt.Print(indent, now.Format(time.UnixDate), "  ", sign)
			if days > 0 {
				fmt.Print(days, "d")
			}
			fmt.Print(remaining, "          \r")
		}
		time.Sleep(50 * time.Millisecond)
	}
}
<|endoftext|>


<commit_before>// Copyright (c) 2018 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package yarpcerrors

import (
	"fmt"
	"strconv"
	"strings"
)

const (
	// CodeOK means no error; returned on success
	CodeOK Code = 0

	// CodeCancelled means the operation was cancelled, typically by the caller.
	CodeCancelled Code = 1

	// CodeUnknown means an unknown error. Errors raised by APIs
	// that do not return enough error information
	// may be converted to this error.
	CodeUnknown Code = 2

	// CodeInvalidArgument means the client specified an invalid argument.
	// Note that this differs from `FailedPrecondition`. `InvalidArgument`
	// indicates arguments that are problematic regardless of the state of
	// the system (e.g., a malformed file name).
	CodeInvalidArgument Code = 3

	// CodeDeadlineExceeded means the deadline expired before the operation could
	// complete. For operations that change the state of the system, this error
	// may be returned even if the operation has completed successfully. For example,
	// a successful response from a server could have been delayed long
	// enough for the deadline to expire.
	CodeDeadlineExceeded Code = 4

	// CodeNotFound means some requested entity (e.g., file or directory) was not found.
	// For privacy reasons, this code *may* be returned when the client
	// does not have the access rights to the entity, though such usage is
	// discouraged.
	CodeNotFound Code = 5

	// CodeAlreadyExists means the entity that a client attempted to create
	// (e.g., file or directory) already exists.
	CodeAlreadyExists Code = 6

	// CodePermissionDenied means the caller does not have permission to execute
	// the specified operation. `PermissionDenied` must not be used for rejections
	// caused by exhausting some resource (use `ResourceExhausted`
	// instead for those errors). `PermissionDenied` must not be
	// used if the caller can not be identified (use `Unauthenticated`
	// instead for those errors).
	CodePermissionDenied Code = 7

	// CodeResourceExhausted means some resource has been exhausted, perhaps a per-user
	// quota, or perhaps the entire file system is out of space.
	CodeResourceExhausted Code = 8

	// CodeFailedPrecondition means the operation was rejected because the system is not
	// in a state required for the operation's execution. For example, the directory
	// to be deleted is non-empty, an rmdir operation is applied to
	// a non-directory, etc.
	//
	// Service implementors can use the following guidelines to decide
	// between `FailedPrecondition`, `Aborted`, and `Unavailable`:
	//  (a) Use `Unavailable` if the client can retry just the failing call.
	//  (b) Use `Aborted` if the client should retry at a higher level
	//      (e.g., restarting a read-modify-write sequence).
	//  (c) Use `FailedPrecondition` if the client should not retry until
	//      the system state has been explicitly fixed. E.g., if an "rmdir"
	//      fails because the directory is non-empty, `FailedPrecondition`
	//      should be returned since the client should not retry unless
	//      the files are deleted from the directory.
	CodeFailedPrecondition Code = 9

	// CodeAborted means the operation was aborted, typically due to a concurrency issue
	// such as a sequencer check failure or transaction abort.
	//
	// See the guidelines above for deciding between `FailedPrecondition`,
	// `Aborted`, and `Unavailable`.
	CodeAborted Code = 10

	// CodeOutOfRange means the operation was attempted past the valid range.
	// E.g., seeking or reading past end-of-file.
	//
	// Unlike `InvalidArgument`, this error indicates a problem that may
	// be fixed if the system state changes. For example, a 32-bit file
	// system will generate `InvalidArgument` if asked to read at an
	// offset that is not in the range [0,2^32-1], but it will generate
	// `OutOfRange` if asked to read from an offset past the current
	// file size.
	//
	// There is a fair bit of overlap between `FailedPrecondition` and
	// `OutOfRange`.  We recommend using `OutOfRange` (the more specific
	// error) when it applies so that callers who are iterating through
	// a space can easily look for an `OutOfRange` error to detect when
	// they are done.
	CodeOutOfRange Code = 11

	// CodeUnimplemented means the operation is not implemented or is not
	// supported/enabled in this service.
	CodeUnimplemented Code = 12

	// CodeInternal means an internal error. This means that some invariants expected
	// by the underlying system have been broken. This error code is reserved
	// for serious errors.
	CodeInternal Code = 13

	// CodeUnavailable means the service is currently unavailable. This is most likely a
	// transient condition, which can be corrected by retrying with a backoff.
	//
	// See the guidelines above for deciding between `FailedPrecondition`,
	// `Aborted`, and `Unavailable`.
	CodeUnavailable Code = 14

	// CodeDataLoss means unrecoverable data loss or corruption.
	CodeDataLoss Code = 15

	// CodeUnauthenticated means the request does not have valid authentication
	// credentials for the operation.
	CodeUnauthenticated Code = 16
)

var (
	_codeToString = map[Code]string{
		CodeOK:                 "ok",
		CodeCancelled:          "cancelled",
		CodeUnknown:            "unknown",
		CodeInvalidArgument:    "invalid-argument",
		CodeDeadlineExceeded:   "deadline-exceeded",
		CodeNotFound:           "not-found",
		CodeAlreadyExists:      "already-exists",
		CodePermissionDenied:   "permission-denied",
		CodeResourceExhausted:  "resource-exhausted",
		CodeFailedPrecondition: "failed-precondition",
		CodeAborted:            "aborted",
		CodeOutOfRange:         "out-of-range",
		CodeUnimplemented:      "unimplemented",
		CodeInternal:           "internal",
		CodeUnavailable:        "unavailable",
		CodeDataLoss:           "data-loss",
		CodeUnauthenticated:    "unauthenticated",
	}
	_stringToCode = map[string]Code{
		"ok":                  CodeOK,
		"cancelled":           CodeCancelled,
		"unknown":             CodeUnknown,
		"invalid-argument":    CodeInvalidArgument,
		"deadline-exceeded":   CodeDeadlineExceeded,
		"not-found":           CodeNotFound,
		"already-exists":      CodeAlreadyExists,
		"permission-denied":   CodePermissionDenied,
		"resource-exhausted":  CodeResourceExhausted,
		"failed-precondition": CodeFailedPrecondition,
		"aborted":             CodeAborted,
		"out-of-range":        CodeOutOfRange,
		"unimplemented":       CodeUnimplemented,
		"internal":            CodeInternal,
		"unavailable":         CodeUnavailable,
		"data-loss":           CodeDataLoss,
		"unauthenticated":     CodeUnauthenticated,
	}
)

// Code represents the type of error for an RPC call.
//
// Sometimes multiple error codes may apply. Services should return
// the most specific error code that applies. For example, prefer
// `OutOfRange` over `FailedPrecondition` if both codes apply.
// Similarly prefer `NotFound` or `AlreadyExists` over `FailedPrecondition`.
type Code int

// String returns the the string representation of the Code.
func (c Code) String() string {
	s, ok := _codeToString[c]
	if ok {
		return s
	}
	return strconv.Itoa(int(c))
}

// MarshalText implements encoding.TextMarshaler.
func (c Code) MarshalText() ([]byte, error) {
	s, ok := _codeToString[c]
	if ok {
		return []byte(s), nil
	}
	return nil, fmt.Errorf("unknown code: %d", int(c))
}

// UnmarshalText implements encoding.TextUnmarshaler.
func (c *Code) UnmarshalText(text []byte) error {
	i, ok := _stringToCode[strings.ToLower(string(text))]
	if !ok {
		return fmt.Errorf("unknown code string: %s", string(text))
	}
	*c = i
	return nil
}

// MarshalJSON implements json.Marshaler.
func (c Code) MarshalJSON() ([]byte, error) {
	s, ok := _codeToString[c]
	if ok {
		return []byte(`"` + s + `"`), nil
	}
	return nil, fmt.Errorf("unknown code: %d", int(c))
}

// UnmarshalJSON implements json.Unmarshaler.
func (c *Code) UnmarshalJSON(text []byte) error {
	s := string(text)
	if len(s) < 3 || s[0] != '"' || s[len(s)-1] != '"' {
		return fmt.Errorf("invalid code string: %s", s)
	}
	i, ok := _stringToCode[strings.ToLower(s[1:len(s)-1])]
	if !ok {
		return fmt.Errorf("unknown code string: %s", s)
	}
	*c = i
	return nil
}
<commit_msg>Add documentation to yarpcerrors.Code to note that the Codes match gRPC status codes (#1430)<commit_after>// Copyright (c) 2018 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package yarpcerrors

import (
	"fmt"
	"strconv"
	"strings"
)

const (
	// CodeOK means no error; returned on success
	CodeOK Code = 0

	// CodeCancelled means the operation was cancelled, typically by the caller.
	CodeCancelled Code = 1

	// CodeUnknown means an unknown error. Errors raised by APIs
	// that do not return enough error information
	// may be converted to this error.
	CodeUnknown Code = 2

	// CodeInvalidArgument means the client specified an invalid argument.
	// Note that this differs from `FailedPrecondition`. `InvalidArgument`
	// indicates arguments that are problematic regardless of the state of
	// the system (e.g., a malformed file name).
	CodeInvalidArgument Code = 3

	// CodeDeadlineExceeded means the deadline expired before the operation could
	// complete. For operations that change the state of the system, this error
	// may be returned even if the operation has completed successfully. For example,
	// a successful response from a server could have been delayed long
	// enough for the deadline to expire.
	CodeDeadlineExceeded Code = 4

	// CodeNotFound means some requested entity (e.g., file or directory) was not found.
	// For privacy reasons, this code *may* be returned when the client
	// does not have the access rights to the entity, though such usage is
	// discouraged.
	CodeNotFound Code = 5

	// CodeAlreadyExists means the entity that a client attempted to create
	// (e.g., file or directory) already exists.
	CodeAlreadyExists Code = 6

	// CodePermissionDenied means the caller does not have permission to execute
	// the specified operation. `PermissionDenied` must not be used for rejections
	// caused by exhausting some resource (use `ResourceExhausted`
	// instead for those errors). `PermissionDenied` must not be
	// used if the caller can not be identified (use `Unauthenticated`
	// instead for those errors).
	CodePermissionDenied Code = 7

	// CodeResourceExhausted means some resource has been exhausted, perhaps a per-user
	// quota, or perhaps the entire file system is out of space.
	CodeResourceExhausted Code = 8

	// CodeFailedPrecondition means the operation was rejected because the system is not
	// in a state required for the operation's execution. For example, the directory
	// to be deleted is non-empty, an rmdir operation is applied to
	// a non-directory, etc.
	//
	// Service implementors can use the following guidelines to decide
	// between `FailedPrecondition`, `Aborted`, and `Unavailable`:
	//  (a) Use `Unavailable` if the client can retry just the failing call.
	//  (b) Use `Aborted` if the client should retry at a higher level
	//      (e.g., restarting a read-modify-write sequence).
	//  (c) Use `FailedPrecondition` if the client should not retry until
	//      the system state has been explicitly fixed. E.g., if an "rmdir"
	//      fails because the directory is non-empty, `FailedPrecondition`
	//      should be returned since the client should not retry unless
	//      the files are deleted from the directory.
	CodeFailedPrecondition Code = 9

	// CodeAborted means the operation was aborted, typically due to a concurrency issue
	// such as a sequencer check failure or transaction abort.
	//
	// See the guidelines above for deciding between `FailedPrecondition`,
	// `Aborted`, and `Unavailable`.
	CodeAborted Code = 10

	// CodeOutOfRange means the operation was attempted past the valid range.
	// E.g., seeking or reading past end-of-file.
	//
	// Unlike `InvalidArgument`, this error indicates a problem that may
	// be fixed if the system state changes. For example, a 32-bit file
	// system will generate `InvalidArgument` if asked to read at an
	// offset that is not in the range [0,2^32-1], but it will generate
	// `OutOfRange` if asked to read from an offset past the current
	// file size.
	//
	// There is a fair bit of overlap between `FailedPrecondition` and
	// `OutOfRange`.  We recommend using `OutOfRange` (the more specific
	// error) when it applies so that callers who are iterating through
	// a space can easily look for an `OutOfRange` error to detect when
	// they are done.
	CodeOutOfRange Code = 11

	// CodeUnimplemented means the operation is not implemented or is not
	// supported/enabled in this service.
	CodeUnimplemented Code = 12

	// CodeInternal means an internal error. This means that some invariants expected
	// by the underlying system have been broken. This error code is reserved
	// for serious errors.
	CodeInternal Code = 13

	// CodeUnavailable means the service is currently unavailable. This is most likely a
	// transient condition, which can be corrected by retrying with a backoff.
	//
	// See the guidelines above for deciding between `FailedPrecondition`,
	// `Aborted`, and `Unavailable`.
	CodeUnavailable Code = 14

	// CodeDataLoss means unrecoverable data loss or corruption.
	CodeDataLoss Code = 15

	// CodeUnauthenticated means the request does not have valid authentication
	// credentials for the operation.
	CodeUnauthenticated Code = 16
)

var (
	_codeToString = map[Code]string{
		CodeOK:                 "ok",
		CodeCancelled:          "cancelled",
		CodeUnknown:            "unknown",
		CodeInvalidArgument:    "invalid-argument",
		CodeDeadlineExceeded:   "deadline-exceeded",
		CodeNotFound:           "not-found",
		CodeAlreadyExists:      "already-exists",
		CodePermissionDenied:   "permission-denied",
		CodeResourceExhausted:  "resource-exhausted",
		CodeFailedPrecondition: "failed-precondition",
		CodeAborted:            "aborted",
		CodeOutOfRange:         "out-of-range",
		CodeUnimplemented:      "unimplemented",
		CodeInternal:           "internal",
		CodeUnavailable:        "unavailable",
		CodeDataLoss:           "data-loss",
		CodeUnauthenticated:    "unauthenticated",
	}
	_stringToCode = map[string]Code{
		"ok":                  CodeOK,
		"cancelled":           CodeCancelled,
		"unknown":             CodeUnknown,
		"invalid-argument":    CodeInvalidArgument,
		"deadline-exceeded":   CodeDeadlineExceeded,
		"not-found":           CodeNotFound,
		"already-exists":      CodeAlreadyExists,
		"permission-denied":   CodePermissionDenied,
		"resource-exhausted":  CodeResourceExhausted,
		"failed-precondition": CodeFailedPrecondition,
		"aborted":             CodeAborted,
		"out-of-range":        CodeOutOfRange,
		"unimplemented":       CodeUnimplemented,
		"internal":            CodeInternal,
		"unavailable":         CodeUnavailable,
		"data-loss":           CodeDataLoss,
		"unauthenticated":     CodeUnauthenticated,
	}
)

// Code represents the type of error for an RPC call.
//
// Sometimes multiple error codes may apply. Services should return
// the most specific error code that applies. For example, prefer
// `OutOfRange` over `FailedPrecondition` if both codes apply.
// Similarly prefer `NotFound` or `AlreadyExists` over `FailedPrecondition`.
//
// These codes are meant to match gRPC status codes.
// https://godoc.org/google.golang.org/grpc/codes#Code
type Code int

// String returns the the string representation of the Code.
func (c Code) String() string {
	s, ok := _codeToString[c]
	if ok {
		return s
	}
	return strconv.Itoa(int(c))
}

// MarshalText implements encoding.TextMarshaler.
func (c Code) MarshalText() ([]byte, error) {
	s, ok := _codeToString[c]
	if ok {
		return []byte(s), nil
	}
	return nil, fmt.Errorf("unknown code: %d", int(c))
}

// UnmarshalText implements encoding.TextUnmarshaler.
func (c *Code) UnmarshalText(text []byte) error {
	i, ok := _stringToCode[strings.ToLower(string(text))]
	if !ok {
		return fmt.Errorf("unknown code string: %s", string(text))
	}
	*c = i
	return nil
}

// MarshalJSON implements json.Marshaler.
func (c Code) MarshalJSON() ([]byte, error) {
	s, ok := _codeToString[c]
	if ok {
		return []byte(`"` + s + `"`), nil
	}
	return nil, fmt.Errorf("unknown code: %d", int(c))
}

// UnmarshalJSON implements json.Unmarshaler.
func (c *Code) UnmarshalJSON(text []byte) error {
	s := string(text)
	if len(s) < 3 || s[0] != '"' || s[len(s)-1] != '"' {
		return fmt.Errorf("invalid code string: %s", s)
	}
	i, ok := _stringToCode[strings.ToLower(s[1:len(s)-1])]
	if !ok {
		return fmt.Errorf("unknown code string: %s", s)
	}
	*c = i
	return nil
}
<|endoftext|>


<commit_before>package terminfo

import (
	"os"
	"strconv"
	"strings"
)

// ColorLevel is the color level supported by a terminal.
type ColorLevel uint

// ColorLevel values.
const (
	ColorLevelNone ColorLevel = iota
	ColorLevelBasic
	ColorLevelHundreds
	ColorLevelMillions
)

// String satisfies the Stringer interface.
func (c ColorLevel) String() string {
	switch c {
	case ColorLevelBasic:
		return "basic"
	case ColorLevelHundreds:
		return "hundreds"
	case ColorLevelMillions:
		return "millions"
	}
	return "none"
}

// ChromaFormatterName returns the github.com/alecthomas/chroma compatible
// formatter name for the color level.
func (c ColorLevel) ChromaFormatterName() string {
	switch c {
	case ColorLevelBasic:
		return "terminal"
	case ColorLevelHundreds:
		return "terminal256"
	case ColorLevelMillions:
		return "terminal16m"
	}
	return "noop"
}

// ColorLevelFromEnv returns the color level COLORTERM, FORCE_COLOR,
// TERM_PROGRAM, or determined from the TERM environment variable.
func ColorLevelFromEnv() (ColorLevel, error) {
	// check for overriding environment variables
	colorTerm, termProg, forceColor := os.Getenv("COLORTERM"), os.Getenv("TERM_PROGRAM"), os.Getenv("FORCE_COLOR")
	switch {
	case colorTerm == "truecolor" || colorTerm == "24bit" || termProg == "Hyper":
		return ColorLevelMillions, nil
	case colorTerm != "" || forceColor != "":
		return ColorLevelBasic, nil
	case termProg == "Apple_Terminal":
		return ColorLevelHundreds, nil
	case termProg == "iTerm.app":
		ver := os.Getenv("TERM_PROGRAM_VERSION")
		if ver == "" {
			return ColorLevelHundreds, nil
		}
		i, err := strconv.Atoi(strings.Split(ver, ".")[0])
		if err != nil {
			return ColorLevelNone, ErrInvalidTermProgramVersion
		}
		if i == 3 {
			return ColorLevelMillions, nil
		}
		return ColorLevelHundreds, nil
	}

	// otherwise determine from TERM's max_colors capability
	if term := os.Getenv("TERM"); term != "" {
		ti, err := Load(term)
		if err != nil {
			return ColorLevelNone, err
		}

		v, ok := ti.Nums[MaxColors]
		switch {
		case !ok || v <= 16:
			return ColorLevelNone, nil
		case ok && v >= 256:
			return ColorLevelHundreds, nil
		}
	}

	return ColorLevelBasic, nil
}
<commit_msg>Minor change to color detection<commit_after>package terminfo

import (
	"os"
	"strconv"
	"strings"
)

// ColorLevel is the color level supported by a terminal.
type ColorLevel uint

// ColorLevel values.
const (
	ColorLevelNone ColorLevel = iota
	ColorLevelBasic
	ColorLevelHundreds
	ColorLevelMillions
)

// String satisfies the Stringer interface.
func (c ColorLevel) String() string {
	switch c {
	case ColorLevelBasic:
		return "basic"
	case ColorLevelHundreds:
		return "hundreds"
	case ColorLevelMillions:
		return "millions"
	}
	return "none"
}

// ChromaFormatterName returns the github.com/alecthomas/chroma compatible
// formatter name for the color level.
func (c ColorLevel) ChromaFormatterName() string {
	switch c {
	case ColorLevelBasic:
		return "terminal"
	case ColorLevelHundreds:
		return "terminal256"
	case ColorLevelMillions:
		return "terminal16m"
	}
	return "noop"
}

// ColorLevelFromEnv returns the color level COLORTERM, FORCE_COLOR,
// TERM_PROGRAM, or determined from the TERM environment variable.
func ColorLevelFromEnv() (ColorLevel, error) {
	// check for overriding environment variables
	colorTerm, termProg, forceColor := os.Getenv("COLORTERM"), os.Getenv("TERM_PROGRAM"), os.Getenv("FORCE_COLOR")
	switch {
	case strings.Contains(colorTerm, "truecolor") || strings.Contains(colorTerm, "24bit") || termProg == "Hyper":
		return ColorLevelMillions, nil
	case colorTerm != "" || forceColor != "":
		return ColorLevelBasic, nil
	case termProg == "Apple_Terminal":
		return ColorLevelHundreds, nil
	case termProg == "iTerm.app":
		ver := os.Getenv("TERM_PROGRAM_VERSION")
		if ver == "" {
			return ColorLevelHundreds, nil
		}
		i, err := strconv.Atoi(strings.Split(ver, ".")[0])
		if err != nil {
			return ColorLevelNone, ErrInvalidTermProgramVersion
		}
		if i == 3 {
			return ColorLevelMillions, nil
		}
		return ColorLevelHundreds, nil
	}

	// otherwise determine from TERM's max_colors capability
	if term := os.Getenv("TERM"); term != "" {
		ti, err := Load(term)
		if err != nil {
			return ColorLevelNone, err
		}

		v, ok := ti.Nums[MaxColors]
		switch {
		case !ok || v <= 16:
			return ColorLevelNone, nil
		case ok && v >= 256:
			return ColorLevelHundreds, nil
		}
	}

	return ColorLevelBasic, nil
}
<|endoftext|>


<commit_before>package main

const NGINX_BUILD_VERSION = "0.3.3"

// nginx
const (
	NGINX_VERSION             = "1.9.0"
	NGINX_DOWNLOAD_URL_PREFIX = "http://nginx.org/download"
)

// pcre
const (
	PCRE_VERSION             = "8.36"
	PCRE_DOWNLOAD_URL_PREFIX = "http://ftp.csx.cam.ac.uk/pub/software/programming/pcre"
)

// openssl
const (
	OPENSSL_VERSION             = "1.0.2a"
	OPENSSL_DOWNLOAD_URL_PREFIX = "http://www.openssl.org/source"
)

// zlib
const (
	ZLIB_VERSION             = "1.2.8"
	ZLIB_DOWNLOAD_URL_PREFIX = "http://zlib.net"
)

// openResty
const (
	OPENRESTY_VERSION             = "1.7.10.1"
	OPENRESTY_DOWNLOAD_URL_PREFIX = "http://openresty.org/download"
)

// tengine
const (
	TENGINE_VERSION             = "2.1.0"
	TENGINE_DOWNLOAD_URL_PREFIX = "http://tengine.taobao.org/download"
)

// component enumerations
const (
	COMPONENT_NGINX = iota
	COMPONENT_OPENRESTY
	COMPONENT_TENGINE
	COMPONENT_PCRE
	COMPONENT_OPENSSL
	COMPONENT_ZLIB
	COMPONENT_MAX
)
<commit_msg>bumped version to 0.3.4.<commit_after>package main

const NGINX_BUILD_VERSION = "0.3.4"

// nginx
const (
	NGINX_VERSION             = "1.9.0"
	NGINX_DOWNLOAD_URL_PREFIX = "http://nginx.org/download"
)

// pcre
const (
	PCRE_VERSION             = "8.36"
	PCRE_DOWNLOAD_URL_PREFIX = "http://ftp.csx.cam.ac.uk/pub/software/programming/pcre"
)

// openssl
const (
	OPENSSL_VERSION             = "1.0.2a"
	OPENSSL_DOWNLOAD_URL_PREFIX = "http://www.openssl.org/source"
)

// zlib
const (
	ZLIB_VERSION             = "1.2.8"
	ZLIB_DOWNLOAD_URL_PREFIX = "http://zlib.net"
)

// openResty
const (
	OPENRESTY_VERSION             = "1.7.10.1"
	OPENRESTY_DOWNLOAD_URL_PREFIX = "http://openresty.org/download"
)

// tengine
const (
	TENGINE_VERSION             = "2.1.0"
	TENGINE_DOWNLOAD_URL_PREFIX = "http://tengine.taobao.org/download"
)

// component enumerations
const (
	COMPONENT_NGINX = iota
	COMPONENT_OPENRESTY
	COMPONENT_TENGINE
	COMPONENT_PCRE
	COMPONENT_OPENSSL
	COMPONENT_ZLIB
	COMPONENT_MAX
)
<|endoftext|>