<commit_before>// Copyright The OpenTelemetry 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 cortex

import (
	"bytes"
	"context"
	"fmt"
	"log"
	"net/http"

	"github.com/gogo/protobuf/proto"
	"github.com/golang/snappy"
	"github.com/prometheus/prometheus/prompb"

	"go.opentelemetry.io/otel/api/global"
	"go.opentelemetry.io/otel/api/label"
	apimetric "go.opentelemetry.io/otel/api/metric"
	"go.opentelemetry.io/otel/sdk/export/metric"
	export "go.opentelemetry.io/otel/sdk/export/metric"
	"go.opentelemetry.io/otel/sdk/export/metric/aggregation"
	"go.opentelemetry.io/otel/sdk/metric/controller/push"
	"go.opentelemetry.io/otel/sdk/metric/selector/simple"
)

// Exporter forwards metrics to a Cortex instance
type Exporter struct {
	config Config
}

// ExportKindFor returns CumulativeExporter so the Processor correctly aggregates data
func (e *Exporter) ExportKindFor(*apimetric.Descriptor, aggregation.Kind) metric.ExportKind {
	return metric.CumulativeExporter
}

// Export forwards metrics to Cortex from the SDK
func (e *Exporter) Export(_ context.Context, checkpointSet metric.CheckpointSet) error {
	timeseries, err := e.ConvertToTimeSeries(checkpointSet)
	if err != nil {
		return err
	}

	message, buildMessageErr := e.buildMessage(timeseries)
	if buildMessageErr != nil {
		return buildMessageErr
	}

	request, buildRequestErr := e.buildRequest(message)
	if buildRequestErr != nil {
		return buildRequestErr
	}

	sendRequestErr := e.sendRequest(request)
	if sendRequestErr != nil {
		return sendRequestErr
	}

	return nil
}

// NewRawExporter validates the Config struct and creates an Exporter with it.
func NewRawExporter(config Config) (*Exporter, error) {
	// This is redundant when the user creates the Config struct with the NewConfig
	// function.
	if err := config.Validate(); err != nil {
		return nil, err
	}

	exporter := Exporter{config}
	return &exporter, nil
}

// NewExportPipeline sets up a complete export pipeline with a push Controller and
// Exporter.
func NewExportPipeline(config Config, options ...push.Option) (*push.Controller, error) {
	exporter, err := NewRawExporter(config)
	if err != nil {
		return nil, err
	}

	pusher := push.New(
		simple.NewWithExactDistribution(),
		exporter,
		options...,
	)
	pusher.Start()
	return pusher, nil
}

// InstallNewPipeline registers a push Controller's Provider globally.
func InstallNewPipeline(config Config, options ...push.Option) (*push.Controller, error) {
	pusher, err := NewExportPipeline(config, options...)
	if err != nil {
		return nil, err
	}
	global.SetMeterProvider(pusher.Provider())
	return pusher, nil
}

// ConvertToTimeSeries converts a CheckpointSet to a slice of TimeSeries pointers
func (e *Exporter) ConvertToTimeSeries(checkpointSet export.CheckpointSet) ([]*prompb.TimeSeries, error) {
	var aggError error
	var timeSeries []*prompb.TimeSeries

	// Iterate over each record in the checkpoint set and convert to TimeSeries
	aggError = checkpointSet.ForEach(e, func(record metric.Record) error {
		// Convert based on aggregation type
		agg := record.Aggregation()

		// Check if aggregation has Sum value
		if sum, ok := agg.(aggregation.Sum); ok {
			tSeries, err := convertFromSum(record, sum)
			if err != nil {
				return err
			}

			timeSeries = append(timeSeries, tSeries)
		}

		// Check if aggregation has MinMaxSumCount value
		if minMaxSumCount, ok := agg.(aggregation.MinMaxSumCount); ok {
			tSeries, err := convertFromMinMaxSumCount(record, minMaxSumCount)
			if err != nil {
				return err
			}

			timeSeries = append(timeSeries, tSeries...)
		} else if lastValue, ok := agg.(aggregation.LastValue); ok {
			tSeries, err := convertFromLastValue(record, lastValue)
			if err != nil {
				return err
			}

			timeSeries = append(timeSeries, tSeries)
		}

		return nil
	})

	// Check if error was returned in checkpointSet.ForEach()
	if aggError != nil {
		return nil, aggError
	}

	return timeSeries, nil
}

// convertFromSum returns a single TimeSeries based on a Record with a Sum aggregation
func convertFromSum(record metric.Record, sum aggregation.Sum) (*prompb.TimeSeries, error) {
	// Get Sum value
	value, err := sum.Sum()
	if err != nil {
		return nil, err
	}
	// Create sample from Sum value
	sample := prompb.Sample{
		Value:     value.CoerceToFloat64(record.Descriptor().NumberKind()),
		Timestamp: record.EndTime().Unix(), // Convert time to Unix (int64)
	}

	// Create labels, including metric name
	name := sanitize(record.Descriptor().Name())
	labels := createLabelSet(record, "__name__", name)

	// Create TimeSeries and return
	tSeries := &prompb.TimeSeries{
		Samples: []prompb.Sample{sample},
		Labels:  labels,
	}

	return tSeries, nil
}

// convertFromLastValue returns a single TimeSeries based on a Record with a LastValue aggregation
func convertFromLastValue(record metric.Record, lastValue aggregation.LastValue) (*prompb.TimeSeries, error) {
	// Get value
	value, _, err := lastValue.LastValue()
	if err != nil {
		return nil, err
	}

	// Create sample from Last value
	sample := prompb.Sample{
		Value:     value.CoerceToFloat64(record.Descriptor().NumberKind()),
		Timestamp: record.EndTime().Unix(), // Convert time to Unix (int64)
	}

	// Create labels, including metric name
	name := sanitize(record.Descriptor().Name())
	labels := createLabelSet(record, "__name__", name)

	// Create TimeSeries and return
	tSeries := &prompb.TimeSeries{
		Samples: []prompb.Sample{sample},
		Labels:  labels,
	}

	return tSeries, nil
}

// convertFromMinMaxSumCount returns 4 TimeSeries for the min, max, sum, and count from the mmsc aggregation
func convertFromMinMaxSumCount(record metric.Record, minMaxSumCount aggregation.MinMaxSumCount) ([]*prompb.TimeSeries, error) {
	// Convert Min
	min, err := minMaxSumCount.Min()
	if err != nil {
		return nil, err
	}
	minSample := prompb.Sample{
		Value:     min.CoerceToFloat64(record.Descriptor().NumberKind()),
		Timestamp: record.EndTime().Unix(), // Convert time to Unix (int64)
	}

	// Create labels, including metric name
	name := sanitize(record.Descriptor().Name() + "_min")
	labels := createLabelSet(record, "__name__", name)

	// Create TimeSeries
	minTimeSeries := &prompb.TimeSeries{
		Samples: []prompb.Sample{minSample},
		Labels:  labels,
	}

	// Convert Max
	max, err := minMaxSumCount.Max()
	if err != nil {
		return nil, err
	}
	maxSample := prompb.Sample{
		Value:     max.CoerceToFloat64(record.Descriptor().NumberKind()),
		Timestamp: record.EndTime().Unix(), // Convert time to Unix (int64)
	}

	// Create labels, including metric name
	name = sanitize(record.Descriptor().Name() + "_max")
	labels = createLabelSet(record, "__name__", name)

	// Create TimeSeries
	maxTimeSeries := &prompb.TimeSeries{
		Samples: []prompb.Sample{maxSample},
		Labels:  labels,
	}

	// Convert Count
	count, err := minMaxSumCount.Count()
	if err != nil {
		return nil, err
	}
	countSample := prompb.Sample{
		Value:     float64(count),
		Timestamp: record.EndTime().Unix(), // Convert time to Unix (int64)
	}

	// Create labels, including metric name
	name = sanitize(record.Descriptor().Name() + "_count")
	labels = createLabelSet(record, "__name__", name)

	// Create TimeSeries
	countTimeSeries := &prompb.TimeSeries{
		Samples: []prompb.Sample{countSample},
		Labels:  labels,
	}

	tSeries := []*prompb.TimeSeries{
		minTimeSeries, maxTimeSeries, countTimeSeries,
	}

	return tSeries, nil
}

// createLabelSet combines labels from a Record, resource, and extra labels to
// create a slice of prompb.Label
func createLabelSet(record metric.Record, extras ...string) []*prompb.Label {
	// Map ensure no duplicate label names
	labelMap := map[string]prompb.Label{}

	// mergeLabels merges Record and Resource labels into a single set, giving
	// precedence to the record's labels.
	mi := label.NewMergeIterator(record.Labels(), record.Resource().LabelSet())
	for mi.Next() {
		label := mi.Label()
		key := string(label.Key)
		labelMap[key] = prompb.Label{
			Name:  sanitize(key),
			Value: label.Value.Emit(),
		}
	}

	// Add extra labels created by the exporter like the metric name
	// or labels to represent histogram buckets
	for i := 0; i < len(extras); i += 2 {
		// Ensure even number of extras (key : value)
		if i+1 >= len(extras) {
			break
		}

		// Ensure label doesn't exist. If it does, notify user that a user created label
		// is being overwritten by a Prometheus reserved label (e.g. 'le' for histograms)
		_, found := labelMap[extras[i]]
		if found {
			log.Printf("Label %s is overwritten. Check if Prometheus reserved labels are used.\n", extras[i])
		}
		labelMap[extras[i]] = prompb.Label{
			Name:  sanitize(extras[i]),
			Value: extras[i+1],
		}
	}

	// Create slice of labels from labelMap and return
	res := make([]*prompb.Label, 0, len(labelMap))
	for _, lb := range labelMap {
		currentLabel := lb
		res = append(res, &currentLabel)
	}

	return res
}

// AddHeaders adds required headers as well as all headers in Header map to a http request.
func (e *Exporter) addHeaders(req *http.Request) {
	// Cortex expects Snappy-compressed protobuf messages. These two headers are hard-coded as they
	// should be on every request.
	req.Header.Add("X-Prometheus-Remote-Write-Version", "0.1.0")
	req.Header.Add("Content-Encoding", "snappy")
	req.Header.Set("Content-Type", "application/x-protobuf")

	// Add all user-supplied headers to the request.
	for name, field := range e.config.Headers {
		req.Header.Add(name, field)
	}
}

// buildMessage creates a Snappy-compressed protobuf message from a slice of TimeSeries.
func (e *Exporter) buildMessage(timeseries []*prompb.TimeSeries) ([]byte, error) {
	// Wrap the TimeSeries as a WriteRequest since Cortex requires it.
	writeRequest := &prompb.WriteRequest{
		Timeseries: timeseries,
	}

	// Convert the struct to a slice of bytes and then compress it.
	message, err := proto.Marshal(writeRequest)
	if err != nil {
		return nil, err
	}
	compressed := snappy.Encode(nil, message)

	return compressed, nil
}

// buildRequest creates an http POST request with a Snappy-compressed protocol buffer
// message as the body and with all the headers attached.
func (e *Exporter) buildRequest(message []byte) (*http.Request, error) {
	req, err := http.NewRequest(
		http.MethodPost,
		e.config.Endpoint,
		bytes.NewBuffer(message),
	)
	if err != nil {
		return nil, err
	}

	// Add the required headers and the headers from Config.Headers.
	e.addHeaders(req)

	return req, nil
}

// sendRequest sends an http request using the Exporter's http Client.
func (e *Exporter) sendRequest(req *http.Request) error {
	// Set a client if the user didn't provide one.
	if e.config.Client == nil {
		e.config.Client = http.DefaultClient
		e.config.Client.Timeout = e.config.RemoteTimeout
	}

	// Attempt to send request.
	res, err := e.config.Client.Do(req)
	if err != nil {
		return err
	}
	defer res.Body.Close()

	// The response should have a status code of 200.
	if res.StatusCode != http.StatusOK {
		return fmt.Errorf("%v", res.Status)
	}
	return nil
}
<commit_msg>Cortex: Remove sanitize call (#215)<commit_after>// Copyright The OpenTelemetry 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 cortex

import (
	"bytes"
	"context"
	"fmt"
	"log"
	"net/http"

	"github.com/gogo/protobuf/proto"
	"github.com/golang/snappy"
	"github.com/prometheus/prometheus/prompb"

	"go.opentelemetry.io/otel/api/global"
	"go.opentelemetry.io/otel/api/label"
	apimetric "go.opentelemetry.io/otel/api/metric"
	"go.opentelemetry.io/otel/sdk/export/metric"
	export "go.opentelemetry.io/otel/sdk/export/metric"
	"go.opentelemetry.io/otel/sdk/export/metric/aggregation"
	"go.opentelemetry.io/otel/sdk/metric/controller/push"
	"go.opentelemetry.io/otel/sdk/metric/selector/simple"
)

// Exporter forwards metrics to a Cortex instance
type Exporter struct {
	config Config
}

// ExportKindFor returns CumulativeExporter so the Processor correctly aggregates data
func (e *Exporter) ExportKindFor(*apimetric.Descriptor, aggregation.Kind) metric.ExportKind {
	return metric.CumulativeExporter
}

// Export forwards metrics to Cortex from the SDK
func (e *Exporter) Export(_ context.Context, checkpointSet metric.CheckpointSet) error {
	timeseries, err := e.ConvertToTimeSeries(checkpointSet)
	if err != nil {
		return err
	}

	message, buildMessageErr := e.buildMessage(timeseries)
	if buildMessageErr != nil {
		return buildMessageErr
	}

	request, buildRequestErr := e.buildRequest(message)
	if buildRequestErr != nil {
		return buildRequestErr
	}

	sendRequestErr := e.sendRequest(request)
	if sendRequestErr != nil {
		return sendRequestErr
	}

	return nil
}

// NewRawExporter validates the Config struct and creates an Exporter with it.
func NewRawExporter(config Config) (*Exporter, error) {
	// This is redundant when the user creates the Config struct with the NewConfig
	// function.
	if err := config.Validate(); err != nil {
		return nil, err
	}

	exporter := Exporter{config}
	return &exporter, nil
}

// NewExportPipeline sets up a complete export pipeline with a push Controller and
// Exporter.
func NewExportPipeline(config Config, options ...push.Option) (*push.Controller, error) {
	exporter, err := NewRawExporter(config)
	if err != nil {
		return nil, err
	}

	pusher := push.New(
		simple.NewWithExactDistribution(),
		exporter,
		options...,
	)
	pusher.Start()
	return pusher, nil
}

// InstallNewPipeline registers a push Controller's Provider globally.
func InstallNewPipeline(config Config, options ...push.Option) (*push.Controller, error) {
	pusher, err := NewExportPipeline(config, options...)
	if err != nil {
		return nil, err
	}
	global.SetMeterProvider(pusher.Provider())
	return pusher, nil
}

// ConvertToTimeSeries converts a CheckpointSet to a slice of TimeSeries pointers
func (e *Exporter) ConvertToTimeSeries(checkpointSet export.CheckpointSet) ([]*prompb.TimeSeries, error) {
	var aggError error
	var timeSeries []*prompb.TimeSeries

	// Iterate over each record in the checkpoint set and convert to TimeSeries
	aggError = checkpointSet.ForEach(e, func(record metric.Record) error {
		// Convert based on aggregation type
		agg := record.Aggregation()

		// Check if aggregation has Sum value
		if sum, ok := agg.(aggregation.Sum); ok {
			tSeries, err := convertFromSum(record, sum)
			if err != nil {
				return err
			}

			timeSeries = append(timeSeries, tSeries)
		}

		// Check if aggregation has MinMaxSumCount value
		if minMaxSumCount, ok := agg.(aggregation.MinMaxSumCount); ok {
			tSeries, err := convertFromMinMaxSumCount(record, minMaxSumCount)
			if err != nil {
				return err
			}

			timeSeries = append(timeSeries, tSeries...)
		} else if lastValue, ok := agg.(aggregation.LastValue); ok {
			tSeries, err := convertFromLastValue(record, lastValue)
			if err != nil {
				return err
			}

			timeSeries = append(timeSeries, tSeries)
		}

		return nil
	})

	// Check if error was returned in checkpointSet.ForEach()
	if aggError != nil {
		return nil, aggError
	}

	return timeSeries, nil
}

// convertFromSum returns a single TimeSeries based on a Record with a Sum aggregation
func convertFromSum(record metric.Record, sum aggregation.Sum) (*prompb.TimeSeries, error) {
	// Get Sum value
	value, err := sum.Sum()
	if err != nil {
		return nil, err
	}
	// Create sample from Sum value
	sample := prompb.Sample{
		Value:     value.CoerceToFloat64(record.Descriptor().NumberKind()),
		Timestamp: record.EndTime().Unix(), // Convert time to Unix (int64)
	}

	// Create labels, including metric name
	name := sanitize(record.Descriptor().Name())
	labels := createLabelSet(record, "__name__", name)

	// Create TimeSeries and return
	tSeries := &prompb.TimeSeries{
		Samples: []prompb.Sample{sample},
		Labels:  labels,
	}

	return tSeries, nil
}

// convertFromLastValue returns a single TimeSeries based on a Record with a LastValue aggregation
func convertFromLastValue(record metric.Record, lastValue aggregation.LastValue) (*prompb.TimeSeries, error) {
	// Get value
	value, _, err := lastValue.LastValue()
	if err != nil {
		return nil, err
	}

	// Create sample from Last value
	sample := prompb.Sample{
		Value:     value.CoerceToFloat64(record.Descriptor().NumberKind()),
		Timestamp: record.EndTime().Unix(), // Convert time to Unix (int64)
	}

	// Create labels, including metric name
	name := sanitize(record.Descriptor().Name())
	labels := createLabelSet(record, "__name__", name)

	// Create TimeSeries and return
	tSeries := &prompb.TimeSeries{
		Samples: []prompb.Sample{sample},
		Labels:  labels,
	}

	return tSeries, nil
}

// convertFromMinMaxSumCount returns 4 TimeSeries for the min, max, sum, and count from the mmsc aggregation
func convertFromMinMaxSumCount(record metric.Record, minMaxSumCount aggregation.MinMaxSumCount) ([]*prompb.TimeSeries, error) {
	// Convert Min
	min, err := minMaxSumCount.Min()
	if err != nil {
		return nil, err
	}
	minSample := prompb.Sample{
		Value:     min.CoerceToFloat64(record.Descriptor().NumberKind()),
		Timestamp: record.EndTime().Unix(), // Convert time to Unix (int64)
	}

	// Create labels, including metric name
	name := sanitize(record.Descriptor().Name() + "_min")
	labels := createLabelSet(record, "__name__", name)

	// Create TimeSeries
	minTimeSeries := &prompb.TimeSeries{
		Samples: []prompb.Sample{minSample},
		Labels:  labels,
	}

	// Convert Max
	max, err := minMaxSumCount.Max()
	if err != nil {
		return nil, err
	}
	maxSample := prompb.Sample{
		Value:     max.CoerceToFloat64(record.Descriptor().NumberKind()),
		Timestamp: record.EndTime().Unix(), // Convert time to Unix (int64)
	}

	// Create labels, including metric name
	name = sanitize(record.Descriptor().Name() + "_max")
	labels = createLabelSet(record, "__name__", name)

	// Create TimeSeries
	maxTimeSeries := &prompb.TimeSeries{
		Samples: []prompb.Sample{maxSample},
		Labels:  labels,
	}

	// Convert Count
	count, err := minMaxSumCount.Count()
	if err != nil {
		return nil, err
	}
	countSample := prompb.Sample{
		Value:     float64(count),
		Timestamp: record.EndTime().Unix(), // Convert time to Unix (int64)
	}

	// Create labels, including metric name
	name = sanitize(record.Descriptor().Name() + "_count")
	labels = createLabelSet(record, "__name__", name)

	// Create TimeSeries
	countTimeSeries := &prompb.TimeSeries{
		Samples: []prompb.Sample{countSample},
		Labels:  labels,
	}

	tSeries := []*prompb.TimeSeries{
		minTimeSeries, maxTimeSeries, countTimeSeries,
	}

	return tSeries, nil
}

// createLabelSet combines labels from a Record, resource, and extra labels to
// create a slice of prompb.Label
func createLabelSet(record metric.Record, extras ...string) []*prompb.Label {
	// Map ensure no duplicate label names
	labelMap := map[string]prompb.Label{}

	// mergeLabels merges Record and Resource labels into a single set, giving
	// precedence to the record's labels.
	mi := label.NewMergeIterator(record.Labels(), record.Resource().LabelSet())
	for mi.Next() {
		label := mi.Label()
		key := string(label.Key)
		labelMap[key] = prompb.Label{
			Name:  sanitize(key),
			Value: label.Value.Emit(),
		}
	}

	// Add extra labels created by the exporter like the metric name
	// or labels to represent histogram buckets
	for i := 0; i < len(extras); i += 2 {
		// Ensure even number of extras (key : value)
		if i+1 >= len(extras) {
			break
		}

		// Ensure label doesn't exist. If it does, notify user that a user created label
		// is being overwritten by a Prometheus reserved label (e.g. 'le' for histograms)
		_, found := labelMap[extras[i]]
		if found {
			log.Printf("Label %s is overwritten. Check if Prometheus reserved labels are used.\n", extras[i])
		}
		labelMap[extras[i]] = prompb.Label{
			Name:  extras[i],
			Value: extras[i+1],
		}
	}

	// Create slice of labels from labelMap and return
	res := make([]*prompb.Label, 0, len(labelMap))
	for _, lb := range labelMap {
		currentLabel := lb
		res = append(res, &currentLabel)
	}

	return res
}

// AddHeaders adds required headers as well as all headers in Header map to a http request.
func (e *Exporter) addHeaders(req *http.Request) {
	// Cortex expects Snappy-compressed protobuf messages. These two headers are hard-coded as they
	// should be on every request.
	req.Header.Add("X-Prometheus-Remote-Write-Version", "0.1.0")
	req.Header.Add("Content-Encoding", "snappy")
	req.Header.Set("Content-Type", "application/x-protobuf")

	// Add all user-supplied headers to the request.
	for name, field := range e.config.Headers {
		req.Header.Add(name, field)
	}
}

// buildMessage creates a Snappy-compressed protobuf message from a slice of TimeSeries.
func (e *Exporter) buildMessage(timeseries []*prompb.TimeSeries) ([]byte, error) {
	// Wrap the TimeSeries as a WriteRequest since Cortex requires it.
	writeRequest := &prompb.WriteRequest{
		Timeseries: timeseries,
	}

	// Convert the struct to a slice of bytes and then compress it.
	message, err := proto.Marshal(writeRequest)
	if err != nil {
		return nil, err
	}
	compressed := snappy.Encode(nil, message)

	return compressed, nil
}

// buildRequest creates an http POST request with a Snappy-compressed protocol buffer
// message as the body and with all the headers attached.
func (e *Exporter) buildRequest(message []byte) (*http.Request, error) {
	req, err := http.NewRequest(
		http.MethodPost,
		e.config.Endpoint,
		bytes.NewBuffer(message),
	)
	if err != nil {
		return nil, err
	}

	// Add the required headers and the headers from Config.Headers.
	e.addHeaders(req)

	return req, nil
}

// sendRequest sends an http request using the Exporter's http Client.
func (e *Exporter) sendRequest(req *http.Request) error {
	// Set a client if the user didn't provide one.
	if e.config.Client == nil {
		e.config.Client = http.DefaultClient
		e.config.Client.Timeout = e.config.RemoteTimeout
	}

	// Attempt to send request.
	res, err := e.config.Client.Do(req)
	if err != nil {
		return err
	}
	defer res.Body.Close()

	// The response should have a status code of 200.
	if res.StatusCode != http.StatusOK {
		return fmt.Errorf("%v", res.Status)
	}
	return nil
}
<|endoftext|>


<commit_before>//  Copyright 2019 Google Inc. All Rights Reserved.
//
//  Licensed under the Apache License, Version 2.0 (the "License");
//  you may not use this file except in compliance with the License.
//  You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
//  Unless required by applicable law or agreed to in writing, software
//  distributed under the License is distributed on an "AS IS" BASIS,
//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//  See the License for the specific language governing permissions and
//  limitations under the License.

// Package patch contains end to end tests for patch management
package patch

import (
	"context"
	"errors"
	"fmt"
	"log"
	"path"
	"regexp"
	"strconv"
	"sync"
	"time"

	daisyCompute "github.com/GoogleCloudPlatform/compute-image-tools/daisy/compute"
	"github.com/GoogleCloudPlatform/compute-image-tools/osconfig_tests/compute"
	gcpclients "github.com/GoogleCloudPlatform/compute-image-tools/osconfig_tests/gcp_clients"
	"github.com/GoogleCloudPlatform/compute-image-tools/osconfig_tests/junitxml"
	testconfig "github.com/GoogleCloudPlatform/compute-image-tools/osconfig_tests/test_config"
	"github.com/GoogleCloudPlatform/compute-image-tools/osconfig_tests/utils"
	"github.com/golang/protobuf/ptypes/duration"
	"github.com/kylelemons/godebug/pretty"
	api "google.golang.org/api/compute/v1"

	osconfigpb "github.com/GoogleCloudPlatform/osconfig/_internal/gapi-cloud-osconfig-go/google.golang.org/genproto/googleapis/cloud/osconfig/v1alpha1"
)

const (
	testSuiteName = "OSPatch"
)

var (
	dump       = &pretty.Config{IncludeUnexported: true}
	testSuffix = utils.RandString(3)
)

// TestSuite is a OSPatch test suite.
func TestSuite(ctx context.Context, tswg *sync.WaitGroup, testSuites chan *junitxml.TestSuite, logger *log.Logger, testSuiteRegex, testCaseRegex *regexp.Regexp, testProjectConfig *testconfig.Project) {
	defer tswg.Done()

	if testSuiteRegex != nil && !testSuiteRegex.MatchString(testSuiteName) {
		return
	}

	testSuite := junitxml.NewTestSuite(testSuiteName)
	defer testSuite.Finish(testSuites)

	logger.Printf("Running TestSuite %q", testSuite.Name)

	var wg sync.WaitGroup
	tests := make(chan *junitxml.TestCase)
	// Basic functionality smoke test against all latest images.
	for _, setup := range headImageTestSetup() {
		wg.Add(1)
		s := setup
		tc := junitxml.NewTestCase(testSuiteName, fmt.Sprintf("[Execute PatchJob] [%s]", s.testName))
		f := func() { runExecutePatchJobTest(ctx, tc, s, testProjectConfig, nil) }
		go runTestCase(tc, f, tests, &wg, logger, testCaseRegex)
	}
	// Test that updates trigger reboot as expected.
	for _, setup := range oldImageTestSetup() {
		wg.Add(1)
		s := setup
		tc := junitxml.NewTestCase(testSuiteName, fmt.Sprintf("[PatchJob triggers reboot] [%s]", s.testName))
		shouldReboot := true
		f := func() {
			runRebootPatchTest(ctx, tc, s, testProjectConfig, &osconfigpb.PatchConfig{Apt: &osconfigpb.AptSettings{Type: osconfigpb.AptSettings_DIST}}, shouldReboot)
		}
		go runTestCase(tc, f, tests, &wg, logger, testCaseRegex)
	}
	// Test that PatchConfig_NEVER prevents reboot.
	for _, setup := range oldImageTestSetup() {
		wg.Add(1)
		s := setup
		tc := junitxml.NewTestCase(testSuiteName, fmt.Sprintf("[PatchJob does not reboot] [%s]", s.testName))
		pc := &osconfigpb.PatchConfig{RebootConfig: osconfigpb.PatchConfig_NEVER, Apt: &osconfigpb.AptSettings{Type: osconfigpb.AptSettings_DIST}}
		shouldReboot := false
		f := func() { runRebootPatchTest(ctx, tc, s, testProjectConfig, pc, shouldReboot) }
		go runTestCase(tc, f, tests, &wg, logger, testCaseRegex)
	}
	// Test APT specific functionality, this just tests that using these settings doesn't break anything.
	for _, setup := range aptHeadImageTestSetup() {
		wg.Add(1)
		s := setup
		tc := junitxml.NewTestCase(testSuiteName, fmt.Sprintf("[APT dist-upgrade] [%s]", s.testName))
		f := func() {
			runExecutePatchJobTest(ctx, tc, s, testProjectConfig, &osconfigpb.PatchConfig{Apt: &osconfigpb.AptSettings{Type: osconfigpb.AptSettings_DIST}})
		}
		go runTestCase(tc, f, tests, &wg, logger, testCaseRegex)
	}
	// Test YUM specific functionality, this just tests that using these settings doesn't break anything.
	for _, setup := range yumHeadImageTestSetup() {
		wg.Add(1)
		s := setup
		tc := junitxml.NewTestCase(testSuiteName, fmt.Sprintf("[YUM security, minimal and excludes] [%s]", s.testName))
		f := func() {
			runExecutePatchJobTest(ctx, tc, s, testProjectConfig, &osconfigpb.PatchConfig{Yum: &osconfigpb.YumSettings{Security: true, Minimal: true, Excludes: []string{"pkg1", "pkg2"}}})
		}
		go runTestCase(tc, f, tests, &wg, logger, testCaseRegex)
	}

	go func() {
		wg.Wait()
		close(tests)
	}()

	for ret := range tests {
		testSuite.TestCase = append(testSuite.TestCase, ret)
	}

	logger.Printf("Finished TestSuite %q", testSuite.Name)
}

func awaitPatchJob(ctx context.Context, job *osconfigpb.PatchJob, timeout time.Duration) (*osconfigpb.PatchJob, error) {
	client, err := gcpclients.GetOsConfigClient(ctx)
	if err != nil {
		return nil, err
	}
	tick := time.Tick(10 * time.Second)
	timedout := time.Tick(timeout)
	for {
		select {
		case <-timedout:
			return nil, errors.New("timed out")
		case <-tick:
			req := &osconfigpb.GetPatchJobRequest{
				Name: job.GetName(),
			}
			res, err := client.GetPatchJob(ctx, req)
			if err != nil {
				return nil, fmt.Errorf("error while fetching patch job: %s", utils.GetStatusFromError(err))
			}

			if isPatchJobFailureState(res.State) {
				return nil, fmt.Errorf("failure status %v with message '%s'", res.State, job.GetErrorMessage())
			}

			if res.State == osconfigpb.PatchJob_SUCCEEDED {
				if res.InstanceDetailsSummary.GetInstancesSucceeded() < 1 && res.InstanceDetailsSummary.GetInstancesSucceededRebootRequired() < 1 {
					return nil, errors.New("completed with no instances patched")
				}
				return res, nil
			}
		}
	}
}

func runExecutePatchJobTest(ctx context.Context, testCase *junitxml.TestCase, testSetup *patchTestSetup, testProjectConfig *testconfig.Project, pc *osconfigpb.PatchConfig) {
	client, err := daisyCompute.NewClient(ctx)
	if err != nil {
		testCase.WriteFailure("error creating client: %v", err)
		return
	}

	testCase.Logf("Creating instance with image %q", testSetup.image)
	var metadataItems []*api.MetadataItems
	metadataItems = append(metadataItems, testSetup.startup)
	metadataItems = append(metadataItems, compute.BuildInstanceMetadataItem("os-config-enabled-prerelease-features", "ospatch"))
	name := fmt.Sprintf("patch-test-%s-%s-%s", path.Base(testSetup.testName), testSuffix, utils.RandString(5))
	inst, err := utils.CreateComputeInstance(metadataItems, client, testSetup.machineType, testSetup.image, name, testProjectConfig.TestProjectID, testProjectConfig.GetZone(), testProjectConfig.ServiceAccountEmail, testProjectConfig.ServiceAccountScopes)
	if err != nil {
		testCase.WriteFailure("Error creating instance: %v", utils.GetStatusFromError(err))
		return
	}
	defer inst.Cleanup()

	testCase.Logf("Waiting for agent install to complete")
	if _, err := inst.WaitForGuestAttributes("osconfig_tests/", "install_done", 5*time.Second, 5*time.Minute); err != nil {
		testCase.WriteFailure("Error waiting for osconfig agent install: %v", err)
		return
	}

	testCase.Logf("Agent installed successfully")

	parent := fmt.Sprintf("projects/%s", testProjectConfig.TestProjectID)
	osconfigClient, err := gcpclients.GetOsConfigClient(ctx)

	req := &osconfigpb.ExecutePatchJobRequest{
		Parent:      parent,
		Description: "testing patch job run",
		Filter:      fmt.Sprintf("name=\"%s\"", name),
		Duration:    &duration.Duration{Seconds: int64(testSetup.assertTimeout / time.Second)},
		PatchConfig: pc,
	}
	job, err := osconfigClient.ExecutePatchJob(ctx, req)

	if err != nil {
		testCase.WriteFailure("error while executing patch job: \n%s\n", utils.GetStatusFromError(err))
		return
	}

	testCase.Logf("Started patch job '%s'", job.GetName())
	if _, err := awaitPatchJob(ctx, job, testSetup.assertTimeout); err != nil {
		testCase.WriteFailure("Patch job '%s' error: %v", job.GetName(), err)
	}
}

func runRebootPatchTest(ctx context.Context, testCase *junitxml.TestCase, testSetup *patchTestSetup, testProjectConfig *testconfig.Project, pc *osconfigpb.PatchConfig, shouldReboot bool) {
	client, err := daisyCompute.NewClient(ctx)
	if err != nil {
		testCase.WriteFailure("error creating client: %v", err)
		return
	}

	testCase.Logf("Creating instance with image %q", testSetup.image)
	var metadataItems []*api.MetadataItems
	metadataItems = append(metadataItems, testSetup.startup)
	metadataItems = append(metadataItems, compute.BuildInstanceMetadataItem("os-config-enabled-prerelease-features", "ospatch"))
	name := fmt.Sprintf("patch-reboot-%s-%s-%s", path.Base(testSetup.testName), testSuffix, utils.RandString(5))
	inst, err := utils.CreateComputeInstance(metadataItems, client, testSetup.machineType, testSetup.image, name, testProjectConfig.TestProjectID, testProjectConfig.GetZone(), testProjectConfig.ServiceAccountEmail, testProjectConfig.ServiceAccountScopes)
	if err != nil {
		testCase.WriteFailure("Error creating instance: %v", utils.GetStatusFromError(err))
		return
	}
	defer inst.Cleanup()

	testCase.Logf("Waiting for agent install to complete")
	if _, err := inst.WaitForGuestAttributes("osconfig_tests/", "install_done", 5*time.Second, 5*time.Minute); err != nil {
		testCase.WriteFailure("Error waiting for osconfig agent install: %v", err)
		return
	}

	testCase.Logf("Agent installed successfully")

	parent := fmt.Sprintf("projects/%s", testProjectConfig.TestProjectID)
	osconfigClient, err := gcpclients.GetOsConfigClient(ctx)

	req := &osconfigpb.ExecutePatchJobRequest{
		Parent:      parent,
		Description: "testing patch job reboot",
		Filter:      fmt.Sprintf("name=\"%s\"", name),
		Duration:    &duration.Duration{Seconds: int64(testSetup.assertTimeout / time.Second)},
		PatchConfig: pc,
	}
	job, err := osconfigClient.ExecutePatchJob(ctx, req)

	if err != nil {
		testCase.WriteFailure("error while executing patch job: \n%s\n", utils.GetStatusFromError(err))
		return
	}

	testCase.Logf("Started patch job '%s'", job.GetName())
	pj, err := awaitPatchJob(ctx, job, testSetup.assertTimeout)
	if err != nil {
		testCase.WriteFailure("Patch job '%s' error: %v", job.GetName(), err)
	}

	// If shouldReboot is true that instance should not report a pending reboot.
	if shouldReboot && pj.InstanceDetailsSummary.GetInstancesSucceededRebootRequired() > 0 {
		testCase.WriteFailure("PatchJob finished with status InstancesSucceededRebootRequired.")
		return
	}
	// If shouldReboot is false that instance should report a pending reboot.
	if !shouldReboot && pj.InstanceDetailsSummary.GetInstancesSucceededRebootRequired() == 0 {
		testCase.WriteFailure("PatchJob should have finished with status InstancesSucceededRebootRequired.")
		return
	}

	testCase.Logf("Checking reboot count")
	attr, err := inst.GetGuestAttributes("osconfig_tests/", "boot_count")
	if err != nil {
		testCase.WriteFailure("Error retrieving boot count: %v", err)
		return
	}
	if len(attr) == 0 {
		testCase.WriteFailure("Error retrieving boot count: osconfig_tests/boot_count attribute empty")
		return
	}
	num, err := strconv.Atoi(attr[0].Value)
	if err != nil {
		testCase.WriteFailure("Error parsing boot count: %v", err)
		return
	}
	if shouldReboot && num < 2 {
		testCase.WriteFailure("Instance should have booted at least 2 times, boot num: %d.", num)
		return
	}
	if !shouldReboot && num > 1 {
		testCase.WriteFailure("Instance should not have booted more that 1 time, boot num: %d.", num)
		return
	}
}

func isPatchJobFailureState(state osconfigpb.PatchJob_State) bool {
	return state == osconfigpb.PatchJob_COMPLETED_WITH_ERRORS ||
		state == osconfigpb.PatchJob_TIMED_OUT ||
		state == osconfigpb.PatchJob_CANCELED
}

func runTestCase(tc *junitxml.TestCase, f func(), tests chan *junitxml.TestCase, wg *sync.WaitGroup, logger *log.Logger, regex *regexp.Regexp) {
	defer wg.Done()

	if tc.FilterTestCase(regex) {
		tc.Finish(tests)
	} else {
		logger.Printf("Running TestCase %q", tc.Name)
		f()
		tc.Finish(tests)
		logger.Printf("TestCase %q finished in %fs", tc.Name, tc.Time)
	}
}
<commit_msg>osconfig_tests: fix nil pointer bug (#819)<commit_after>//  Copyright 2019 Google Inc. All Rights Reserved.
//
//  Licensed under the Apache License, Version 2.0 (the "License");
//  you may not use this file except in compliance with the License.
//  You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
//  Unless required by applicable law or agreed to in writing, software
//  distributed under the License is distributed on an "AS IS" BASIS,
//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//  See the License for the specific language governing permissions and
//  limitations under the License.

// Package patch contains end to end tests for patch management
package patch

import (
	"context"
	"errors"
	"fmt"
	"log"
	"path"
	"regexp"
	"strconv"
	"sync"
	"time"

	daisyCompute "github.com/GoogleCloudPlatform/compute-image-tools/daisy/compute"
	"github.com/GoogleCloudPlatform/compute-image-tools/osconfig_tests/compute"
	gcpclients "github.com/GoogleCloudPlatform/compute-image-tools/osconfig_tests/gcp_clients"
	"github.com/GoogleCloudPlatform/compute-image-tools/osconfig_tests/junitxml"
	testconfig "github.com/GoogleCloudPlatform/compute-image-tools/osconfig_tests/test_config"
	"github.com/GoogleCloudPlatform/compute-image-tools/osconfig_tests/utils"
	"github.com/golang/protobuf/ptypes/duration"
	"github.com/kylelemons/godebug/pretty"
	api "google.golang.org/api/compute/v1"

	osconfigpb "github.com/GoogleCloudPlatform/osconfig/_internal/gapi-cloud-osconfig-go/google.golang.org/genproto/googleapis/cloud/osconfig/v1alpha1"
)

const (
	testSuiteName = "OSPatch"
)

var (
	dump       = &pretty.Config{IncludeUnexported: true}
	testSuffix = utils.RandString(3)
)

// TestSuite is a OSPatch test suite.
func TestSuite(ctx context.Context, tswg *sync.WaitGroup, testSuites chan *junitxml.TestSuite, logger *log.Logger, testSuiteRegex, testCaseRegex *regexp.Regexp, testProjectConfig *testconfig.Project) {
	defer tswg.Done()

	if testSuiteRegex != nil && !testSuiteRegex.MatchString(testSuiteName) {
		return
	}

	testSuite := junitxml.NewTestSuite(testSuiteName)
	defer testSuite.Finish(testSuites)

	logger.Printf("Running TestSuite %q", testSuite.Name)

	var wg sync.WaitGroup
	tests := make(chan *junitxml.TestCase)
	// Basic functionality smoke test against all latest images.
	for _, setup := range headImageTestSetup() {
		wg.Add(1)
		s := setup
		tc := junitxml.NewTestCase(testSuiteName, fmt.Sprintf("[Execute PatchJob] [%s]", s.testName))
		f := func() { runExecutePatchJobTest(ctx, tc, s, testProjectConfig, nil) }
		go runTestCase(tc, f, tests, &wg, logger, testCaseRegex)
	}
	// Test that updates trigger reboot as expected.
	for _, setup := range oldImageTestSetup() {
		wg.Add(1)
		s := setup
		tc := junitxml.NewTestCase(testSuiteName, fmt.Sprintf("[PatchJob triggers reboot] [%s]", s.testName))
		shouldReboot := true
		f := func() {
			runRebootPatchTest(ctx, tc, s, testProjectConfig, &osconfigpb.PatchConfig{Apt: &osconfigpb.AptSettings{Type: osconfigpb.AptSettings_DIST}}, shouldReboot)
		}
		go runTestCase(tc, f, tests, &wg, logger, testCaseRegex)
	}
	// Test that PatchConfig_NEVER prevents reboot.
	for _, setup := range oldImageTestSetup() {
		wg.Add(1)
		s := setup
		tc := junitxml.NewTestCase(testSuiteName, fmt.Sprintf("[PatchJob does not reboot] [%s]", s.testName))
		pc := &osconfigpb.PatchConfig{RebootConfig: osconfigpb.PatchConfig_NEVER, Apt: &osconfigpb.AptSettings{Type: osconfigpb.AptSettings_DIST}}
		shouldReboot := false
		f := func() { runRebootPatchTest(ctx, tc, s, testProjectConfig, pc, shouldReboot) }
		go runTestCase(tc, f, tests, &wg, logger, testCaseRegex)
	}
	// Test APT specific functionality, this just tests that using these settings doesn't break anything.
	for _, setup := range aptHeadImageTestSetup() {
		wg.Add(1)
		s := setup
		tc := junitxml.NewTestCase(testSuiteName, fmt.Sprintf("[APT dist-upgrade] [%s]", s.testName))
		f := func() {
			runExecutePatchJobTest(ctx, tc, s, testProjectConfig, &osconfigpb.PatchConfig{Apt: &osconfigpb.AptSettings{Type: osconfigpb.AptSettings_DIST}})
		}
		go runTestCase(tc, f, tests, &wg, logger, testCaseRegex)
	}
	// Test YUM specific functionality, this just tests that using these settings doesn't break anything.
	for _, setup := range yumHeadImageTestSetup() {
		wg.Add(1)
		s := setup
		tc := junitxml.NewTestCase(testSuiteName, fmt.Sprintf("[YUM security, minimal and excludes] [%s]", s.testName))
		f := func() {
			runExecutePatchJobTest(ctx, tc, s, testProjectConfig, &osconfigpb.PatchConfig{Yum: &osconfigpb.YumSettings{Security: true, Minimal: true, Excludes: []string{"pkg1", "pkg2"}}})
		}
		go runTestCase(tc, f, tests, &wg, logger, testCaseRegex)
	}

	go func() {
		wg.Wait()
		close(tests)
	}()

	for ret := range tests {
		testSuite.TestCase = append(testSuite.TestCase, ret)
	}

	logger.Printf("Finished TestSuite %q", testSuite.Name)
}

func awaitPatchJob(ctx context.Context, job *osconfigpb.PatchJob, timeout time.Duration) (*osconfigpb.PatchJob, error) {
	client, err := gcpclients.GetOsConfigClient(ctx)
	if err != nil {
		return nil, err
	}
	tick := time.Tick(10 * time.Second)
	timedout := time.Tick(timeout)
	for {
		select {
		case <-timedout:
			return nil, errors.New("timed out")
		case <-tick:
			req := &osconfigpb.GetPatchJobRequest{
				Name: job.GetName(),
			}
			res, err := client.GetPatchJob(ctx, req)
			if err != nil {
				return nil, fmt.Errorf("error while fetching patch job: %s", utils.GetStatusFromError(err))
			}

			if isPatchJobFailureState(res.State) {
				return nil, fmt.Errorf("failure status %v with message '%s'", res.State, job.GetErrorMessage())
			}

			if res.State == osconfigpb.PatchJob_SUCCEEDED {
				if res.InstanceDetailsSummary.GetInstancesSucceeded() < 1 && res.InstanceDetailsSummary.GetInstancesSucceededRebootRequired() < 1 {
					return nil, errors.New("completed with no instances patched")
				}
				return res, nil
			}
		}
	}
}

func runExecutePatchJobTest(ctx context.Context, testCase *junitxml.TestCase, testSetup *patchTestSetup, testProjectConfig *testconfig.Project, pc *osconfigpb.PatchConfig) {
	client, err := daisyCompute.NewClient(ctx)
	if err != nil {
		testCase.WriteFailure("error creating client: %v", err)
		return
	}

	testCase.Logf("Creating instance with image %q", testSetup.image)
	var metadataItems []*api.MetadataItems
	metadataItems = append(metadataItems, testSetup.startup)
	metadataItems = append(metadataItems, compute.BuildInstanceMetadataItem("os-config-enabled-prerelease-features", "ospatch"))
	name := fmt.Sprintf("patch-test-%s-%s-%s", path.Base(testSetup.testName), testSuffix, utils.RandString(5))
	inst, err := utils.CreateComputeInstance(metadataItems, client, testSetup.machineType, testSetup.image, name, testProjectConfig.TestProjectID, testProjectConfig.GetZone(), testProjectConfig.ServiceAccountEmail, testProjectConfig.ServiceAccountScopes)
	if err != nil {
		testCase.WriteFailure("Error creating instance: %v", utils.GetStatusFromError(err))
		return
	}
	defer inst.Cleanup()

	testCase.Logf("Waiting for agent install to complete")
	if _, err := inst.WaitForGuestAttributes("osconfig_tests/", "install_done", 5*time.Second, 5*time.Minute); err != nil {
		testCase.WriteFailure("Error waiting for osconfig agent install: %v", err)
		return
	}

	testCase.Logf("Agent installed successfully")

	parent := fmt.Sprintf("projects/%s", testProjectConfig.TestProjectID)
	osconfigClient, err := gcpclients.GetOsConfigClient(ctx)

	req := &osconfigpb.ExecutePatchJobRequest{
		Parent:      parent,
		Description: "testing patch job run",
		Filter:      fmt.Sprintf("name=\"%s\"", name),
		Duration:    &duration.Duration{Seconds: int64(testSetup.assertTimeout / time.Second)},
		PatchConfig: pc,
	}
	job, err := osconfigClient.ExecutePatchJob(ctx, req)

	if err != nil {
		testCase.WriteFailure("error while executing patch job: \n%s\n", utils.GetStatusFromError(err))
		return
	}

	testCase.Logf("Started patch job '%s'", job.GetName())
	if _, err := awaitPatchJob(ctx, job, testSetup.assertTimeout); err != nil {
		testCase.WriteFailure("Patch job '%s' error: %v", job.GetName(), err)
	}
}

func runRebootPatchTest(ctx context.Context, testCase *junitxml.TestCase, testSetup *patchTestSetup, testProjectConfig *testconfig.Project, pc *osconfigpb.PatchConfig, shouldReboot bool) {
	client, err := daisyCompute.NewClient(ctx)
	if err != nil {
		testCase.WriteFailure("error creating client: %v", err)
		return
	}

	testCase.Logf("Creating instance with image %q", testSetup.image)
	var metadataItems []*api.MetadataItems
	metadataItems = append(metadataItems, testSetup.startup)
	metadataItems = append(metadataItems, compute.BuildInstanceMetadataItem("os-config-enabled-prerelease-features", "ospatch"))
	name := fmt.Sprintf("patch-reboot-%s-%s-%s", path.Base(testSetup.testName), testSuffix, utils.RandString(5))
	inst, err := utils.CreateComputeInstance(metadataItems, client, testSetup.machineType, testSetup.image, name, testProjectConfig.TestProjectID, testProjectConfig.GetZone(), testProjectConfig.ServiceAccountEmail, testProjectConfig.ServiceAccountScopes)
	if err != nil {
		testCase.WriteFailure("Error creating instance: %v", utils.GetStatusFromError(err))
		return
	}
	defer inst.Cleanup()

	testCase.Logf("Waiting for agent install to complete")
	if _, err := inst.WaitForGuestAttributes("osconfig_tests/", "install_done", 5*time.Second, 5*time.Minute); err != nil {
		testCase.WriteFailure("Error waiting for osconfig agent install: %v", err)
		return
	}

	testCase.Logf("Agent installed successfully")

	parent := fmt.Sprintf("projects/%s", testProjectConfig.TestProjectID)
	osconfigClient, err := gcpclients.GetOsConfigClient(ctx)

	req := &osconfigpb.ExecutePatchJobRequest{
		Parent:      parent,
		Description: "testing patch job reboot",
		Filter:      fmt.Sprintf("name=\"%s\"", name),
		Duration:    &duration.Duration{Seconds: int64(testSetup.assertTimeout / time.Second)},
		PatchConfig: pc,
	}
	job, err := osconfigClient.ExecutePatchJob(ctx, req)

	if err != nil {
		testCase.WriteFailure("error while executing patch job: \n%s\n", utils.GetStatusFromError(err))
		return
	}

	testCase.Logf("Started patch job '%s'", job.GetName())
	pj, err := awaitPatchJob(ctx, job, testSetup.assertTimeout)
	if err != nil {
		testCase.WriteFailure("Patch job '%s' error: %v", job.GetName(), err)
	}

	// If shouldReboot is true that instance should not report a pending reboot.
	if shouldReboot && pj.GetInstanceDetailsSummary().GetInstancesSucceededRebootRequired() > 0 {
		testCase.WriteFailure("PatchJob finished with status InstancesSucceededRebootRequired.")
		return
	}
	// If shouldReboot is false that instance should report a pending reboot.
	if !shouldReboot && pj.GetInstanceDetailsSummary().GetInstancesSucceededRebootRequired() == 0 {
		testCase.WriteFailure("PatchJob should have finished with status InstancesSucceededRebootRequired.")
		return
	}

	testCase.Logf("Checking reboot count")
	attr, err := inst.GetGuestAttributes("osconfig_tests/", "boot_count")
	if err != nil {
		testCase.WriteFailure("Error retrieving boot count: %v", err)
		return
	}
	if len(attr) == 0 {
		testCase.WriteFailure("Error retrieving boot count: osconfig_tests/boot_count attribute empty")
		return
	}
	num, err := strconv.Atoi(attr[0].Value)
	if err != nil {
		testCase.WriteFailure("Error parsing boot count: %v", err)
		return
	}
	if shouldReboot && num < 2 {
		testCase.WriteFailure("Instance should have booted at least 2 times, boot num: %d.", num)
		return
	}
	if !shouldReboot && num > 1 {
		testCase.WriteFailure("Instance should not have booted more that 1 time, boot num: %d.", num)
		return
	}
}

func isPatchJobFailureState(state osconfigpb.PatchJob_State) bool {
	return state == osconfigpb.PatchJob_COMPLETED_WITH_ERRORS ||
		state == osconfigpb.PatchJob_TIMED_OUT ||
		state == osconfigpb.PatchJob_CANCELED
}

func runTestCase(tc *junitxml.TestCase, f func(), tests chan *junitxml.TestCase, wg *sync.WaitGroup, logger *log.Logger, regex *regexp.Regexp) {
	defer wg.Done()

	if tc.FilterTestCase(regex) {
		tc.Finish(tests)
	} else {
		logger.Printf("Running TestCase %q", tc.Name)
		f()
		tc.Finish(tests)
		logger.Printf("TestCase %q finished in %fs", tc.Name, tc.Time)
	}
}
<|endoftext|>


<commit_before>package pstoremem

import (
	"context"
	"sort"
	"sync"
	"time"

	logging "github.com/ipfs/go-log"
	peer "github.com/libp2p/go-libp2p-peer"
	ma "github.com/multiformats/go-multiaddr"

	pstore "github.com/libp2p/go-libp2p-peerstore"
	addr "github.com/libp2p/go-libp2p-peerstore/addr"
)

var log = logging.Logger("peerstore")

type expiringAddr struct {
	Addr    ma.Multiaddr
	TTL     time.Duration
	Expires time.Time
}

func (e *expiringAddr) ExpiredBy(t time.Time) bool {
	return t.After(e.Expires)
}

type addrSlice []expiringAddr

var _ pstore.AddrBook = (*memoryAddrBook)(nil)

// memoryAddrBook manages addresses.
type memoryAddrBook struct {
	addrmu sync.Mutex
	addrs  map[peer.ID]addrSlice

	subManager *AddrSubManager
}

func NewAddrBook() pstore.AddrBook {
	return &memoryAddrBook{
		addrs:      make(map[peer.ID]addrSlice),
		subManager: NewAddrSubManager(),
	}
}

func (mab *memoryAddrBook) PeersWithAddrs() peer.IDSlice {
	mab.addrmu.Lock()
	defer mab.addrmu.Unlock()
	if mab.addrs == nil {
		return nil
	}

	pids := make(peer.IDSlice, 0, len(mab.addrs))
	for pid := range mab.addrs {
		pids = append(pids, pid)
	}
	return pids
}

// AddAddr calls AddAddrs(p, []ma.Multiaddr{addr}, ttl)
func (mab *memoryAddrBook) AddAddr(p peer.ID, addr ma.Multiaddr, ttl time.Duration) {
	mab.AddAddrs(p, []ma.Multiaddr{addr}, ttl)
}

// AddAddrs gives memoryAddrBook addresses to use, with a given ttl
// (time-to-live), after which the address is no longer valid.
// If the manager has a longer TTL, the operation is a no-op for that address
func (mab *memoryAddrBook) AddAddrs(p peer.ID, addrs []ma.Multiaddr, ttl time.Duration) {
	mab.addrmu.Lock()
	defer mab.addrmu.Unlock()

	// if ttl is zero, exit. nothing to do.
	if ttl <= 0 {
		return
	}

	oldAddrs := mab.addrs[p]
	amap := make(map[string]expiringAddr, len(oldAddrs))
	for _, ea := range oldAddrs {
		amap[string(ea.Addr.Bytes())] = ea
	}

	// only expand ttls
	exp := time.Now().Add(ttl)
	for _, addr := range addrs {
		if addr == nil {
			log.Warningf("was passed nil multiaddr for %s", p)
			continue
		}

		addrstr := string(addr.Bytes())
		a, found := amap[addrstr]
		if !found || exp.After(a.Expires) {
			amap[addrstr] = expiringAddr{Addr: addr, Expires: exp, TTL: ttl}

			mab.subManager.BroadcastAddr(p, addr)
		}
	}
	newAddrs := make([]expiringAddr, 0, len(amap))
	for _, ea := range amap {
		newAddrs = append(newAddrs, ea)
	}
	mab.addrs[p] = newAddrs
}

// SetAddr calls mgr.SetAddrs(p, addr, ttl)
func (mab *memoryAddrBook) SetAddr(p peer.ID, addr ma.Multiaddr, ttl time.Duration) {
	mab.SetAddrs(p, []ma.Multiaddr{addr}, ttl)
}

// SetAddrs sets the ttl on addresses. This clears any TTL there previously.
// This is used when we receive the best estimate of the validity of an address.
func (mab *memoryAddrBook) SetAddrs(p peer.ID, addrs []ma.Multiaddr, ttl time.Duration) {
	mab.addrmu.Lock()
	defer mab.addrmu.Unlock()

	oldAddrs := mab.addrs[p]
	amap := make(map[string]expiringAddr, len(oldAddrs))
	for _, ea := range oldAddrs {
		amap[string(ea.Addr.Bytes())] = ea
	}

	exp := time.Now().Add(ttl)
	for _, addr := range addrs {
		if addr == nil {
			log.Warningf("was passed nil multiaddr for %s", p)
			continue
		}
		// re-set all of them for new ttl.
		addrs := string(addr.Bytes())

		if ttl > 0 {
			amap[addrs] = expiringAddr{Addr: addr, Expires: exp, TTL: ttl}

			mab.subManager.BroadcastAddr(p, addr)
		} else {
			delete(amap, addrs)
		}
	}
	newAddrs := make([]expiringAddr, 0, len(amap))
	for _, ea := range amap {
		newAddrs = append(newAddrs, ea)
	}
	mab.addrs[p] = newAddrs
}

// UpdateAddrs updates the addresses associated with the given peer that have
// the given oldTTL to have the given newTTL.
func (mab *memoryAddrBook) UpdateAddrs(p peer.ID, oldTTL time.Duration, newTTL time.Duration) {
	mab.addrmu.Lock()
	defer mab.addrmu.Unlock()

	if mab.addrs == nil {
		return
	}

	addrs, found := mab.addrs[p]
	if !found {
		return
	}

	exp := time.Now().Add(newTTL)
	// TODO: RK - Shorthand.
	for i := range addrs {
		aexp := &addrs[i]
		if oldTTL == aexp.TTL {
			aexp.TTL = newTTL
			aexp.Expires = exp
		}
	}
}

// Addresses returns all known (and valid) addresses for a given
func (mab *memoryAddrBook) Addrs(p peer.ID) []ma.Multiaddr {
	mab.addrmu.Lock()
	defer mab.addrmu.Unlock()

	// not initialized? nothing to give.
	if mab.addrs == nil {
		return nil
	}

	maddrs, found := mab.addrs[p]
	if !found {
		return nil
	}

	now := time.Now()
	good := make([]ma.Multiaddr, 0, len(maddrs))
	cleaned := make([]expiringAddr, 0, len(maddrs))
	for _, m := range maddrs {
		if !m.ExpiredBy(now) {
			cleaned = append(cleaned, m)
			good = append(good, m.Addr)
		}
	}

	// clean up the expired ones.
	if len(cleaned) == 0 {
		delete(mab.addrs, p)
	} else {
		mab.addrs[p] = cleaned
	}
	return good
}

// ClearAddrs removes all previously stored addresses
func (mab *memoryAddrBook) ClearAddrs(p peer.ID) {
	mab.addrmu.Lock()
	defer mab.addrmu.Unlock()

	delete(mab.addrs, p)
}

// AddrStream returns a channel on which all new addresses discovered for a
// given peer ID will be published.
func (mab *memoryAddrBook) AddrStream(ctx context.Context, p peer.ID) <-chan ma.Multiaddr {
	mab.addrmu.Lock()
	defer mab.addrmu.Unlock()

	baseaddrslice := mab.addrs[p]
	initial := make([]ma.Multiaddr, 0, len(baseaddrslice))
	for _, a := range baseaddrslice {
		initial = append(initial, a.Addr)
	}

	return mab.subManager.AddrStream(ctx, p, initial)
}

type addrSub struct {
	pubch  chan ma.Multiaddr
	lk     sync.Mutex
	buffer []ma.Multiaddr
	ctx    context.Context
}

func (s *addrSub) pubAddr(a ma.Multiaddr) {
	select {
	case s.pubch <- a:
	case <-s.ctx.Done():
	}
}

// An abstracted, pub-sub manager for address streams. Extracted from
// memoryAddrBook in order to support additional implementations.
type AddrSubManager struct {
	mu   sync.RWMutex
	subs map[peer.ID][]*addrSub
}

// NewAddrSubManager initializes an AddrSubManager.
func NewAddrSubManager() *AddrSubManager {
	return &AddrSubManager{
		subs: make(map[peer.ID][]*addrSub),
	}
}

// Used internally by the address stream coroutine to remove a subscription
// from the manager.
func (mgr *AddrSubManager) removeSub(p peer.ID, s *addrSub) {
	mgr.mu.Lock()
	defer mgr.mu.Unlock()

	subs := mgr.subs[p]
	if len(subs) == 1 {
		if subs[0] != s {
			return
		}
		delete(mgr.subs, p)
		return
	}

	for i, v := range subs {
		if v == s {
			subs[i] = subs[len(subs)-1]
			subs[len(subs)-1] = nil
			mgr.subs[p] = subs[:len(subs)-1]
			return
		}
	}
}

// BroadcastAddr broadcasts a new address to all subscribed streams.
func (mgr *AddrSubManager) BroadcastAddr(p peer.ID, addr ma.Multiaddr) {
	mgr.mu.RLock()
	defer mgr.mu.RUnlock()

	if subs, ok := mgr.subs[p]; ok {
		for _, sub := range subs {
			sub.pubAddr(addr)
		}
	}
}

// AddrStream creates a new subscription for a given peer ID, pre-populating the
// channel with any addresses we might already have on file.
func (mgr *AddrSubManager) AddrStream(ctx context.Context, p peer.ID, initial []ma.Multiaddr) <-chan ma.Multiaddr {
	sub := &addrSub{pubch: make(chan ma.Multiaddr), ctx: ctx}
	out := make(chan ma.Multiaddr)

	mgr.mu.Lock()
	if _, ok := mgr.subs[p]; ok {
		mgr.subs[p] = append(mgr.subs[p], sub)
	} else {
		mgr.subs[p] = []*addrSub{sub}
	}
	mgr.mu.Unlock()

	sort.Sort(addr.AddrList(initial))

	go func(buffer []ma.Multiaddr) {
		defer close(out)

		sent := make(map[string]bool, len(buffer))
		var outch chan ma.Multiaddr

		for _, a := range buffer {
			sent[string(a.Bytes())] = true
		}

		var next ma.Multiaddr
		if len(buffer) > 0 {
			next = buffer[0]
			buffer = buffer[1:]
			outch = out
		}

		for {
			select {
			case outch <- next:
				if len(buffer) > 0 {
					next = buffer[0]
					buffer = buffer[1:]
				} else {
					outch = nil
					next = nil
				}
			case naddr := <-sub.pubch:
				if sent[string(naddr.Bytes())] {
					continue
				}

				sent[string(naddr.Bytes())] = true
				if next == nil {
					next = naddr
					outch = out
				} else {
					buffer = append(buffer, naddr)
				}
			case <-ctx.Done():
				mgr.removeSub(p, sub)
				return
			}
		}

	}(initial)

	return out
}
<commit_msg>Resolve TODO.<commit_after>package pstoremem

import (
	"context"
	"sort"
	"sync"
	"time"

	logging "github.com/ipfs/go-log"
	peer "github.com/libp2p/go-libp2p-peer"
	ma "github.com/multiformats/go-multiaddr"

	pstore "github.com/libp2p/go-libp2p-peerstore"
	addr "github.com/libp2p/go-libp2p-peerstore/addr"
)

var log = logging.Logger("peerstore")

type expiringAddr struct {
	Addr    ma.Multiaddr
	TTL     time.Duration
	Expires time.Time
}

func (e *expiringAddr) ExpiredBy(t time.Time) bool {
	return t.After(e.Expires)
}

type addrSlice []expiringAddr

var _ pstore.AddrBook = (*memoryAddrBook)(nil)

// memoryAddrBook manages addresses.
type memoryAddrBook struct {
	addrmu sync.Mutex
	addrs  map[peer.ID]addrSlice

	subManager *AddrSubManager
}

func NewAddrBook() pstore.AddrBook {
	return &memoryAddrBook{
		addrs:      make(map[peer.ID]addrSlice),
		subManager: NewAddrSubManager(),
	}
}

func (mab *memoryAddrBook) PeersWithAddrs() peer.IDSlice {
	mab.addrmu.Lock()
	defer mab.addrmu.Unlock()
	if mab.addrs == nil {
		return nil
	}

	pids := make(peer.IDSlice, 0, len(mab.addrs))
	for pid := range mab.addrs {
		pids = append(pids, pid)
	}
	return pids
}

// AddAddr calls AddAddrs(p, []ma.Multiaddr{addr}, ttl)
func (mab *memoryAddrBook) AddAddr(p peer.ID, addr ma.Multiaddr, ttl time.Duration) {
	mab.AddAddrs(p, []ma.Multiaddr{addr}, ttl)
}

// AddAddrs gives memoryAddrBook addresses to use, with a given ttl
// (time-to-live), after which the address is no longer valid.
// If the manager has a longer TTL, the operation is a no-op for that address
func (mab *memoryAddrBook) AddAddrs(p peer.ID, addrs []ma.Multiaddr, ttl time.Duration) {
	mab.addrmu.Lock()
	defer mab.addrmu.Unlock()

	// if ttl is zero, exit. nothing to do.
	if ttl <= 0 {
		return
	}

	oldAddrs := mab.addrs[p]
	amap := make(map[string]expiringAddr, len(oldAddrs))
	for _, ea := range oldAddrs {
		amap[string(ea.Addr.Bytes())] = ea
	}

	// only expand ttls
	exp := time.Now().Add(ttl)
	for _, addr := range addrs {
		if addr == nil {
			log.Warningf("was passed nil multiaddr for %s", p)
			continue
		}

		addrstr := string(addr.Bytes())
		a, found := amap[addrstr]
		if !found || exp.After(a.Expires) {
			amap[addrstr] = expiringAddr{Addr: addr, Expires: exp, TTL: ttl}

			mab.subManager.BroadcastAddr(p, addr)
		}
	}
	newAddrs := make([]expiringAddr, 0, len(amap))
	for _, ea := range amap {
		newAddrs = append(newAddrs, ea)
	}
	mab.addrs[p] = newAddrs
}

// SetAddr calls mgr.SetAddrs(p, addr, ttl)
func (mab *memoryAddrBook) SetAddr(p peer.ID, addr ma.Multiaddr, ttl time.Duration) {
	mab.SetAddrs(p, []ma.Multiaddr{addr}, ttl)
}

// SetAddrs sets the ttl on addresses. This clears any TTL there previously.
// This is used when we receive the best estimate of the validity of an address.
func (mab *memoryAddrBook) SetAddrs(p peer.ID, addrs []ma.Multiaddr, ttl time.Duration) {
	mab.addrmu.Lock()
	defer mab.addrmu.Unlock()

	oldAddrs := mab.addrs[p]
	amap := make(map[string]expiringAddr, len(oldAddrs))
	for _, ea := range oldAddrs {
		amap[string(ea.Addr.Bytes())] = ea
	}

	exp := time.Now().Add(ttl)
	for _, addr := range addrs {
		if addr == nil {
			log.Warningf("was passed nil multiaddr for %s", p)
			continue
		}
		// re-set all of them for new ttl.
		addrs := string(addr.Bytes())

		if ttl > 0 {
			amap[addrs] = expiringAddr{Addr: addr, Expires: exp, TTL: ttl}

			mab.subManager.BroadcastAddr(p, addr)
		} else {
			delete(amap, addrs)
		}
	}
	newAddrs := make([]expiringAddr, 0, len(amap))
	for _, ea := range amap {
		newAddrs = append(newAddrs, ea)
	}
	mab.addrs[p] = newAddrs
}

// UpdateAddrs updates the addresses associated with the given peer that have
// the given oldTTL to have the given newTTL.
func (mab *memoryAddrBook) UpdateAddrs(p peer.ID, oldTTL time.Duration, newTTL time.Duration) {
	mab.addrmu.Lock()
	defer mab.addrmu.Unlock()

	if mab.addrs == nil {
		return
	}

	addrs, found := mab.addrs[p]
	if !found {
		return
	}

	exp := time.Now().Add(newTTL)
	for i := range addrs {
		if aexp := &addrs[i]; oldTTL == aexp.TTL {
			aexp.TTL = newTTL
			aexp.Expires = exp
		}
	}
}

// Addresses returns all known (and valid) addresses for a given
func (mab *memoryAddrBook) Addrs(p peer.ID) []ma.Multiaddr {
	mab.addrmu.Lock()
	defer mab.addrmu.Unlock()

	// not initialized? nothing to give.
	if mab.addrs == nil {
		return nil
	}

	maddrs, found := mab.addrs[p]
	if !found {
		return nil
	}

	now := time.Now()
	good := make([]ma.Multiaddr, 0, len(maddrs))
	cleaned := make([]expiringAddr, 0, len(maddrs))
	for _, m := range maddrs {
		if !m.ExpiredBy(now) {
			cleaned = append(cleaned, m)
			good = append(good, m.Addr)
		}
	}

	// clean up the expired ones.
	if len(cleaned) == 0 {
		delete(mab.addrs, p)
	} else {
		mab.addrs[p] = cleaned
	}
	return good
}

// ClearAddrs removes all previously stored addresses
func (mab *memoryAddrBook) ClearAddrs(p peer.ID) {
	mab.addrmu.Lock()
	defer mab.addrmu.Unlock()

	delete(mab.addrs, p)
}

// AddrStream returns a channel on which all new addresses discovered for a
// given peer ID will be published.
func (mab *memoryAddrBook) AddrStream(ctx context.Context, p peer.ID) <-chan ma.Multiaddr {
	mab.addrmu.Lock()
	defer mab.addrmu.Unlock()

	baseaddrslice := mab.addrs[p]
	initial := make([]ma.Multiaddr, 0, len(baseaddrslice))
	for _, a := range baseaddrslice {
		initial = append(initial, a.Addr)
	}

	return mab.subManager.AddrStream(ctx, p, initial)
}

type addrSub struct {
	pubch  chan ma.Multiaddr
	lk     sync.Mutex
	buffer []ma.Multiaddr
	ctx    context.Context
}

func (s *addrSub) pubAddr(a ma.Multiaddr) {
	select {
	case s.pubch <- a:
	case <-s.ctx.Done():
	}
}

// An abstracted, pub-sub manager for address streams. Extracted from
// memoryAddrBook in order to support additional implementations.
type AddrSubManager struct {
	mu   sync.RWMutex
	subs map[peer.ID][]*addrSub
}

// NewAddrSubManager initializes an AddrSubManager.
func NewAddrSubManager() *AddrSubManager {
	return &AddrSubManager{
		subs: make(map[peer.ID][]*addrSub),
	}
}

// Used internally by the address stream coroutine to remove a subscription
// from the manager.
func (mgr *AddrSubManager) removeSub(p peer.ID, s *addrSub) {
	mgr.mu.Lock()
	defer mgr.mu.Unlock()

	subs := mgr.subs[p]
	if len(subs) == 1 {
		if subs[0] != s {
			return
		}
		delete(mgr.subs, p)
		return
	}

	for i, v := range subs {
		if v == s {
			subs[i] = subs[len(subs)-1]
			subs[len(subs)-1] = nil
			mgr.subs[p] = subs[:len(subs)-1]
			return
		}
	}
}

// BroadcastAddr broadcasts a new address to all subscribed streams.
func (mgr *AddrSubManager) BroadcastAddr(p peer.ID, addr ma.Multiaddr) {
	mgr.mu.RLock()
	defer mgr.mu.RUnlock()

	if subs, ok := mgr.subs[p]; ok {
		for _, sub := range subs {
			sub.pubAddr(addr)
		}
	}
}

// AddrStream creates a new subscription for a given peer ID, pre-populating the
// channel with any addresses we might already have on file.
func (mgr *AddrSubManager) AddrStream(ctx context.Context, p peer.ID, initial []ma.Multiaddr) <-chan ma.Multiaddr {
	sub := &addrSub{pubch: make(chan ma.Multiaddr), ctx: ctx}
	out := make(chan ma.Multiaddr)

	mgr.mu.Lock()
	if _, ok := mgr.subs[p]; ok {
		mgr.subs[p] = append(mgr.subs[p], sub)
	} else {
		mgr.subs[p] = []*addrSub{sub}
	}
	mgr.mu.Unlock()

	sort.Sort(addr.AddrList(initial))

	go func(buffer []ma.Multiaddr) {
		defer close(out)

		sent := make(map[string]bool, len(buffer))
		var outch chan ma.Multiaddr

		for _, a := range buffer {
			sent[string(a.Bytes())] = true
		}

		var next ma.Multiaddr
		if len(buffer) > 0 {
			next = buffer[0]
			buffer = buffer[1:]
			outch = out
		}

		for {
			select {
			case outch <- next:
				if len(buffer) > 0 {
					next = buffer[0]
					buffer = buffer[1:]
				} else {
					outch = nil
					next = nil
				}
			case naddr := <-sub.pubch:
				if sent[string(naddr.Bytes())] {
					continue
				}

				sent[string(naddr.Bytes())] = true
				if next == nil {
					next = naddr
					outch = out
				} else {
					buffer = append(buffer, naddr)
				}
			case <-ctx.Done():
				mgr.removeSub(p, sub)
				return
			}
		}

	}(initial)

	return out
}
<|endoftext|>


<commit_before>// Copyright (c) 2017 Uber Technologies, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// 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 json

import (
	"testing"

	"github.com/stretchr/testify/assert"

	"github.com/jaegertracing/jaeger/model"
)

func TestProcessHashtable(t *testing.T) {
	ht := &processHashtable{}

	p1 := model.NewProcess("s1", []model.KeyValue{
		model.String("host", "google.com"),
	})
	p1dup := model.NewProcess("s1", []model.KeyValue{
		model.String("host", "google.com"),
	})
	p2 := model.NewProcess("s2", []model.KeyValue{
		model.String("host", "facebook.com"),
	})

	assert.Equal(t, "p1", ht.getKey(p1))
	assert.Equal(t, "p1", ht.getKey(p1))
	assert.Equal(t, "p1", ht.getKey(p1dup))
	assert.Equal(t, "p2", ht.getKey(p2))

	expectedMapping := map[string]*model.Process{
		"p1": p1,
		"p2": p2,
	}
	assert.Equal(t, expectedMapping, ht.getMapping())
}

func TestProcessHashtableCollision(t *testing.T) {
	ht := &processHashtable{}
	// hash all processes to the same number
	ht.extHash = func(*model.Process) uint64 {
		return 42
	}

	p1 := model.NewProcess("s1", []model.KeyValue{
		model.String("host", "google.com"),
	})
	p2 := model.NewProcess("s2", []model.KeyValue{
		model.String("host", "facebook.com"),
	})
	assert.Equal(t, "p1", ht.getKey(p1))
	assert.Equal(t, "p2", ht.getKey(p2))
}
<commit_msg>Test that processes with identical tags are deduped (#708)<commit_after>// Copyright (c) 2017 Uber Technologies, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// 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 json

import (
	"testing"

	"github.com/stretchr/testify/assert"

	"github.com/jaegertracing/jaeger/model"
)

func TestProcessHashtable(t *testing.T) {
	ht := &processHashtable{}

	p1 := model.NewProcess("s1", []model.KeyValue{
		model.String("ip", "1.2.3.4"),
		model.String("host", "google.com"),
	})
	// same process but with different order of tags
	p1dup := model.NewProcess("s1", []model.KeyValue{
		model.String("host", "google.com"),
		model.String("ip", "1.2.3.4"),
	})
	p2 := model.NewProcess("s2", []model.KeyValue{
		model.String("host", "facebook.com"),
	})

	assert.Equal(t, "p1", ht.getKey(p1))
	assert.Equal(t, "p1", ht.getKey(p1))
	assert.Equal(t, "p1", ht.getKey(p1dup))
	assert.Equal(t, "p2", ht.getKey(p2))

	expectedMapping := map[string]*model.Process{
		"p1": p1,
		"p2": p2,
	}
	assert.Equal(t, expectedMapping, ht.getMapping())
}

func TestProcessHashtableCollision(t *testing.T) {
	ht := &processHashtable{}
	// hash all processes to the same number
	ht.extHash = func(*model.Process) uint64 {
		return 42
	}

	p1 := model.NewProcess("s1", []model.KeyValue{
		model.String("host", "google.com"),
	})
	p2 := model.NewProcess("s2", []model.KeyValue{
		model.String("host", "facebook.com"),
	})
	assert.Equal(t, "p1", ht.getKey(p1))
	assert.Equal(t, "p2", ht.getKey(p2))
}
<|endoftext|>


<commit_before>package buildbox

import (
	"io"
	"net/http"
	"os"
	"path/filepath"
	"strings"
)

type Download struct {
	// The actual URL to get the file from
	URL string

	// The root directory of the download
	Destination string

	// The relative path that should be preserved in the download folder
	Path string
}

func StartDownload(quit chan string, download Download) {
	// If we're downloading a file with a path of "pkg/foo.txt" to a folder
	// called "pkg", we should merge the two paths together. So, instead of it
	// downloading to: destination/pkg/pkg/foo.txt, it will just download to
	// destination/pkg/foo.txt
	destinationPaths := strings.Split(download.Destination, string(os.PathSeparator))
	downloadPaths := strings.Split(download.Path, string(os.PathSeparator))

	for i := 0; i < len(downloadPaths); i += 100 {
		// If the last part of the destination path matches
		// this path in the download, then cut it out.
		lastIndex := len(destinationPaths) - 1

		// Break if we've gone too far.
		if lastIndex == -1 {
			break
		}

		lastPathInDestination := destinationPaths[lastIndex]
		if lastPathInDestination == downloadPaths[i] {
			destinationPaths = destinationPaths[:lastIndex]
		}
	}

	finalizedDestination := strings.Join(destinationPaths, string(os.PathSeparator))

	targetFile := filepath.Join(finalizedDestination, download.Path)
	targetDirectory, _ := filepath.Split(targetFile)

	// Show a nice message that we're starting to download the file
	Logger.Infof("Downloading %s to %s", download.URL, targetFile)

	// Start by downloading the file
	response, err := http.Get(download.URL)
	if err != nil {
		Logger.Errorf("Error while downloading %s (%T: %v)", download.URL, err, err)
		return
	}
	defer response.Body.Close()

	// Now make the folder for our file
	err = os.MkdirAll(targetDirectory, 0777)
	if err != nil {
		Logger.Errorf("Failed to create folder for %s (%T: %v)", targetFile, err, err)
		return
	}

	// Create a file to handle the file
	fileBuffer, err := os.Create(targetFile)
	if err != nil {
		Logger.Errorf("Failed to create file %s (%T: %v)", targetFile, err, err)
		return
	}
	defer fileBuffer.Close()

	// Copy the data to the file
	bytes, err := io.Copy(fileBuffer, response.Body)
	if err != nil {
		Logger.Errorf("Error when copying data %s (%T: %v)", download.URL, err, err)
		return
	}

	Logger.Infof("Successfully downloaded %s (%d bytes)", download.URL, bytes)

	// We can notify the channel that this routine has finished now
	quit <- "finished"
}
<commit_msg>Tweaked logging around downloading artifacts.<commit_after>package buildbox

import (
	"io"
	"net/http"
	"os"
	"path/filepath"
	"strings"
)

type Download struct {
	// The actual URL to get the file from
	URL string

	// The root directory of the download
	Destination string

	// The relative path that should be preserved in the download folder
	Path string
}

func StartDownload(quit chan string, download Download) {
	// If we're downloading a file with a path of "pkg/foo.txt" to a folder
	// called "pkg", we should merge the two paths together. So, instead of it
	// downloading to: destination/pkg/pkg/foo.txt, it will just download to
	// destination/pkg/foo.txt
	destinationPaths := strings.Split(download.Destination, string(os.PathSeparator))
	downloadPaths := strings.Split(download.Path, string(os.PathSeparator))

	for i := 0; i < len(downloadPaths); i += 100 {
		// If the last part of the destination path matches
		// this path in the download, then cut it out.
		lastIndex := len(destinationPaths) - 1

		// Break if we've gone too far.
		if lastIndex == -1 {
			break
		}

		lastPathInDestination := destinationPaths[lastIndex]
		if lastPathInDestination == downloadPaths[i] {
			destinationPaths = destinationPaths[:lastIndex]
		}
	}

	finalizedDestination := strings.Join(destinationPaths, string(os.PathSeparator))

	targetFile := filepath.Join(finalizedDestination, download.Path)
	targetDirectory, _ := filepath.Split(targetFile)

	// Show a nice message that we're starting to download the file
	Logger.Debugf("Downloading %s to %s", download.URL, targetFile)

	// Start by downloading the file
	response, err := http.Get(download.URL)
	if err != nil {
		Logger.Errorf("Error while downloading %s (%T: %v)", download.URL, err, err)
		return
	}
	defer response.Body.Close()

	// Now make the folder for our file
	err = os.MkdirAll(targetDirectory, 0777)
	if err != nil {
		Logger.Errorf("Failed to create folder for %s (%T: %v)", targetFile, err, err)
		return
	}

	// Create a file to handle the file
	fileBuffer, err := os.Create(targetFile)
	if err != nil {
		Logger.Errorf("Failed to create file %s (%T: %v)", targetFile, err, err)
		return
	}
	defer fileBuffer.Close()

	// Copy the data to the file
	bytes, err := io.Copy(fileBuffer, response.Body)
	if err != nil {
		Logger.Errorf("Error when copying data %s (%T: %v)", download.URL, err, err)
		return
	}

	Logger.Infof("Successfully downloaded %s (%d bytes)", download.Path, bytes)

	// We can notify the channel that this routine has finished now
	quit <- "finished"
}
<|endoftext|>


<commit_before>package anylisp

import (
	"fmt"
	"math/big"
	"os"
)

type Lister interface {
	Car() interface{}
	SetCar(v interface{}) interface{}
	Cdr() Lister
	SetCdr(v Lister) Lister
	Last() Lister
}

type List struct {
	car interface{}
	cdr Lister
}

type Inter interface {
	Add(x, y *big.Int) *big.Int
	Cmp(y *big.Int) (r int)
	Int64() int64
	Mul(x, y *big.Int) *big.Int
}

var (
	Ps_      Lister
	C_       Lister
	TempRoot Lister
)

func Parse(code string) {
	TempRoot = &List{"sx'", nil}
	Ps_ = &List{TempRoot, nil}
	C_ = &List{&List{TempRoot, nil}, nil}
	tok := ""
	for i := 0; i < len(code); i++ {
		if code[i] == ' ' || code[i] == '\t' || code[i] == '\n' {
			if tok == ")" {
				Assert(Ps_.Cdr() != nil, "Parse WTF! Too many )s")
				Ps_ = Ps_.Cdr()
			} else if len(tok) > 0 {
				var ls Lister
				if tok == "(" { // list
					ls = &List{nil, nil}
				} else if tok[0] == '[' && tok[len(tok)-1] == ']' { // number
					for j := 1; j < len(tok)-1; j++ {
						Assert(tok[j] == '-' || (tok[j] >= '0' && tok[j] <= '9') || (tok[j] >= 'a' && tok[j] <= 'f'),
							"Parse WTF! Bad character in number")
					}
					bi := new(big.Int)
					_, err := fmt.Sscanf(tok[1:len(tok)-1], "%x", bi)
					Assert(err == nil, "Parse WTF! Bad number")
					ls = &List{bi, nil}
				} else { // symbol
					ls = &List{tok, nil}
				}
				if Ps_.Car() == nil {
					Ps_.Cdr().Car().(Lister).SetCar(ls) // 1st token in list
				} else {
					Ps_.Car().(Lister).SetCdr(ls)
				}
				Ps_.SetCar(ls)
				if tok == "(" {
					Ps_ = &List{nil, Ps_}
				}
			}
			tok = ""
		} else {
			tok += string(code[i])
		}
	}
	Assert(Ps_.Cdr() == nil, "Parse WTF! Too few )s")
}

func Run() {
	for C_ != nil {
		f, ok := C_.Car().(Lister)
		Assert(ok, "WTF! Bad stack frame")
		e, ok := f.Car().(Lister)
		if !ok {
			fmt.Print("0 ")
			Ret(f.Car())
		} else {
			switch t := e.Car().(type) {
			case nil:
				Assert(false, "WTF! Can't call the empty list")
			case Inter:
				Assert(false, "WTF! Can't call an int")
			case Lister:
				Assert(false, "WTF! Can't call a list")
				// I kind of like the behavior below, but it causes strange error messages if there's a bug
				/*if f.Cdr() == nil {
					fmt.Println("a")
					C_ = &List{&List{t, nil}, C_}
				} else {
					fmt.Println("b")
					f.SetCar(f.Cdr().Car())
				}*/
			case string:
				switch t {
				case "sx'": // sx', arg, ret
					if f.Cdr() == nil {
						fmt.Print("{0 ")
						f.SetCdr(&List{e.Cdr(), nil})
					} else if f.Cdr().Car() != nil {
						fmt.Print("{1 ")
						C_ = &List{&List{NCarL(f, 1).Car(), nil}, C_}
						f.SetCdr(&List{NCarL(f, 1).Cdr(), nil})
					} else {
						fmt.Print("{2 ")
						Ret(f.Last().Car())
					}
				case "q'":
					fmt.Print("' ")
					Assert(e.Cdr() != nil, "WTF! Missing argument to quote")
					Ret(e.Cdr().Car())
				case ":^'", ":>'", ":|'": // op, ret
					if f.Cdr() == nil {
						fmt.Print(":1 ")
						Assert(e.Cdr() != nil, "WTF! Missing argument to "+t)
						C_ = &List{&List{e.Cdr().Car(), nil}, C_}
					} else if f.Cdr().Car() == nil {
						fmt.Print(":2 ")
						Ret(nil)
					} else {
						fmt.Print(":3 ")
						arg := NCarLA(f, 1, "WTF! "+t+" takes a list")
						switch t {
						case ":^'":
							Ret(arg.Car())
						case ":>'":
							Ret(arg.Cdr())
						case ":|'":
							Ret(arg.Last())
						}
					}
				case "lt'": // lt', arg, ret...
					if f.Cdr() == nil {
						fmt.Print("lt0 ")
						f.SetCdr(&List{e.Cdr(), nil})
					} else if f.Cdr().Car() != nil {
						fmt.Print("lt1 ")
						C_ = &List{&List{NCarL(f, 1).Car(), nil}, C_}
						f.Cdr().SetCar(NCarL(f, 1).Cdr())
					} else {
						fmt.Print("lt2 ")
						SetCdrA(NCdr(f, -2), f.Last().Car(), "WTF! Last argument to lt' must be a list")
						Ret(NCdr(f, 2))
					}
				case "?'":
					// ?', if part, then part, ret
					// ?', then part, nil, ret
					if e.Cdr() == nil {
						fmt.Print("?0 ")
						Ret(nil)
					} else if NCdr(f, 1) == nil {
						fmt.Print("?1 ")
						f.SetCdr(&List{e.Cdr(), &List{e.Cdr().Cdr(), nil}})
					} else if NCdr(f, 3) == nil {
						fmt.Print("?2 ")
						C_ = &List{&List{NCarL(f, 1).Car(), nil}, C_}
					} else if NCar(f, 2) == nil {
						fmt.Print("?3 ")
						Ret(NCar(f, 3))
					} else if NCar(f, 3) != nil {
						fmt.Print("?4 ")
						f.SetCdr(&List{NCarL(f, 1).Cdr(), &List{nil, nil}})
					} else if NCarL(f, 2).Cdr() == nil {
						fmt.Print("?5 ")
						Ret(nil)
					} else {
						fmt.Print("?6 ")
						f.SetCdr(&List{NCarL(f, 2).Cdr(), &List{NCarL(f, 2).Cdr().Cdr(), nil}})
					}
				case "+'", "*'": // op, arg, sum, ret
					if f.Cdr() == nil {
						fmt.Print("+0 ")
						var bi *big.Int
						if t == "+'" {
							bi = big.NewInt(0)
						} else if t == "*'" {
							bi = big.NewInt(1)
						}
						f.SetCdr(&List{e.Cdr(), &List{bi, nil}})
					} else if NCdr(f, 3) != nil {
						fmt.Print("+1 ")
						if t == "+'" {
							NCarI(f, 2).Add(NCarI(f, 2), NCarIA(f, 3, "WTF! +' takes numbers"))
						} else if t == "*'" {
							NCarI(f, 2).Mul(NCarI(f, 2), NCarIA(f, 3, "WTF! *' takes numbers"))
						}
						f.Cdr().Cdr().SetCdr(nil)
					} else if f.Cdr().Car() != nil {
						fmt.Print("+2 ")
						C_ = &List{&List{NCarL(f, 1).Car(), nil}, C_}
						f.Cdr().SetCar(NCarL(f, 1).Cdr())
					} else {
						fmt.Print("+3 ")
						Ret(NCar(f, 2))
					}
				case "pr'": // pr', arg, ret...
					if f.Cdr() == nil {
						fmt.Print("pr0 ")
						f.SetCdr(&List{e.Cdr(), nil})
					} else {
						if f.Cdr().Cdr() != nil {
							fmt.Print("pr1 ")
							SetCdrA(NCdr(f, -2), f.Last().Car(), "WTF! pr' takes a string")
						}
						if f.Cdr().Car() != nil {
							fmt.Print("pr2 ")
							C_ = &List{&List{NCarL(f, 1).Car(), nil}, C_}
							f.Cdr().SetCar(NCarL(f, 1).Cdr())
						} else {
							fmt.Print("pr3 ")
							s := make([]uint8, Len(f.Cdr().Cdr()))
							for i, arg := 0, f.Cdr().Cdr(); arg != nil; i, arg = i+1, arg.Cdr() {
								c, ok := arg.Car().(Inter)
								Assert(ok && c.Cmp(big.NewInt(-1)) == 1 && c.Cmp(big.NewInt(256)) == -1,
									"WTF! Bad byte passed to pr'")
								s[i] = uint8(c.Int64())
							}
							fmt.Print(string(s))
							Ret(f.Cdr().Cdr())
						}
					}
				default:
					Assert(false, "WTF! Can't call undefined function \""+t+"\"")
				}
			default:
				Assert(false, "WTF! Unrecognized function type (probably an interpreter bug)")
			}
		}
	}
}

func PrintTree(ls interface{}) {
	switch t := ls.(type) {
	case nil:
		fmt.Print("( ) ")
	case Inter:
		fmt.Printf("[%x] ", t)
	case Lister:
		fmt.Print("( ")
		for ls != nil {
			PrintTree(ls.(Lister).Car())
			ls = ls.(Lister).Cdr()
		}
		fmt.Print(") ")
	case string:
		fmt.Print(t + " ")
	default:
		Assert(false, "Unrecognized object in tree")
	}
}

func Ret(v interface{}) {
	if C_.Cdr() != nil {
		C_.Cdr().Car().(Lister).Last().SetCdr(&List{v, nil})
	}
	C_ = C_.Cdr()
}

func Len(ls Lister) int {
	if ls == nil {
		return 0
	}
	if ls.Cdr() == nil {
		return 1
	}
	return Len(ls.Cdr()) + 1
}

func NCar(ls Lister, n int) interface{} {
	nCdr := NCdr(ls, n)
	Assert(nCdr != nil, "WTF! Out of bounds when calling :'")
	return nCdr.Car()
}

func NCarL(ls Lister, n int) Lister {
	return NCarLA(ls, n, "WTF! Requested list element isn't a list")
}

func NCarLA(ls Lister, n int, msg string) Lister {
	nCar, ok := NCar(ls, n).(Lister)
	Assert(ok, msg)
	return nCar
}

func NCarI(ls Lister, n int) *big.Int {
	return NCarIA(ls, n, "WTF! Requested list element isn't an int")
}

func NCarIA(ls Lister, n int, msg string) *big.Int {
	nCar, ok := NCar(ls, n).(*big.Int)
	Assert(ok, msg)
	return nCar
}

func NCdr(ls Lister, n int) Lister {
	if ls == nil {
		return nil
	}
	if n > 0 {
		return NCdr(ls.Cdr(), n-1)
	}
	if n < 0 {
		return NCdr(ls, Len(ls)+n)
	}
	return ls
}

func SetCdrA(ls Lister, v interface{}, msg string) Lister {
	switch t := v.(type) {
	case nil:
		return ls.SetCdr(nil)
	case Lister:
		return ls.SetCdr(t)
	}
	Assert(false, msg)
	return nil
}

func (ls *List) Car() interface{} {
	return ls.car
}

func (ls *List) SetCar(v interface{}) interface{} {
	ls.car = v
	return v
}

func (ls *List) Cdr() Lister {
	return ls.cdr
}

func (ls *List) SetCdr(v Lister) Lister {
	ls.cdr = v
	return v
}

func (ls *List) Last() Lister {
	if ls.cdr == nil {
		return ls
	}
	return ls.cdr.Last()
}

func Assert(cond bool, msg string) {
	if !cond {
		fmt.Println(msg)
		os.Exit(2)
	}
}
<commit_msg>fix : debug numbering<commit_after>package anylisp

import (
	"fmt"
	"math/big"
	"os"
)

type Lister interface {
	Car() interface{}
	SetCar(v interface{}) interface{}
	Cdr() Lister
	SetCdr(v Lister) Lister
	Last() Lister
}

type List struct {
	car interface{}
	cdr Lister
}

type Inter interface {
	Add(x, y *big.Int) *big.Int
	Cmp(y *big.Int) (r int)
	Int64() int64
	Mul(x, y *big.Int) *big.Int
}

var (
	Ps_      Lister
	C_       Lister
	TempRoot Lister
)

func Parse(code string) {
	TempRoot = &List{"sx'", nil}
	Ps_ = &List{TempRoot, nil}
	C_ = &List{&List{TempRoot, nil}, nil}
	tok := ""
	for i := 0; i < len(code); i++ {
		if code[i] == ' ' || code[i] == '\t' || code[i] == '\n' {
			if tok == ")" {
				Assert(Ps_.Cdr() != nil, "Parse WTF! Too many )s")
				Ps_ = Ps_.Cdr()
			} else if len(tok) > 0 {
				var ls Lister
				if tok == "(" { // list
					ls = &List{nil, nil}
				} else if tok[0] == '[' && tok[len(tok)-1] == ']' { // number
					for j := 1; j < len(tok)-1; j++ {
						Assert(tok[j] == '-' || (tok[j] >= '0' && tok[j] <= '9') || (tok[j] >= 'a' && tok[j] <= 'f'),
							"Parse WTF! Bad character in number")
					}
					bi := new(big.Int)
					_, err := fmt.Sscanf(tok[1:len(tok)-1], "%x", bi)
					Assert(err == nil, "Parse WTF! Bad number")
					ls = &List{bi, nil}
				} else { // symbol
					ls = &List{tok, nil}
				}
				if Ps_.Car() == nil {
					Ps_.Cdr().Car().(Lister).SetCar(ls) // 1st token in list
				} else {
					Ps_.Car().(Lister).SetCdr(ls)
				}
				Ps_.SetCar(ls)
				if tok == "(" {
					Ps_ = &List{nil, Ps_}
				}
			}
			tok = ""
		} else {
			tok += string(code[i])
		}
	}
	Assert(Ps_.Cdr() == nil, "Parse WTF! Too few )s")
}

func Run() {
	for C_ != nil {
		f, ok := C_.Car().(Lister)
		Assert(ok, "WTF! Bad stack frame")
		e, ok := f.Car().(Lister)
		if !ok {
			fmt.Print("0 ")
			Ret(f.Car())
		} else {
			switch t := e.Car().(type) {
			case nil:
				Assert(false, "WTF! Can't call the empty list")
			case Inter:
				Assert(false, "WTF! Can't call an int")
			case Lister:
				Assert(false, "WTF! Can't call a list")
				// I kind of like the behavior below, but it causes strange error messages if there's a bug
				/*if f.Cdr() == nil {
					fmt.Println("a")
					C_ = &List{&List{t, nil}, C_}
				} else {
					fmt.Println("b")
					f.SetCar(f.Cdr().Car())
				}*/
			case string:
				switch t {
				case "sx'": // sx', arg, ret
					if f.Cdr() == nil {
						fmt.Print("{0 ")
						f.SetCdr(&List{e.Cdr(), nil})
					} else if f.Cdr().Car() != nil {
						fmt.Print("{1 ")
						C_ = &List{&List{NCarL(f, 1).Car(), nil}, C_}
						f.SetCdr(&List{NCarL(f, 1).Cdr(), nil})
					} else {
						fmt.Print("{2 ")
						Ret(f.Last().Car())
					}
				case "q'":
					fmt.Print("' ")
					Assert(e.Cdr() != nil, "WTF! Missing argument to quote")
					Ret(e.Cdr().Car())
				case ":^'", ":>'", ":|'": // op, ret
					if f.Cdr() == nil {
						fmt.Print(":0 ")
						Assert(e.Cdr() != nil, "WTF! Missing argument to "+t)
						C_ = &List{&List{e.Cdr().Car(), nil}, C_}
					} else if f.Cdr().Car() == nil {
						fmt.Print(":1 ")
						Ret(nil)
					} else {
						fmt.Print(":2 ")
						arg := NCarLA(f, 1, "WTF! "+t+" takes a list")
						switch t {
						case ":^'":
							Ret(arg.Car())
						case ":>'":
							Ret(arg.Cdr())
						case ":|'":
							Ret(arg.Last())
						}
					}
				case "lt'": // lt', arg, ret...
					if f.Cdr() == nil {
						fmt.Print("lt0 ")
						f.SetCdr(&List{e.Cdr(), nil})
					} else if f.Cdr().Car() != nil {
						fmt.Print("lt1 ")
						C_ = &List{&List{NCarL(f, 1).Car(), nil}, C_}
						f.Cdr().SetCar(NCarL(f, 1).Cdr())
					} else {
						fmt.Print("lt2 ")
						SetCdrA(NCdr(f, -2), f.Last().Car(), "WTF! Last argument to lt' must be a list")
						Ret(NCdr(f, 2))
					}
				case "?'":
					// ?', if part, then part, ret
					// ?', then part, nil, ret
					if e.Cdr() == nil {
						fmt.Print("?0 ")
						Ret(nil)
					} else if NCdr(f, 1) == nil {
						fmt.Print("?1 ")
						f.SetCdr(&List{e.Cdr(), &List{e.Cdr().Cdr(), nil}})
					} else if NCdr(f, 3) == nil {
						fmt.Print("?2 ")
						C_ = &List{&List{NCarL(f, 1).Car(), nil}, C_}
					} else if NCar(f, 2) == nil {
						fmt.Print("?3 ")
						Ret(NCar(f, 3))
					} else if NCar(f, 3) != nil {
						fmt.Print("?4 ")
						f.SetCdr(&List{NCarL(f, 1).Cdr(), &List{nil, nil}})
					} else if NCarL(f, 2).Cdr() == nil {
						fmt.Print("?5 ")
						Ret(nil)
					} else {
						fmt.Print("?6 ")
						f.SetCdr(&List{NCarL(f, 2).Cdr(), &List{NCarL(f, 2).Cdr().Cdr(), nil}})
					}
				case "+'", "*'": // op, arg, sum, ret
					if f.Cdr() == nil {
						fmt.Print("+0 ")
						var bi *big.Int
						if t == "+'" {
							bi = big.NewInt(0)
						} else if t == "*'" {
							bi = big.NewInt(1)
						}
						f.SetCdr(&List{e.Cdr(), &List{bi, nil}})
					} else if NCdr(f, 3) != nil {
						fmt.Print("+1 ")
						if t == "+'" {
							NCarI(f, 2).Add(NCarI(f, 2), NCarIA(f, 3, "WTF! +' takes numbers"))
						} else if t == "*'" {
							NCarI(f, 2).Mul(NCarI(f, 2), NCarIA(f, 3, "WTF! *' takes numbers"))
						}
						f.Cdr().Cdr().SetCdr(nil)
					} else if f.Cdr().Car() != nil {
						fmt.Print("+2 ")
						C_ = &List{&List{NCarL(f, 1).Car(), nil}, C_}
						f.Cdr().SetCar(NCarL(f, 1).Cdr())
					} else {
						fmt.Print("+3 ")
						Ret(NCar(f, 2))
					}
				case "pr'": // pr', arg, ret...
					if f.Cdr() == nil {
						fmt.Print("pr0 ")
						f.SetCdr(&List{e.Cdr(), nil})
					} else {
						if f.Cdr().Cdr() != nil {
							fmt.Print("pr1 ")
							SetCdrA(NCdr(f, -2), f.Last().Car(), "WTF! pr' takes a string")
						}
						if f.Cdr().Car() != nil {
							fmt.Print("pr2 ")
							C_ = &List{&List{NCarL(f, 1).Car(), nil}, C_}
							f.Cdr().SetCar(NCarL(f, 1).Cdr())
						} else {
							fmt.Print("pr3 ")
							s := make([]uint8, Len(f.Cdr().Cdr()))
							for i, arg := 0, f.Cdr().Cdr(); arg != nil; i, arg = i+1, arg.Cdr() {
								c, ok := arg.Car().(Inter)
								Assert(ok && c.Cmp(big.NewInt(-1)) == 1 && c.Cmp(big.NewInt(256)) == -1,
									"WTF! Bad byte passed to pr'")
								s[i] = uint8(c.Int64())
							}
							fmt.Print(string(s))
							Ret(f.Cdr().Cdr())
						}
					}
				default:
					Assert(false, "WTF! Can't call undefined function \""+t+"\"")
				}
			default:
				Assert(false, "WTF! Unrecognized function type (probably an interpreter bug)")
			}
		}
	}
}

func PrintTree(ls interface{}) {
	switch t := ls.(type) {
	case nil:
		fmt.Print("( ) ")
	case Inter:
		fmt.Printf("[%x] ", t)
	case Lister:
		fmt.Print("( ")
		for ls != nil {
			PrintTree(ls.(Lister).Car())
			ls = ls.(Lister).Cdr()
		}
		fmt.Print(") ")
	case string:
		fmt.Print(t + " ")
	default:
		Assert(false, "Unrecognized object in tree")
	}
}

func Ret(v interface{}) {
	if C_.Cdr() != nil {
		C_.Cdr().Car().(Lister).Last().SetCdr(&List{v, nil})
	}
	C_ = C_.Cdr()
}

func Len(ls Lister) int {
	if ls == nil {
		return 0
	}
	if ls.Cdr() == nil {
		return 1
	}
	return Len(ls.Cdr()) + 1
}

func NCar(ls Lister, n int) interface{} {
	nCdr := NCdr(ls, n)
	Assert(nCdr != nil, "WTF! Out of bounds when calling :'")
	return nCdr.Car()
}

func NCarL(ls Lister, n int) Lister {
	return NCarLA(ls, n, "WTF! Requested list element isn't a list")
}

func NCarLA(ls Lister, n int, msg string) Lister {
	nCar, ok := NCar(ls, n).(Lister)
	Assert(ok, msg)
	return nCar
}

func NCarI(ls Lister, n int) *big.Int {
	return NCarIA(ls, n, "WTF! Requested list element isn't an int")
}

func NCarIA(ls Lister, n int, msg string) *big.Int {
	nCar, ok := NCar(ls, n).(*big.Int)
	Assert(ok, msg)
	return nCar
}

func NCdr(ls Lister, n int) Lister {
	if ls == nil {
		return nil
	}
	if n > 0 {
		return NCdr(ls.Cdr(), n-1)
	}
	if n < 0 {
		return NCdr(ls, Len(ls)+n)
	}
	return ls
}

func SetCdrA(ls Lister, v interface{}, msg string) Lister {
	switch t := v.(type) {
	case nil:
		return ls.SetCdr(nil)
	case Lister:
		return ls.SetCdr(t)
	}
	Assert(false, msg)
	return nil
}

func (ls *List) Car() interface{} {
	return ls.car
}

func (ls *List) SetCar(v interface{}) interface{} {
	ls.car = v
	return v
}

func (ls *List) Cdr() Lister {
	return ls.cdr
}

func (ls *List) SetCdr(v Lister) Lister {
	ls.cdr = v
	return v
}

func (ls *List) Last() Lister {
	if ls.cdr == nil {
		return ls
	}
	return ls.cdr.Last()
}

func Assert(cond bool, msg string) {
	if !cond {
		fmt.Println(msg)
		os.Exit(2)
	}
}
<|endoftext|>