<commit_before>package agent

import (
	"context"
	"encoding/hex"
	"encoding/json"
	"errors"
	"io"
	"time"

	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/status"

	pb "github.com/fnproject/fn/api/agent/grpc"
	"github.com/fnproject/fn/api/common"
	"github.com/fnproject/fn/api/models"
	pool "github.com/fnproject/fn/api/runnerpool"
	"github.com/fnproject/fn/grpcutil"
	"github.com/sirupsen/logrus"
)

var (
	ErrorRunnerClosed    = errors.New("Runner is closed")
	ErrorPureRunnerNoEOF = errors.New("Purerunner missing EOF response")
)

const (
	// max buffer size for grpc data messages, 10K
	MaxDataChunk = 10 * 1024
)

type gRPCRunner struct {
	shutWg  *common.WaitGroup
	address string
	conn    *grpc.ClientConn
	client  pb.RunnerProtocolClient
}

func SecureGRPCRunnerFactory(addr, runnerCertCN string, pki *pool.PKIData) (pool.Runner, error) {
	conn, client, err := runnerConnection(addr, runnerCertCN, pki)
	if err != nil {
		return nil, err
	}

	return &gRPCRunner{
		shutWg:  common.NewWaitGroup(),
		address: addr,
		conn:    conn,
		client:  client,
	}, nil
}

func (r *gRPCRunner) Close(context.Context) error {
	r.shutWg.CloseGroup()
	return r.conn.Close()
}

func runnerConnection(address, runnerCertCN string, pki *pool.PKIData) (*grpc.ClientConn, pb.RunnerProtocolClient, error) {
	ctx := context.Background()

	var creds credentials.TransportCredentials
	if pki != nil {
		var err error
		creds, err = grpcutil.CreateCredentials(pki.Cert, pki.Key, pki.Ca, runnerCertCN)
		if err != nil {
			logrus.WithError(err).Error("Unable to create credentials to connect to runner node")
			return nil, nil, err
		}
	}

	// we want to set a very short timeout to fail-fast if something goes wrong
	conn, err := grpcutil.DialWithBackoff(ctx, address, creds, 100*time.Millisecond, grpc.DefaultBackoffConfig)
	if err != nil {
		logrus.WithError(err).Error("Unable to connect to runner node")
	}

	protocolClient := pb.NewRunnerProtocolClient(conn)
	logrus.WithField("runner_addr", address).Info("Connected to runner")

	return conn, protocolClient, nil
}

func (r *gRPCRunner) Address() string {
	return r.address
}

func isRetriable(err error) bool {
	// A formal API error returned from pure-runner
	if models.GetAPIErrorCode(err) == models.GetAPIErrorCode(models.ErrCallTimeoutServerBusy) {
		return true
	}
	if err != nil {
		// engagement/recv errors could also be a 503.
		st := status.Convert(err)
		if int(st.Code()) == models.GetAPIErrorCode(models.ErrCallTimeoutServerBusy) {
			return true
		}
	}
	return false
}

func (r *gRPCRunner) TryExec(ctx context.Context, call pool.RunnerCall) (bool, error) {
	logrus.WithField("runner_addr", r.address).Debug("Attempting to place call")
	if !r.shutWg.AddSession(1) {
		// try another runner if this one is closed.
		return false, ErrorRunnerClosed
	}
	defer r.shutWg.DoneSession()

	// extract the call's model data to pass on to the pure runner
	modelJSON, err := json.Marshal(call.Model())
	if err != nil {
		logrus.WithError(err).Error("Failed to encode model as JSON")
		// If we can't encode the model, no runner will ever be able to run this. Give up.
		return true, err
	}

	runnerConnection, err := r.client.Engage(ctx)
	if err != nil {
		logrus.WithError(err).Error("Unable to create client to runner node")
		// Try on next runner
		return false, err
	}

	// After this point, we assume "COMMITTED" unless pure runner
	// send explicit NACK
	err = runnerConnection.Send(&pb.ClientMsg{Body: &pb.ClientMsg_Try{Try: &pb.TryCall{
		ModelsCallJson: string(modelJSON),
		SlotHashId:     hex.EncodeToString([]byte(call.SlotHashId())),
	}}})
	if err != nil {
		logrus.WithError(err).Error("Failed to send message to runner node")
		return true, err
	}

	recvDone := make(chan error, 1)

	go receiveFromRunner(runnerConnection, call, recvDone)
	go sendToRunner(runnerConnection, call)

	select {
	case <-ctx.Done():
		logrus.Infof("Engagement Context ended ctxErr=%v", ctx.Err())
		return true, ctx.Err()
	case recvErr := <-recvDone:
		return !isRetriable(recvErr), recvErr
	}
}

func sendToRunner(protocolClient pb.RunnerProtocol_EngageClient, call pool.RunnerCall) {
	bodyReader := call.RequestBody()
	writeBuffer := make([]byte, MaxDataChunk)

	// IMPORTANT: IO Read below can fail in multiple go-routine cases (in retry
	// case especially if receiveFromRunner go-routine receives a NACK while sendToRunner is
	// already blocked on a read) or in the case of reading the http body multiple times (retries.)
	// Normally http.Request.Body can be read once. However runner_client users should implement/add
	// http.Request.GetBody() function and cache the body content in the request.
	// See lb_agent setRequestGetBody() which handles this. With GetBody installed,
	// the 'Read' below is an actually non-blocking operation since GetBody() should hand out
	// a new instance of io.ReadCloser() that allows repetitive reads on the http body.
	for {
		// WARNING: blocking read.
		n, err := bodyReader.Read(writeBuffer)
		if err != nil && err != io.EOF {
			logrus.WithError(err).Error("Failed to receive data from http client body")
		}

		// any IO error or n == 0 is an EOF for pure-runner
		isEOF := err != nil || n == 0
		data := writeBuffer[:n]

		logrus.Debugf("Sending %d bytes of data isEOF=%v to runner", n, isEOF)
		sendErr := protocolClient.Send(&pb.ClientMsg{
			Body: &pb.ClientMsg_Data{
				Data: &pb.DataFrame{
					Data: data,
					Eof:  isEOF,
				},
			},
		})
		if sendErr != nil {
			logrus.WithError(sendErr).Errorf("Failed to send data frame size=%d isEOF=%v", n, isEOF)
			return
		}
		if isEOF {
			return
		}
	}
}

func parseError(msg *pb.CallFinished) error {
	if msg.GetSuccess() {
		return nil
	}
	eCode := msg.GetErrorCode()
	eStr := msg.GetErrorStr()
	if eStr == "" {
		eStr = "Unknown Error From Pure Runner"
	}
	return models.NewAPIError(int(eCode), errors.New(eStr))
}

func tryQueueError(err error, done chan error) {
	select {
	case done <- err:
	default:
	}
}

func receiveFromRunner(protocolClient pb.RunnerProtocol_EngageClient, c pool.RunnerCall, done chan error) {
	w := c.ResponseWriter()
	defer close(done)

	isPartialWrite := false

DataLoop:
	for {
		msg, err := protocolClient.Recv()
		if err != nil {
			logrus.WithError(err).Info("Receive error from runner")
			tryQueueError(err, done)
			return
		}

		switch body := msg.Body.(type) {

		// Process HTTP header/status message. This may not arrive depending on
		// pure runners behavior. (Eg. timeout & no IO received from function)
		case *pb.RunnerMsg_ResultStart:
			switch meta := body.ResultStart.Meta.(type) {
			case *pb.CallResultStart_Http:
				logrus.Debugf("Received meta http result from runner Status=%v", meta.Http.StatusCode)
				for _, header := range meta.Http.Headers {
					w.Header().Set(header.Key, header.Value)
				}
				if meta.Http.StatusCode > 0 {
					w.WriteHeader(int(meta.Http.StatusCode))
				}
			default:
				logrus.Errorf("Unhandled meta type in start message: %v", meta)
			}

		// May arrive if function has output. We ignore EOF.
		case *pb.RunnerMsg_Data:
			logrus.Debugf("Received data from runner len=%d isEOF=%v", len(body.Data.Data), body.Data.Eof)
			if !isPartialWrite {
				// WARNING: blocking write
				n, err := w.Write(body.Data.Data)
				if n != len(body.Data.Data) {
					isPartialWrite = true
					logrus.WithError(err).Infof("Failed to write full response (%d of %d) to client", n, len(body.Data.Data))
					if err == nil {
						err = io.ErrShortWrite
					}
					tryQueueError(err, done)
				}
			}

		// Finish messages required for finish/finalize the processing.
		case *pb.RunnerMsg_Finished:
			logrus.Infof("Call finished Success=%v %v", body.Finished.Success, body.Finished.Details)
			if !body.Finished.Success {
				err := parseError(body.Finished)
				tryQueueError(err, done)
			}
			break DataLoop

		default:
			logrus.Error("Ignoring unknown message type %T from runner, possible client/server mismatch", body)
		}
	}

	// There should be an EOF following the last packet
	for {
		msg, err := protocolClient.Recv()
		if err == io.EOF {
			break
		}
		if err != nil {
			logrus.WithError(err).Infof("Call Waiting EOF received error")
			tryQueueError(err, done)
			break
		}

		switch body := msg.Body.(type) {
		default:
			logrus.Infof("Call Waiting EOF ignoring message %T", body)
		}
		tryQueueError(ErrorPureRunnerNoEOF, done)
	}
}
<commit_msg>fn: LB agent: EOF from runner is normal in nack cases (#1032)<commit_after>package agent

import (
	"context"
	"encoding/hex"
	"encoding/json"
	"errors"
	"io"
	"time"

	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/status"

	pb "github.com/fnproject/fn/api/agent/grpc"
	"github.com/fnproject/fn/api/common"
	"github.com/fnproject/fn/api/models"
	pool "github.com/fnproject/fn/api/runnerpool"
	"github.com/fnproject/fn/grpcutil"
	"github.com/sirupsen/logrus"
)

var (
	ErrorRunnerClosed    = errors.New("Runner is closed")
	ErrorPureRunnerNoEOF = errors.New("Purerunner missing EOF response")
)

const (
	// max buffer size for grpc data messages, 10K
	MaxDataChunk = 10 * 1024
)

type gRPCRunner struct {
	shutWg  *common.WaitGroup
	address string
	conn    *grpc.ClientConn
	client  pb.RunnerProtocolClient
}

func SecureGRPCRunnerFactory(addr, runnerCertCN string, pki *pool.PKIData) (pool.Runner, error) {
	conn, client, err := runnerConnection(addr, runnerCertCN, pki)
	if err != nil {
		return nil, err
	}

	return &gRPCRunner{
		shutWg:  common.NewWaitGroup(),
		address: addr,
		conn:    conn,
		client:  client,
	}, nil
}

func (r *gRPCRunner) Close(context.Context) error {
	r.shutWg.CloseGroup()
	return r.conn.Close()
}

func runnerConnection(address, runnerCertCN string, pki *pool.PKIData) (*grpc.ClientConn, pb.RunnerProtocolClient, error) {
	ctx := context.Background()

	var creds credentials.TransportCredentials
	if pki != nil {
		var err error
		creds, err = grpcutil.CreateCredentials(pki.Cert, pki.Key, pki.Ca, runnerCertCN)
		if err != nil {
			logrus.WithError(err).Error("Unable to create credentials to connect to runner node")
			return nil, nil, err
		}
	}

	// we want to set a very short timeout to fail-fast if something goes wrong
	conn, err := grpcutil.DialWithBackoff(ctx, address, creds, 100*time.Millisecond, grpc.DefaultBackoffConfig)
	if err != nil {
		logrus.WithError(err).Error("Unable to connect to runner node")
	}

	protocolClient := pb.NewRunnerProtocolClient(conn)
	logrus.WithField("runner_addr", address).Info("Connected to runner")

	return conn, protocolClient, nil
}

func (r *gRPCRunner) Address() string {
	return r.address
}

func isRetriable(err error) bool {
	// A formal API error returned from pure-runner
	if models.GetAPIErrorCode(err) == models.GetAPIErrorCode(models.ErrCallTimeoutServerBusy) {
		return true
	}
	if err != nil {
		// engagement/recv errors could also be a 503.
		st := status.Convert(err)
		if int(st.Code()) == models.GetAPIErrorCode(models.ErrCallTimeoutServerBusy) {
			return true
		}
	}
	return false
}

func (r *gRPCRunner) TryExec(ctx context.Context, call pool.RunnerCall) (bool, error) {
	logrus.WithField("runner_addr", r.address).Debug("Attempting to place call")
	if !r.shutWg.AddSession(1) {
		// try another runner if this one is closed.
		return false, ErrorRunnerClosed
	}
	defer r.shutWg.DoneSession()

	// extract the call's model data to pass on to the pure runner
	modelJSON, err := json.Marshal(call.Model())
	if err != nil {
		logrus.WithError(err).Error("Failed to encode model as JSON")
		// If we can't encode the model, no runner will ever be able to run this. Give up.
		return true, err
	}

	runnerConnection, err := r.client.Engage(ctx)
	if err != nil {
		logrus.WithError(err).Error("Unable to create client to runner node")
		// Try on next runner
		return false, err
	}

	// After this point, we assume "COMMITTED" unless pure runner
	// send explicit NACK
	err = runnerConnection.Send(&pb.ClientMsg{Body: &pb.ClientMsg_Try{Try: &pb.TryCall{
		ModelsCallJson: string(modelJSON),
		SlotHashId:     hex.EncodeToString([]byte(call.SlotHashId())),
	}}})
	if err != nil {
		logrus.WithError(err).Error("Failed to send message to runner node")
		return true, err
	}

	recvDone := make(chan error, 1)

	go receiveFromRunner(runnerConnection, call, recvDone)
	go sendToRunner(runnerConnection, call)

	select {
	case <-ctx.Done():
		logrus.Infof("Engagement Context ended ctxErr=%v", ctx.Err())
		return true, ctx.Err()
	case recvErr := <-recvDone:
		return !isRetriable(recvErr), recvErr
	}
}

func sendToRunner(protocolClient pb.RunnerProtocol_EngageClient, call pool.RunnerCall) {
	bodyReader := call.RequestBody()
	writeBuffer := make([]byte, MaxDataChunk)

	// IMPORTANT: IO Read below can fail in multiple go-routine cases (in retry
	// case especially if receiveFromRunner go-routine receives a NACK while sendToRunner is
	// already blocked on a read) or in the case of reading the http body multiple times (retries.)
	// Normally http.Request.Body can be read once. However runner_client users should implement/add
	// http.Request.GetBody() function and cache the body content in the request.
	// See lb_agent setRequestGetBody() which handles this. With GetBody installed,
	// the 'Read' below is an actually non-blocking operation since GetBody() should hand out
	// a new instance of io.ReadCloser() that allows repetitive reads on the http body.
	for {
		// WARNING: blocking read.
		n, err := bodyReader.Read(writeBuffer)
		if err != nil && err != io.EOF {
			logrus.WithError(err).Error("Failed to receive data from http client body")
		}

		// any IO error or n == 0 is an EOF for pure-runner
		isEOF := err != nil || n == 0
		data := writeBuffer[:n]

		logrus.Debugf("Sending %d bytes of data isEOF=%v to runner", n, isEOF)
		sendErr := protocolClient.Send(&pb.ClientMsg{
			Body: &pb.ClientMsg_Data{
				Data: &pb.DataFrame{
					Data: data,
					Eof:  isEOF,
				},
			},
		})
		if sendErr != nil {
			// It's often normal to receive an EOF here as we optimistically start sending body until a NACK
			// from the runner. Let's ignore EOF and rely on recv side to catch premature EOF.
			if sendErr != io.EOF {
				logrus.WithError(sendErr).Errorf("Failed to send data frame size=%d isEOF=%v", n, isEOF)
			}
			return
		}
		if isEOF {
			return
		}
	}
}

func parseError(msg *pb.CallFinished) error {
	if msg.GetSuccess() {
		return nil
	}
	eCode := msg.GetErrorCode()
	eStr := msg.GetErrorStr()
	if eStr == "" {
		eStr = "Unknown Error From Pure Runner"
	}
	return models.NewAPIError(int(eCode), errors.New(eStr))
}

func tryQueueError(err error, done chan error) {
	select {
	case done <- err:
	default:
	}
}

func receiveFromRunner(protocolClient pb.RunnerProtocol_EngageClient, c pool.RunnerCall, done chan error) {
	w := c.ResponseWriter()
	defer close(done)

	isPartialWrite := false

DataLoop:
	for {
		msg, err := protocolClient.Recv()
		if err != nil {
			logrus.WithError(err).Info("Receive error from runner")
			tryQueueError(err, done)
			return
		}

		switch body := msg.Body.(type) {

		// Process HTTP header/status message. This may not arrive depending on
		// pure runners behavior. (Eg. timeout & no IO received from function)
		case *pb.RunnerMsg_ResultStart:
			switch meta := body.ResultStart.Meta.(type) {
			case *pb.CallResultStart_Http:
				logrus.Debugf("Received meta http result from runner Status=%v", meta.Http.StatusCode)
				for _, header := range meta.Http.Headers {
					w.Header().Set(header.Key, header.Value)
				}
				if meta.Http.StatusCode > 0 {
					w.WriteHeader(int(meta.Http.StatusCode))
				}
			default:
				logrus.Errorf("Unhandled meta type in start message: %v", meta)
			}

		// May arrive if function has output. We ignore EOF.
		case *pb.RunnerMsg_Data:
			logrus.Debugf("Received data from runner len=%d isEOF=%v", len(body.Data.Data), body.Data.Eof)
			if !isPartialWrite {
				// WARNING: blocking write
				n, err := w.Write(body.Data.Data)
				if n != len(body.Data.Data) {
					isPartialWrite = true
					logrus.WithError(err).Infof("Failed to write full response (%d of %d) to client", n, len(body.Data.Data))
					if err == nil {
						err = io.ErrShortWrite
					}
					tryQueueError(err, done)
				}
			}

		// Finish messages required for finish/finalize the processing.
		case *pb.RunnerMsg_Finished:
			logrus.Infof("Call finished Success=%v %v", body.Finished.Success, body.Finished.Details)
			if !body.Finished.Success {
				err := parseError(body.Finished)
				tryQueueError(err, done)
			}
			break DataLoop

		default:
			logrus.Error("Ignoring unknown message type %T from runner, possible client/server mismatch", body)
		}
	}

	// There should be an EOF following the last packet
	for {
		msg, err := protocolClient.Recv()
		if err == io.EOF {
			break
		}
		if err != nil {
			logrus.WithError(err).Infof("Call Waiting EOF received error")
			tryQueueError(err, done)
			break
		}

		switch body := msg.Body.(type) {
		default:
			logrus.Infof("Call Waiting EOF ignoring message %T", body)
		}
		tryQueueError(ErrorPureRunnerNoEOF, done)
	}
}
<|endoftext|>


<commit_before>/*
 * Copyright 2018 Arm Limited.
 * SPDX-License-Identifier: Apache-2.0
 *
 * 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 core

import (
	"fmt"
	"reflect"
	"strings"

	"github.com/google/blueprint/proptools"

	"github.com/ARM-software/bob-build/utils"
)

// featurePropertyName returns name of feature. Name needs to start from capital letter because
// this is how it works in go exported/unexported properties
// e.g. Android, Foo_bar
func featurePropertyName(name string) string {
	result := strings.ToLower(name) // e.g. android, foo_bar
	if len(name) > 0 {
		return strings.ToUpper(name[0:1]) + result[1:]
	}
	return result
}

// Features must be embedded in each modules property structure to support the
// use of features in the module. The feature must be initialised with a call
// to Init().
type Features struct {
	// 'BlueprintEmbed' is a special case in Blueprint which makes it interpret
	// a runtime-generated type as being embedded in its parent struct.
	BlueprintEmbed interface{}
}

type singleFeature struct {
	BlueprintEmbed interface{}
}

func typesOf(list ...interface{}) []reflect.Type {
	types := make([]reflect.Type, len(list))
	for i, element := range list {
		types[i] = reflect.TypeOf(element)
	}
	return types
}

// Init generates and initializes a struct containing a field of type
// 'propsType' for every available feature. 'propsType' will be constructed
// from list of types. By constructed we mean properties of each
// type will be merged together. It is important to set here
// every available feature not only enabled ones, because blueprint will
// fail during reading .bp files. The generated object is embedded
// in each module types' properties instance, and is used by Blueprint to
// decide what properties can be set inside features in each module type.
//
// An example generated type:
// type BlueprintEmbedType struct {
//         Debug PropsType
//         Enable_something PropsType
//         Some_other_feature PropsType
// }
// Name of each property in this struct is custom feature name.
// Blueprint will inflate this structure with data read from .bp files.
// Only exported properties can be set so property name MUST start from capital letter.
func (f *Features) Init(availableFeatures []string, list ...interface{}) {
	if len(list) == 0 {
		panic("List can't be empty")
	}

	propsType := coalesceTypes(typesOf(list...)...)
	fields := make([]reflect.StructField, len(availableFeatures))

	for i, featureName := range availableFeatures {
		fields[i] = reflect.StructField{
			Name: featurePropertyName(featureName),
			Type: reflect.TypeOf(singleFeature{}),
		}
	}

	bpFeatureStruct := reflect.StructOf(fields)
	instancePtr := reflect.New(bpFeatureStruct)
	f.BlueprintEmbed = instancePtr.Interface()

	instance := reflect.Indirect(instancePtr)
	for i, _ := range availableFeatures {
		propsInFeature := instance.Field(i).Addr().Interface().(*singleFeature)
		propsInFeature.BlueprintEmbed = reflect.New(propsType).Interface()
	}

}

// coalesceTypes will squash multiple types to new type. This has different result
// than Go composition of structs.
//
// Example (go composition):
// type compositeStruct struct {
//     testPropsGroupA
//     testPropsGroupB
// }
// Debug print:
// {
//     testPropsGroupA: core.testPropsGroupA
//     {
//       Field_a: string
//       Field_c: string
//       Field_f: string
//     }
//     testPropsGroupB: core.testPropsGroupB
//     {
//       Field_b: string
//     }
// }
//
// Example for: coalesceTypes([]reflect.Type{
//    testPropsGroupA{},
//    testPropsGroupB{},
// })
// Debug print:
// {
//     Field_a: string
//     Field_c: string
//     Field_f: string
//     Field_b: string
// }
func coalesceTypes(list ...reflect.Type) reflect.Type {
	if len(list) == 0 {
		panic("List can't be empty")
	}
	if len(list) == 1 {
		return list[0]
	}

	fieldsKeys := map[string]bool{}
	fields := []reflect.StructField{}

	for _, elementType := range list {
		for i := 0; i < elementType.NumField(); i++ {
			field := elementType.Field(i)
			fieldName := field.Name
			if _, ok := fieldsKeys[fieldName]; ok {
				panic(fmt.Sprintf("Name collision: '%v'\n", fieldName))
			} else {
				fieldsKeys[fieldName] = true
				fields = append(fields, field)
			}
		}
	}

	return reflect.StructOf(fields)
}

// AppendProps merges properties from BlueprintEmbed to dst, but only for enabled features
// expect that Features are inited (before using this function we should call Features.Init)
// expect that properties.Features should contain all available features (whenever disabled/enabled)
func (f *Features) AppendProps(dst []interface{}, properties *configProperties) error {
	// featuresData is struct created in Features.Init function
	featuresData := reflect.ValueOf(f.BlueprintEmbed).Elem()

	for _, featureKey := range utils.SortedKeysBoolMap(properties.Features) {
		if properties.Features[featureKey] { // Check the feature is enabled
			// Features are matched like "Feature_name" - feature structure
			featureFieldName := featurePropertyName(featureKey)
			featureStruct := featuresData.FieldByName(featureFieldName)
			if !featureStruct.IsValid() {
				panic(fmt.Sprintf("Field returned for property %s isn't valid\n", featureFieldName))
			}
			// AppendProperties expects a pointer to a struct.
			featureStructPointer := featureStruct.FieldByName("BlueprintEmbed").Interface()

			// If featureProps is nil then we've determined that we can skip this,
			// so avoid calling AppendProperties
			if featureStructPointer != nil {
				err := proptools.AppendMatchingProperties(dst, featureStructPointer, nil)
				if err != nil {
					return err
				}
			}
		}
	}
	return nil
}
<commit_msg>Fix creation of embedded fields with reflection<commit_after>/*
 * Copyright 2018 Arm Limited.
 * SPDX-License-Identifier: Apache-2.0
 *
 * 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 core

import (
	"fmt"
	"reflect"
	"strings"

	"github.com/google/blueprint/proptools"

	"github.com/ARM-software/bob-build/utils"
)

// featurePropertyName returns name of feature. Name needs to start from capital letter because
// this is how it works in go exported/unexported properties
// e.g. Android, Foo_bar
func featurePropertyName(name string) string {
	result := strings.ToLower(name) // e.g. android, foo_bar
	if len(name) > 0 {
		return strings.ToUpper(name[0:1]) + result[1:]
	}
	return result
}

// Features must be embedded in each modules property structure to support the
// use of features in the module. The feature must be initialised with a call
// to Init().
type Features struct {
	// 'BlueprintEmbed' is a special case in Blueprint which makes it interpret
	// a runtime-generated type as being embedded in its parent struct.
	BlueprintEmbed interface{}
}

type singleFeature struct {
	BlueprintEmbed interface{}
}

func typesOf(list ...interface{}) []reflect.Type {
	types := make([]reflect.Type, len(list))
	for i, element := range list {
		types[i] = reflect.TypeOf(element)
	}
	return types
}

// Init generates and initializes a struct containing a field of type
// 'propsType' for every available feature. 'propsType' will be constructed
// from list of types. By constructed we mean properties of each
// type will be merged together. It is important to set here
// every available feature not only enabled ones, because blueprint will
// fail during reading .bp files. The generated object is embedded
// in each module types' properties instance, and is used by Blueprint to
// decide what properties can be set inside features in each module type.
//
// An example generated type:
// type BlueprintEmbedType struct {
//         Debug PropsType
//         Enable_something PropsType
//         Some_other_feature PropsType
// }
// Name of each property in this struct is custom feature name.
// Blueprint will inflate this structure with data read from .bp files.
// Only exported properties can be set so property name MUST start from capital letter.
func (f *Features) Init(availableFeatures []string, list ...interface{}) {
	if len(list) == 0 {
		panic("List can't be empty")
	}

	propsType := coalesceTypes(typesOf(list...)...)
	fields := make([]reflect.StructField, len(availableFeatures))

	for i, featureName := range availableFeatures {
		fields[i] = reflect.StructField{
			Name: featurePropertyName(featureName),
			Type: reflect.TypeOf(singleFeature{}),
		}
	}

	bpFeatureStruct := reflect.StructOf(fields)
	instancePtr := reflect.New(bpFeatureStruct)
	f.BlueprintEmbed = instancePtr.Interface()

	instance := reflect.Indirect(instancePtr)
	for i, _ := range availableFeatures {
		propsInFeature := instance.Field(i).Addr().Interface().(*singleFeature)
		propsInFeature.BlueprintEmbed = reflect.New(propsType).Interface()
	}

}

// Some Go versions (e.g. 1.8) ignore the 'Anonymous' property, and require
// that the name is empty to create an anonymous field. This is fixed in 1.10,
// but 1.10 now errors if the name is empty. So for coalesceTypes() to create
// anonymous fields, we need to find out what this version's behavior is.
func isEmptyNameRequiredForAnonymousField() bool {
	// Try creating a new type which contains an embedded struct. If the
	// struct's 'Anonymous' property is dropped when we create an instance
	// of it, we need to use an empty 'Name' field instead.
	type TestEmbeddedStruct struct {
		prop1 string
		prop2 bool
	}
	fields := []reflect.StructField{
		{
			Name:      "TestEmbeddedStruct",
			Type:      reflect.TypeOf(TestEmbeddedStruct{}),
			Anonymous: true,
		},
	}
	newType := reflect.StructOf(fields)
	if newType.Field(0).Anonymous != true { // 'Anonymous' was ignored because the name wasn't empty.
		// Now make sure using an empty name actually works.
		defer func() {
			if r := recover(); r != nil {
				panic(fmt.Errorf("Couldn't create a struct with anonymous fields: %v", r))
			}
		}()
		fields[0].Name = ""
		reflect.StructOf(fields)
		return true
	} else {
		return false
	}
}

var emptyNameRequiredForAnonymousField = isEmptyNameRequiredForAnonymousField()

// coalesceTypes will squash multiple types to new type. This has different result
// than Go composition of structs.
//
// Example (go composition):
// type compositeStruct struct {
//     testPropsGroupA
//     testPropsGroupB
// }
// Debug print:
// {
//     testPropsGroupA: core.testPropsGroupA
//     {
//       Field_a: string
//       Field_c: string
//       Field_f: string
//     }
//     testPropsGroupB: core.testPropsGroupB
//     {
//       Field_b: string
//     }
// }
//
// Example for: coalesceTypes([]reflect.Type{
//    testPropsGroupA{},
//    testPropsGroupB{},
// })
// Debug print:
// {
//     Field_a: string
//     Field_c: string
//     Field_f: string
//     Field_b: string
// }
func coalesceTypes(list ...reflect.Type) reflect.Type {
	if len(list) == 0 {
		panic("List can't be empty")
	}
	if len(list) == 1 {
		return list[0]
	}

	fieldsKeys := map[string]bool{}
	fields := []reflect.StructField{}

	for _, elementType := range list {
		for i := 0; i < elementType.NumField(); i++ {
			field := elementType.Field(i)
			fieldName := field.Name
			if emptyNameRequiredForAnonymousField && field.Anonymous {
				field.Name = ""
			}
			if _, ok := fieldsKeys[fieldName]; ok {
				panic(fmt.Sprintf("Name collision: '%v'\n", fieldName))
			} else {
				fieldsKeys[fieldName] = true
				fields = append(fields, field)
			}
		}
	}

	return reflect.StructOf(fields)
}

// AppendProps merges properties from BlueprintEmbed to dst, but only for enabled features
// expect that Features are inited (before using this function we should call Features.Init)
// expect that properties.Features should contain all available features (whenever disabled/enabled)
func (f *Features) AppendProps(dst []interface{}, properties *configProperties) error {
	// featuresData is struct created in Features.Init function
	featuresData := reflect.ValueOf(f.BlueprintEmbed).Elem()

	for _, featureKey := range utils.SortedKeysBoolMap(properties.Features) {
		if properties.Features[featureKey] { // Check the feature is enabled
			// Features are matched like "Feature_name" - feature structure
			featureFieldName := featurePropertyName(featureKey)
			featureStruct := featuresData.FieldByName(featureFieldName)
			if !featureStruct.IsValid() {
				panic(fmt.Sprintf("Field returned for property %s isn't valid\n", featureFieldName))
			}
			// AppendProperties expects a pointer to a struct.
			featureStructPointer := featureStruct.FieldByName("BlueprintEmbed").Interface()

			// If featureProps is nil then we've determined that we can skip this,
			// so avoid calling AppendProperties
			if featureStructPointer != nil {
				err := proptools.AppendMatchingProperties(dst, featureStructPointer, nil)
				if err != nil {
					return err
				}
			}
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package gofast

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"net/http/httptest"
	"strings"
)

// NewAuthRequest returns a new *http.Request
// and a *Request with the body buffered
// into new NopReader.
func NewAuthRequest(orgl *http.Request) (r *http.Request, req *Request, err error) {

	// new request struct that inherits orgl values
	r = &http.Request{}
	*r = *orgl

	var stdin io.ReadCloser

	// clone the raw request content into r.Body and stdin
	// if there is any body
	if orgl.Body != nil {
		var content []byte
		content, err = ioutil.ReadAll(orgl.Body)
		if err != nil {
			return
		}
		r.Body = ioutil.NopCloser(bytes.NewBuffer(content))
		stdin = ioutil.NopCloser(bytes.NewBuffer(content))
	}

	// generate the request
	req = &Request{
		Raw:    orgl,
		Role:   RoleAuthorizer,
		Params: make(map[string]string),
		Stdin:  stdin,
		Data:   nil,
	}
	return
}

// NewAuthorizer creates an authorizer
func NewAuthorizer(clientFactory ClientFactory, sessionHandler SessionHandler) *Authorizer {
	return &Authorizer{
		clientFactory,
		sessionHandler,
	}
}

// Authorizer guard a given http.Handler
//
// Since this is implemented as a generic http.Handler middleware,
// you may use this with other non-gofast library components as
// long as they implements the http.Handler interface.
type Authorizer struct {
	clientFactory  ClientFactory
	sessionHandler SessionHandler
}

// Wrap method is a generic http.Handler middleware. Requests
// to wrapped hander would go through the fastcgi authorizer
// first. If not authorized, the request will not reach wrapped
// hander.
func (ar Authorizer) Wrap(inner http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {

		// generate auth request
		innerReq, req, err := NewAuthRequest(r)
		if err != nil {
			w.Header().Add("Content-Type", "text/html; charset=utf8")
			w.WriteHeader(http.StatusBadRequest)
			fmt.Fprintf(w, "%s", err)
			return
		}

		// get client to fastcgi application
		c, err := ar.clientFactory()
		if err != nil {
			w.Header().Add("Content-Type", "text/html; charset=utf8")
			w.WriteHeader(http.StatusInternalServerError)
			fmt.Fprintf(w, "unable to connect to authorizer: %s", err)
			return
		}

		// make request with client
		resp, err := ar.sessionHandler(c, req)
		if err != nil {
			w.Header().Add("Content-Type", "text/html; charset=utf8")
			w.WriteHeader(http.StatusInternalServerError)
			fmt.Fprintf(w, "error with authorizer request: %s", err)
			return
		}

		ew := new(bytes.Buffer)
		rw := httptest.NewRecorder() // FIXME: should do this without httptest
		if err = resp.WriteTo(rw, ew); err != nil {
			log.Printf("cannot write to response pipe: %s", err)
			w.WriteHeader(http.StatusInternalServerError)
			fmt.Fprintf(w, http.StatusText(http.StatusInternalServerError))
			return
		}

		// if code is not http.StatusOK (200)
		if rw.Code != http.StatusOK {
			// copy header map
			for k, m := range rw.Header() {
				for _, v := range m {
					w.Header().Add(k, v)
				}
			}
			w.WriteHeader(rw.Code)
			fmt.Fprintf(w, rw.Body.String())

			// if error stream is not empty
			// also write to response
			// TODO: add option to suppress this?
			if ew.Len() > 0 {
				w.Header().Add("Content-Type", "text/html; charset=utf8")
				w.WriteHeader(http.StatusInternalServerError)
				fmt.Fprintf(w, "error reading authorizer response: %s", err)
				log.Printf("gofast: error stream from application process %s",
					ew.String())
				return
			}
			return
		}

		// no problem from authorizer
		// pass down variable to the inner handler
		// and discard the authorizer stdout and stderr
		for k, m := range rw.HeaderMap {
			// looking for header with keys "Variable-*"
			// strip the prefix and pass to the inner header
			if len(k) > 9 && strings.HasPrefix(strings.ToLower(k), "variable-") {
				innerKey := k[9:]
				for _, v := range m {
					log.Printf("k: %s, innerKey: %s, v: %s", k, innerKey, v)
					innerReq.Header.Add(innerKey, v)
				}
			}
		}
		inner.ServeHTTP(w, innerReq)
	})
}
<commit_msg>authorizer: Fix unnecessary Fprintf<commit_after>package gofast

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"net/http/httptest"
	"strings"
)

// NewAuthRequest returns a new *http.Request
// and a *Request with the body buffered
// into new NopReader.
func NewAuthRequest(orgl *http.Request) (r *http.Request, req *Request, err error) {

	// new request struct that inherits orgl values
	r = &http.Request{}
	*r = *orgl

	var stdin io.ReadCloser

	// clone the raw request content into r.Body and stdin
	// if there is any body
	if orgl.Body != nil {
		var content []byte
		content, err = ioutil.ReadAll(orgl.Body)
		if err != nil {
			return
		}
		r.Body = ioutil.NopCloser(bytes.NewBuffer(content))
		stdin = ioutil.NopCloser(bytes.NewBuffer(content))
	}

	// generate the request
	req = &Request{
		Raw:    orgl,
		Role:   RoleAuthorizer,
		Params: make(map[string]string),
		Stdin:  stdin,
		Data:   nil,
	}
	return
}

// NewAuthorizer creates an authorizer
func NewAuthorizer(clientFactory ClientFactory, sessionHandler SessionHandler) *Authorizer {
	return &Authorizer{
		clientFactory,
		sessionHandler,
	}
}

// Authorizer guard a given http.Handler
//
// Since this is implemented as a generic http.Handler middleware,
// you may use this with other non-gofast library components as
// long as they implements the http.Handler interface.
type Authorizer struct {
	clientFactory  ClientFactory
	sessionHandler SessionHandler
}

// Wrap method is a generic http.Handler middleware. Requests
// to wrapped hander would go through the fastcgi authorizer
// first. If not authorized, the request will not reach wrapped
// hander.
func (ar Authorizer) Wrap(inner http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {

		// generate auth request
		innerReq, req, err := NewAuthRequest(r)
		if err != nil {
			w.Header().Add("Content-Type", "text/html; charset=utf8")
			w.WriteHeader(http.StatusBadRequest)
			fmt.Fprintf(w, "%s", err)
			return
		}

		// get client to fastcgi application
		c, err := ar.clientFactory()
		if err != nil {
			w.Header().Add("Content-Type", "text/html; charset=utf8")
			w.WriteHeader(http.StatusInternalServerError)
			fmt.Fprintf(w, "unable to connect to authorizer: %s", err)
			return
		}

		// make request with client
		resp, err := ar.sessionHandler(c, req)
		if err != nil {
			w.Header().Add("Content-Type", "text/html; charset=utf8")
			w.WriteHeader(http.StatusInternalServerError)
			fmt.Fprintf(w, "error with authorizer request: %s", err)
			return
		}

		ew := new(bytes.Buffer)
		rw := httptest.NewRecorder() // FIXME: should do this without httptest
		if err = resp.WriteTo(rw, ew); err != nil {
			log.Printf("cannot write to response pipe: %s", err)
			w.WriteHeader(http.StatusInternalServerError)
			fmt.Fprint(w, http.StatusText(http.StatusInternalServerError))
			return
		}

		// if code is not http.StatusOK (200)
		if rw.Code != http.StatusOK {
			// copy header map
			for k, m := range rw.Header() {
				for _, v := range m {
					w.Header().Add(k, v)
				}
			}
			w.WriteHeader(rw.Code)
			fmt.Fprint(w, rw.Body.String())

			// if error stream is not empty
			// also write to response
			// TODO: add option to suppress this?
			if ew.Len() > 0 {
				w.Header().Add("Content-Type", "text/html; charset=utf8")
				w.WriteHeader(http.StatusInternalServerError)
				fmt.Fprintf(w, "error reading authorizer response: %s", err)
				log.Printf("gofast: error stream from application process %s",
					ew.String())
				return
			}
			return
		}

		// no problem from authorizer
		// pass down variable to the inner handler
		// and discard the authorizer stdout and stderr
		for k, m := range rw.HeaderMap {
			// looking for header with keys "Variable-*"
			// strip the prefix and pass to the inner header
			if len(k) > 9 && strings.HasPrefix(strings.ToLower(k), "variable-") {
				innerKey := k[9:]
				for _, v := range m {
					log.Printf("k: %s, innerKey: %s, v: %s", k, innerKey, v)
					innerReq.Header.Add(innerKey, v)
				}
			}
		}
		inner.ServeHTTP(w, innerReq)
	})
}
<|endoftext|>


<commit_before>package device

import (
	"fmt"
	"os"
	"path/filepath"
	"strconv"
	"strings"

	"golang.org/x/sys/unix"

	"github.com/lxc/lxd/lxd/device/config"
	"github.com/lxc/lxd/lxd/state"
	"github.com/lxc/lxd/shared"
	"github.com/lxc/lxd/shared/idmap"
	"github.com/lxc/lxd/shared/logger"
)

// instanceUnixGetDeviceAttributes returns the UNIX device attributes for an instance device.
// Uses supplied device config for device properties, and if they haven't been set, falls back to
// using UnixGetDeviceAttributes() to directly query an existing device file.
func instanceUnixGetDeviceAttributes(devicesPath string, prefix string, config config.Device) (string, int, int, error) {
	// Check if we've been passed major and minor numbers already.
	var tmp int
	var err error
	dMajor := -1
	if config["major"] != "" {
		tmp, err = strconv.Atoi(config["major"])
		if err == nil {
			dMajor = tmp
		}
	}

	dMinor := -1
	if config["minor"] != "" {
		tmp, err = strconv.Atoi(config["minor"])
		if err == nil {
			dMinor = tmp
		}
	}

	dType := ""
	if config["type"] == "unix-char" {
		dType = "c"
	} else if config["type"] == "unix-block" {
		dType = "b"
	}

	// Figure out the paths
	destPath := config["path"]
	if destPath == "" {
		destPath = config["source"]
	}
	relativeDestPath := strings.TrimPrefix(destPath, "/")
	devName := fmt.Sprintf("%s.%s", strings.Replace(prefix, "/", "-", -1), strings.Replace(relativeDestPath, "/", "-", -1))
	devPath := filepath.Join(devicesPath, devName)

	if dType == "" || dMajor < 0 || dMinor < 0 {
		dType, dMajor, dMinor, err = UnixGetDeviceAttributes(devPath)
		if err != nil {
			return dType, dMajor, dMinor, err
		}
	}

	return dType, dMajor, dMinor, err
}

// UnixGetDeviceAttributes returns the decice type, major and minor numbers for a device.
func UnixGetDeviceAttributes(path string) (string, int, int, error) {
	// Get a stat struct from the provided path
	stat := unix.Stat_t{}
	err := unix.Stat(path, &stat)
	if err != nil {
		return "", 0, 0, err
	}

	// Check what kind of file it is
	dType := ""
	if stat.Mode&unix.S_IFMT == unix.S_IFBLK {
		dType = "b"
	} else if stat.Mode&unix.S_IFMT == unix.S_IFCHR {
		dType = "c"
	} else {
		return "", 0, 0, fmt.Errorf("Not a device")
	}

	// Return the device information
	major := shared.Major(stat.Rdev)
	minor := shared.Minor(stat.Rdev)
	return dType, major, minor, nil
}

func unixGetDeviceModeOct(strmode string) (int, error) {
	// Default mode
	if strmode == "" {
		return 0600, nil
	}

	// Converted mode
	i, err := strconv.ParseInt(strmode, 8, 32)
	if err != nil {
		return 0, fmt.Errorf("Bad device mode: %s", strmode)
	}

	return int(i), nil
}

// UnixCreateDevice Unix devices handling.
func UnixCreateDevice(s *state.State, idmapSet *idmap.IdmapSet, devicesPath string, prefix string, m config.Device, defaultMode bool) ([]string, error) {
	var err error
	var major, minor int

	// Extra checks for nesting
	if s.OS.RunningInUserNS {
		for key, value := range m {
			if shared.StringInSlice(key, []string{"major", "minor", "mode", "uid", "gid"}) && value != "" {
				return nil, fmt.Errorf("The \"%s\" property may not be set when adding a device to a nested container", key)
			}
		}
	}

	srcPath := m["source"]
	if srcPath == "" {
		srcPath = m["path"]
	}
	srcPath = shared.HostPath(srcPath)

	// Get the major/minor of the device we want to create
	if m["major"] == "" && m["minor"] == "" {
		// If no major and minor are set, use those from the device on the host
		_, major, minor, err = UnixGetDeviceAttributes(srcPath)
		if err != nil {
			return nil, fmt.Errorf("Failed to get device attributes for %s: %s", m["path"], err)
		}
	} else if m["major"] == "" || m["minor"] == "" {
		return nil, fmt.Errorf("Both major and minor must be supplied for device: %s", m["path"])
	} else {
		major, err = strconv.Atoi(m["major"])
		if err != nil {
			return nil, fmt.Errorf("Bad major %s in device %s", m["major"], m["path"])
		}

		minor, err = strconv.Atoi(m["minor"])
		if err != nil {
			return nil, fmt.Errorf("Bad minor %s in device %s", m["minor"], m["path"])
		}
	}

	// Get the device mode
	mode := os.FileMode(0660)
	if m["mode"] != "" {
		tmp, err := unixGetDeviceModeOct(m["mode"])
		if err != nil {
			return nil, fmt.Errorf("Bad mode %s in device %s", m["mode"], m["path"])
		}
		mode = os.FileMode(tmp)
	} else if !defaultMode {
		mode, err = shared.GetPathMode(srcPath)
		if err != nil {
			errno, isErrno := shared.GetErrno(err)
			if !isErrno || errno != unix.ENOENT {
				return nil, fmt.Errorf("Failed to retrieve mode of device %s: %s", m["path"], err)
			}
			mode = os.FileMode(0660)
		}
	}

	if m["type"] == "unix-block" {
		mode |= unix.S_IFBLK
	} else {
		mode |= unix.S_IFCHR
	}

	// Get the device owner
	uid := 0
	gid := 0

	if m["uid"] != "" {
		uid, err = strconv.Atoi(m["uid"])
		if err != nil {
			return nil, fmt.Errorf("Invalid uid %s in device %s", m["uid"], m["path"])
		}
	}

	if m["gid"] != "" {
		gid, err = strconv.Atoi(m["gid"])
		if err != nil {
			return nil, fmt.Errorf("Invalid gid %s in device %s", m["gid"], m["path"])
		}
	}

	// Create the devices directory if missing
	if !shared.PathExists(devicesPath) {
		os.Mkdir(devicesPath, 0711)
		if err != nil {
			return nil, fmt.Errorf("Failed to create devices path: %s", err)
		}
	}

	destPath := m["path"]
	if destPath == "" {
		destPath = m["source"]
	}
	relativeDestPath := strings.TrimPrefix(destPath, "/")
	devName := fmt.Sprintf("%s.%s", strings.Replace(prefix, "/", "-", -1), strings.Replace(relativeDestPath, "/", "-", -1))
	devPath := filepath.Join(devicesPath, devName)

	// Create the new entry
	if !s.OS.RunningInUserNS {
		encodedDeviceNumber := (minor & 0xff) | (major << 8) | ((minor & ^0xff) << 12)
		err := unix.Mknod(devPath, uint32(mode), encodedDeviceNumber)
		if err != nil {
			return nil, fmt.Errorf("Failed to create device %s for %s: %s", devPath, m["path"], err)
		}

		err = os.Chown(devPath, uid, gid)
		if err != nil {
			return nil, fmt.Errorf("Failed to chown device %s: %s", devPath, err)
		}

		// Needed as mknod respects the umask
		err = os.Chmod(devPath, mode)
		if err != nil {
			return nil, fmt.Errorf("Failed to chmod device %s: %s", devPath, err)
		}

		if idmapSet != nil {
			err := idmapSet.ShiftFile(devPath)
			if err != nil {
				// uidshift failing is weird, but not a big problem.  Log and proceed
				logger.Debugf("Failed to uidshift device %s: %s\n", m["path"], err)
			}
		}
	} else {
		f, err := os.Create(devPath)
		if err != nil {
			return nil, err
		}
		f.Close()

		err = DiskMount(srcPath, devPath, false, false, "")
		if err != nil {
			return nil, err
		}
	}

	return []string{devPath, relativeDestPath}, nil
}
<commit_msg>device/device/utils/unix: Device management function rework<commit_after>package device

import (
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"strconv"
	"strings"

	"golang.org/x/sys/unix"

	"github.com/lxc/lxd/lxd/device/config"
	"github.com/lxc/lxd/lxd/state"
	"github.com/lxc/lxd/shared"
	"github.com/lxc/lxd/shared/idmap"
	"github.com/lxc/lxd/shared/logger"
)

// unixDeviceInstanceAttributes returns the UNIX device attributes for an instance device.
// Uses supplied device config for device properties, and if they haven't been set, falls back to
// using UnixGetDeviceAttributes() to directly query an existing device file.
func unixDeviceInstanceAttributes(devicesPath string, prefix string, config config.Device) (string, int, int, error) {
	// Check if we've been passed major and minor numbers already.
	var tmp int
	var err error
	dMajor := -1
	if config["major"] != "" {
		tmp, err = strconv.Atoi(config["major"])
		if err == nil {
			dMajor = tmp
		}
	}

	dMinor := -1
	if config["minor"] != "" {
		tmp, err = strconv.Atoi(config["minor"])
		if err == nil {
			dMinor = tmp
		}
	}

	dType := ""
	if config["type"] == "unix-char" {
		dType = "c"
	} else if config["type"] == "unix-block" {
		dType = "b"
	}

	// Figure out the paths
	destPath := config["path"]
	if destPath == "" {
		destPath = config["source"]
	}
	relativeDestPath := strings.TrimPrefix(destPath, "/")
	devName := fmt.Sprintf("%s.%s", strings.Replace(prefix, "/", "-", -1), strings.Replace(relativeDestPath, "/", "-", -1))
	devPath := filepath.Join(devicesPath, devName)

	if dType == "" || dMajor < 0 || dMinor < 0 {
		dType, dMajor, dMinor, err = UnixDeviceAttributes(devPath)
		if err != nil {
			return dType, dMajor, dMinor, err
		}
	}

	return dType, dMajor, dMinor, err
}

// UnixDeviceAttributes returns the decice type, major and minor numbers for a device.
func UnixDeviceAttributes(path string) (string, int, int, error) {
	// Get a stat struct from the provided path
	stat := unix.Stat_t{}
	err := unix.Stat(path, &stat)
	if err != nil {
		return "", 0, 0, err
	}

	// Check what kind of file it is
	dType := ""
	if stat.Mode&unix.S_IFMT == unix.S_IFBLK {
		dType = "b"
	} else if stat.Mode&unix.S_IFMT == unix.S_IFCHR {
		dType = "c"
	} else {
		return "", 0, 0, fmt.Errorf("Not a device")
	}

	// Return the device information
	major := shared.Major(stat.Rdev)
	minor := shared.Minor(stat.Rdev)
	return dType, major, minor, nil
}

func unixDeviceModeOct(strmode string) (int, error) {
	// Default mode
	if strmode == "" {
		return 0600, nil
	}

	// Converted mode
	i, err := strconv.ParseInt(strmode, 8, 32)
	if err != nil {
		return 0, fmt.Errorf("Bad device mode: %s", strmode)
	}

	return int(i), nil
}

// UnixDevice contains information about a created UNIX device.
type UnixDevice struct {
	HostPath     string      // Absolute path to the device on the host.
	RelativePath string      // Relative path where the device will be mounted inside instance.
	Type         string      // Type of device; c (for char) or b for (block).
	Major        int         // Major number.
	Minor        int         // Minor number.
	Mode         os.FileMode // File mode.
	UID          int         // Owner UID.
	GID          int         // Owner GID.
}

// unixDeviceDestPath returns the absolute path for a device inside an instance.
// This is based on the "path" property of the devices' config, or the "source" property if "path"
// not defined.
func unixDeviceDestPath(m config.Device) string {
	destPath := m["path"]
	if destPath == "" {
		destPath = m["source"]
	}

	return destPath
}

// UnixDeviceCreate creates a UNIX device (either block or char). If the supplied device config map
// contains a major and minor number for the device, then a stat is avoided, otherwise this info
// retrieved from the origin device. Similarly, if a mode is supplied in the device config map or
// defaultMode is set as true, then the device is created with the supplied or default mode (0660)
// respectively, otherwise the origin device's mode is used. If the device config doesn't contain a
// type field then it defaults to created a unix-char device. The ownership of the created device
// defaults to root (0) but can be specified with the uid and gid fields in the device config map.
// It returns a UnixDevice containing information about the device created.
func UnixDeviceCreate(s *state.State, idmapSet *idmap.IdmapSet, devicesPath string, prefix string, m config.Device, defaultMode bool) (*UnixDevice, error) {
	var err error
	d := UnixDevice{}

	// Extra checks for nesting.
	if s.OS.RunningInUserNS {
		for key, value := range m {
			if shared.StringInSlice(key, []string{"major", "minor", "mode", "uid", "gid"}) && value != "" {
				return nil, fmt.Errorf("The \"%s\" property may not be set when adding a device to a nested container", key)
			}
		}
	}

	srcPath := m["source"]
	if srcPath == "" {
		srcPath = m["path"]
	}
	srcPath = shared.HostPath(srcPath)

	// Get the major/minor of the device we want to create.
	if m["major"] == "" && m["minor"] == "" {
		// If no major and minor are set, use those from the device on the host.
		_, d.Major, d.Minor, err = UnixDeviceAttributes(srcPath)
		if err != nil {
			return nil, fmt.Errorf("Failed to get device attributes for %s: %s", m["path"], err)
		}
	} else if m["major"] == "" || m["minor"] == "" {
		return nil, fmt.Errorf("Both major and minor must be supplied for device: %s", m["path"])
	} else {
		d.Major, err = strconv.Atoi(m["major"])
		if err != nil {
			return nil, fmt.Errorf("Bad major %s in device %s", m["major"], m["path"])
		}

		d.Minor, err = strconv.Atoi(m["minor"])
		if err != nil {
			return nil, fmt.Errorf("Bad minor %s in device %s", m["minor"], m["path"])
		}
	}

	// Get the device mode (defaults to 0660 if not supplied).
	d.Mode = os.FileMode(0660)
	if m["mode"] != "" {
		tmp, err := unixDeviceModeOct(m["mode"])
		if err != nil {
			return nil, fmt.Errorf("Bad mode %s in device %s", m["mode"], m["path"])
		}
		d.Mode = os.FileMode(tmp)
	} else if !defaultMode {
		d.Mode, err = shared.GetPathMode(srcPath)
		if err != nil {
			errno, isErrno := shared.GetErrno(err)
			if !isErrno || errno != unix.ENOENT {
				return nil, fmt.Errorf("Failed to retrieve mode of device %s: %s", m["path"], err)
			}
			d.Mode = os.FileMode(0660)
		}
	}

	if m["type"] == "unix-block" {
		d.Mode |= unix.S_IFBLK
		d.Type = "b"
	} else {
		d.Mode |= unix.S_IFCHR
		d.Type = "c"
	}

	// Get the device owner.
	if m["uid"] != "" {
		d.UID, err = strconv.Atoi(m["uid"])
		if err != nil {
			return nil, fmt.Errorf("Invalid uid %s in device %s", m["uid"], m["path"])
		}
	}

	if m["gid"] != "" {
		d.GID, err = strconv.Atoi(m["gid"])
		if err != nil {
			return nil, fmt.Errorf("Invalid gid %s in device %s", m["gid"], m["path"])
		}
	}

	// Create the devices directory if missing.
	if !shared.PathExists(devicesPath) {
		os.Mkdir(devicesPath, 0711)
		if err != nil {
			return nil, fmt.Errorf("Failed to create devices path: %s", err)
		}
	}

	destPath := unixDeviceDestPath(m)
	relativeDestPath := strings.TrimPrefix(destPath, "/")
	devName := unixDeviceEncode(unixDeviceJoinPath(prefix, relativeDestPath))
	devPath := filepath.Join(devicesPath, devName)

	// Create the new entry.
	if !s.OS.RunningInUserNS {
		devNum := int(unix.Mkdev(uint32(d.Major), uint32(d.Minor)))
		err := unix.Mknod(devPath, uint32(d.Mode), devNum)
		if err != nil {
			return nil, fmt.Errorf("Failed to create device %s for %s: %s", devPath, m["path"], err)
		}

		err = os.Chown(devPath, d.UID, d.GID)
		if err != nil {
			return nil, fmt.Errorf("Failed to chown device %s: %s", devPath, err)
		}

		// Needed as mknod respects the umask.
		err = os.Chmod(devPath, d.Mode)
		if err != nil {
			return nil, fmt.Errorf("Failed to chmod device %s: %s", devPath, err)
		}

		if idmapSet != nil {
			err := idmapSet.ShiftFile(devPath)
			if err != nil {
				// uidshift failing is weird, but not a big problem. Log and proceed.
				logger.Debugf("Failed to uidshift device %s: %s\n", m["path"], err)
			}
		}
	} else {
		f, err := os.Create(devPath)
		if err != nil {
			return nil, err
		}
		f.Close()

		err = DiskMount(srcPath, devPath, false, false, "")
		if err != nil {
			return nil, err
		}
	}

	d.HostPath = devPath
	d.RelativePath = relativeDestPath
	return &d, nil
}

// unixDeviceSetup creates a UNIX device on host and then configures supplied RunConfig with the
// mount and cgroup rule instructions to have it be attached to the instance. If defaultMode is true
// or mode is supplied in the device config then the origin device does not need to be accessed for
// its file mode.
func unixDeviceSetup(s *state.State, devicesPath string, typePrefix string, deviceName string, m config.Device, defaultMode bool, runConf *RunConfig) error {
	// Before creating the device, check that another existing device isn't using the same mount
	// path inside the instance as our device. If we find an existing device with the same mount
	// path we will skip mounting our device inside the instance. This can happen when multiple
	// LXD devices share the same parent device (such as Nvidia GPUs and Infiniband devices).

	// Convert the requested dest path inside the instance to an encoded relative one.
	ourDestPath := unixDeviceDestPath(m)
	ourEncRelDestFile := unixDeviceEncode(strings.TrimPrefix(ourDestPath, "/"))

	// Load all existing host devices.
	dents, err := ioutil.ReadDir(devicesPath)
	if err != nil {
		if !os.IsNotExist(err) {
			return err
		}
	}

	dupe := false
	for _, ent := range dents {
		devName := ent.Name()

		// Remove the LXD device type and name prefix, leaving just the encoded dest path.
		idx := strings.LastIndex(devName, ".")
		if idx == -1 {
			return fmt.Errorf("Invalid device name \"%s\"", devName)
		}

		encRelDestFile := devName[idx+1:]

		// If the encoded relative path of the device file matches the encoded relative dest
		// path of our new device then return as we do not want to instruct LXD to mount
		// the device and create cgroup rules.
		if encRelDestFile == ourEncRelDestFile {
			dupe = true // There is an existing device using the same mount path.
			break
		}
	}

	// Create the device on the host.
	ourPrefix := unixDeviceEncode(unixDeviceJoinPath(typePrefix, deviceName))
	d, err := UnixDeviceCreate(s, nil, devicesPath, ourPrefix, m, defaultMode)
	if err != nil {
		return err
	}

	// If there was an existing device using the same mount path detected then skip mounting.
	if dupe {
		return nil
	}

	// Instruct liblxc to perform the mount.
	runConf.Mounts = append(runConf.Mounts, MountEntryItem{
		DevPath:    d.HostPath,
		TargetPath: d.RelativePath,
		FSType:     "none",
		Opts:       []string{"bind", "create=file"},
	})

	// Instruct liblxc to setup the cgroup rule.
	runConf.CGroups = append(runConf.CGroups, RunConfigItem{
		Key:   "devices.allow",
		Value: fmt.Sprintf("%s %d:%d rwm", d.Type, d.Major, d.Minor),
	})

	return nil
}

// unixDeviceSetupCharNum calls unixDeviceSetup and overrides the supplied device config with the
// type as "unix-char" and the supplied major and minor numbers. This function can be used when you
// already know the device's major and minor numbers to avoid unixDeviceSetup() having to stat the
// device to ascertain these attributes. If defaultMode is true or mode is supplied in the device
// config then the origin device does not need to be accessed for its file mode.
func unixDeviceSetupCharNum(s *state.State, devicesPath string, typePrefix string, deviceName string, m config.Device, major int, minor int, path string, defaultMode bool, runConf *RunConfig) error {
	configCopy := config.Device{}
	for k, v := range m {
		configCopy[k] = v
	}

	// Overridng these in the config copy should avoid the need for unixDeviceSetup to stat
	// the origin device to ascertain this information.
	configCopy["type"] = "unix-char"
	configCopy["major"] = fmt.Sprintf("%d", major)
	configCopy["minor"] = fmt.Sprintf("%d", minor)
	configCopy["path"] = path

	return unixDeviceSetup(s, devicesPath, typePrefix, deviceName, configCopy, defaultMode, runConf)
}

// UnixDeviceExists checks if the unix device already exists in devices path.
func UnixDeviceExists(devicesPath string, prefix string, path string) bool {
	relativeDestPath := strings.TrimPrefix(path, "/")
	devName := fmt.Sprintf("%s.%s", unixDeviceEncode(prefix), unixDeviceEncode(relativeDestPath))
	devPath := filepath.Join(devicesPath, devName)

	return shared.PathExists(devPath)
}

// unixDeviceEncode encodes a string to be used as part of a file name in the LXD devices path.
// The encoding scheme replaces "-" with "--" and then "/" with "-".
func unixDeviceEncode(text string) string {
	return strings.Replace(strings.Replace(text, "-", "--", -1), "/", "-", -1)
}

// unixDeviceDecode decodes a string used in the LXD devices path back to its original form.
// The decoding scheme converts "-" back to "/" and "--" back to "-".
func unixDeviceDecode(text string) string {
	// This converts "--" to the null character "\0" first, to allow remaining "-" chars to be
	// converted back to "/" before making a final pass to convert "\0" back to original "-".
	return strings.Replace(strings.Replace(strings.Replace(text, "--", "\000", -1), "-", "/", -1), "\000", "-", -1)
}

// unixDeviceJoinPath joins together prefix and text delimited by a "." for device path generation.
func unixDeviceJoinPath(prefix string, text string) string {
	return fmt.Sprintf("%s.%s", prefix, text)
}

// unixRemoveDevice identifies all files related to the supplied typePrefix and deviceName and then
// populates the supplied runConf with the instructions to remove cgroup rules and unmount devices.
// It detects if any other devices attached to the instance that share the same prefix have the same
// relative mount path inside the instance encoded into the file name. If there is another device
// that shares the same mount path then the unmount rule is not added to the runConf as the device
// may still be in use with another LXD device.
func unixDeviceRemove(devicesPath string, typePrefix string, deviceName string, runConf *RunConfig) error {
	// Load all devices.
	dents, err := ioutil.ReadDir(devicesPath)
	if err != nil {
		if !os.IsNotExist(err) {
			return err
		}
	}

	ourPrefix := unixDeviceEncode(unixDeviceJoinPath(typePrefix, deviceName))
	ourDevs := []string{}
	otherDevs := []string{}

	for _, ent := range dents {
		devName := ent.Name()

		// This device file belongs our LXD device.
		if strings.HasPrefix(devName, ourPrefix) {
			ourDevs = append(ourDevs, devName)
			continue
		}

		// This device file belongs to another LXD device.
		otherDevs = append(otherDevs, devName)
	}

	// It is possible for some LXD devices to share the same device on the same mount point
	// inside the instance. We extract the relative path of the device that is encoded into its
	// name on the host so that we can compare the device files for our own device and check
	// none of them use the same mount point.
	encRelDevFiles := []string{}
	for _, otherDev := range otherDevs {
		// Remove the LXD device type and name prefix, leaving just the encoded dest path.
		idx := strings.LastIndex(otherDev, ".")
		if idx == -1 {
			return fmt.Errorf("Invalid device name \"%s\"", otherDev)
		}

		encRelDestFile := otherDev[idx+1:]
		encRelDevFiles = append(encRelDevFiles, encRelDestFile)
	}

	// Check that none of our devices are in use by another LXD device.
	for _, ourDev := range ourDevs {
		// Remove the LXD device type and name prefix, leaving just the encoded dest path.
		idx := strings.LastIndex(ourDev, ".")
		if idx == -1 {
			return fmt.Errorf("Invalid device name \"%s\"", ourDev)
		}

		ourEncRelDestFile := ourDev[idx+1:]

		// Look for devices for other LXD devices that match the same path.
		dupe := false
		for _, encRelDevFile := range encRelDevFiles {
			if encRelDevFile == ourEncRelDestFile {
				dupe = true
				break
			}
		}

		// If a device has been found that points to the same device inside the instance
		// then we cannot request it be umounted inside the instance as it's still in use.
		if dupe {
			continue
		}

		// Append this device to the mount rules (these will be unmounted).
		runConf.Mounts = append(runConf.Mounts, MountEntryItem{
			TargetPath: unixDeviceDecode(ourEncRelDestFile),
		})

		absDevPath := filepath.Join(devicesPath, ourDev)
		dType, dMajor, dMinor, err := UnixDeviceAttributes(absDevPath)
		if err != nil {
			return fmt.Errorf("Failed to get UNIX device attributes for '%s': %v", absDevPath, err)
		}

		// Append a deny cgroup fule for this device.
		runConf.CGroups = append(runConf.CGroups, RunConfigItem{
			Key:   "devices.deny",
			Value: fmt.Sprintf("%s %d:%d rwm", dType, dMajor, dMinor),
		})
	}

	return nil
}

// unixDeviceDeleteFiles removes all host side device files for a particular LXD device.
// This should be run after the files have been detached from the instance using unixDeviceRemove().
func unixDeviceDeleteFiles(s *state.State, devicesPath string, typePrefix string, deviceName string) error {
	ourPrefix := unixDeviceEncode(unixDeviceJoinPath(typePrefix, deviceName))

	// Load all devices.
	dents, err := ioutil.ReadDir(devicesPath)
	if err != nil {
		if !os.IsNotExist(err) {
			return err
		}
	}

	// Remove our host side device files.
	for _, ent := range dents {
		devName := ent.Name()

		// This device file belongs our LXD device.
		if strings.HasPrefix(devName, ourPrefix) {
			devPath := filepath.Join(devicesPath, devName)

			// Remove the host side mount.
			if s.OS.RunningInUserNS {
				unix.Unmount(devPath, unix.MNT_DETACH)
			}

			// Remove the host side device file.
			err = os.Remove(devPath)
			if err != nil {
				return err
			}
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package mpb

import (
	"container/heap"
	"io"
	"io/ioutil"
	"os"
	"sync"
	"time"

	"github.com/vbauerster/mpb/cwriter"
)

const (
	// default RefreshRate
	prr = 120 * time.Millisecond
	// default width
	pwidth = 80
	// default format
	pformat = "[=>-]"
)

// Progress represents the container that renders Progress bars
type Progress struct {
	wg           *sync.WaitGroup
	uwg          *sync.WaitGroup
	operateState chan func(*pState)
	done         chan struct{}
}

type (
	// progress state, which may contain several bars
	pState struct {
		bHeap           *priorityQueue
		shutdownPending []*Bar
		heapUpdated     bool
		zeroWait        bool
		idCounter       int
		width           int
		format          string
		rr              time.Duration
		cw              *cwriter.Writer
		ticker          *time.Ticker

		// following are provided by user
		uwg              *sync.WaitGroup
		cancel           <-chan struct{}
		shutdownNotifier chan struct{}
		interceptors     []func(io.Writer)
		waitBars         map[*Bar]*Bar
		debugOut         io.Writer
	}
	widthSyncer struct {
		// Public for easy testing
		Accumulator []chan int
		Distributor []chan int
	}
)

// New creates new Progress instance, which orchestrates bars rendering process.
// Accepts mpb.ProgressOption funcs for customization.
func New(options ...ProgressOption) *Progress {
	pq := make(priorityQueue, 0)
	heap.Init(&pq)
	s := &pState{
		bHeap:    &pq,
		width:    pwidth,
		format:   pformat,
		cw:       cwriter.New(os.Stdout),
		rr:       prr,
		ticker:   time.NewTicker(prr),
		waitBars: make(map[*Bar]*Bar),
		debugOut: ioutil.Discard,
	}

	for _, opt := range options {
		if opt != nil {
			opt(s)
		}
	}

	p := &Progress{
		uwg:          s.uwg,
		wg:           new(sync.WaitGroup),
		operateState: make(chan func(*pState)),
		done:         make(chan struct{}),
	}
	go p.serve(s)
	return p
}

// AddBar creates a new progress bar and adds to the container.
func (p *Progress) AddBar(total int64, options ...BarOption) *Bar {
	p.wg.Add(1)
	result := make(chan *Bar, 1)
	select {
	case p.operateState <- func(s *pState) {
		options = append(options, barWidth(s.width), barFormat(s.format))
		b := newBar(p.wg, s.idCounter, total, s.cancel, options...)
		if b.runningBar != nil {
			s.waitBars[b.runningBar] = b
		} else {
			heap.Push(s.bHeap, b)
			s.heapUpdated = true
		}
		s.idCounter++
		result <- b
	}:
		return <-result
	case <-p.done:
		p.wg.Done()
		return nil
	}
}

// Abort is only effective while bar progress is running,
// it means remove bar now without waiting for its completion.
// If bar is already completed, there is nothing to abort.
// If you need to remove bar after completion, use BarRemoveOnComplete BarOption.
func (p *Progress) Abort(b *Bar) {
	select {
	case p.operateState <- func(s *pState) {
		if b.index < 0 {
			return
		}
		s.heapUpdated = heap.Remove(s.bHeap, b.index) != nil
		s.shutdownPending = append(s.shutdownPending, b)
	}:
	case <-p.done:
	}
}

// UpdateBarPriority provides a way to change bar's order position.
// Zero is highest priority, i.e. bar will be on top.
func (p *Progress) UpdateBarPriority(b *Bar, priority int) {
	select {
	case p.operateState <- func(s *pState) { s.bHeap.update(b, priority) }:
	case <-p.done:
	}
}

// BarCount returns bars count
func (p *Progress) BarCount() int {
	result := make(chan int, 1)
	select {
	case p.operateState <- func(s *pState) { result <- s.bHeap.Len() }:
		return <-result
	case <-p.done:
		return 0
	}
}

// Wait first waits for user provided *sync.WaitGroup, if any,
// then waits far all bars to complete and finally shutdowns master goroutine.
// After this method has been called, there is no way to reuse *Progress instance.
func (p *Progress) Wait() {
	if p.uwg != nil {
		p.uwg.Wait()
	}

	p.wg.Wait()

	select {
	case p.operateState <- func(s *pState) { s.zeroWait = true }:
		<-p.done
	case <-p.done:
	}
}

func newWidthSyncer(timeout <-chan struct{}, numBars, numColumn int) *widthSyncer {
	ws := &widthSyncer{
		Accumulator: make([]chan int, numColumn),
		Distributor: make([]chan int, numColumn),
	}
	for i := 0; i < numColumn; i++ {
		ws.Accumulator[i] = make(chan int, numBars)
		ws.Distributor[i] = make(chan int, numBars)
	}
	for i := 0; i < numColumn; i++ {
		go func(accumulator <-chan int, distributor chan<- int) {
			defer close(distributor)
			widths := make([]int, 0, numBars)
		loop:
			for {
				select {
				case w := <-accumulator:
					widths = append(widths, w)
					if len(widths) == numBars {
						break loop
					}
				case <-timeout:
					if len(widths) == 0 {
						return
					}
					break loop
				}
			}
			maxWidth := calcMax(widths)
			for i := 0; i < len(widths); i++ {
				distributor <- maxWidth
			}
		}(ws.Accumulator[i], ws.Distributor[i])
	}
	return ws
}

func (s *pState) writeAndFlush(tw, numP, numA int) (err error) {
	timeout := make(chan struct{})
	pSyncer := newWidthSyncer(timeout, s.bHeap.Len(), numP)
	aSyncer := newWidthSyncer(timeout, s.bHeap.Len(), numA)
	time.AfterFunc(s.rr-s.rr/12, func() {
		close(timeout)
	})

	for _, ch := range s.renderByPriority(tw, pSyncer, aSyncer) {
		bf := <-ch
		_, err = s.cw.ReadFrom(bf.reader)
		if !bf.bar.completed && bf.toComplete {
			bf.bar.completed = true
			if bf.bar.removeOnComplete {
				s.heapUpdated = heap.Remove(s.bHeap, bf.bar.index) != nil
			}
			if replacementBar, ok := s.waitBars[bf.bar]; ok {
				heap.Push(s.bHeap, replacementBar)
				s.heapUpdated = true
				delete(s.waitBars, bf.bar)
			}
			// defer is required to make removeOnComplete visually happen
			defer func() {
				s.shutdownPending = append(s.shutdownPending, bf.bar)
			}()
		}
	}

	for _, interceptor := range s.interceptors {
		interceptor(s.cw)
	}

	if e := s.cw.Flush(); err == nil {
		err = e
	}

	for i := len(s.shutdownPending) - 1; i >= 0; i-- {
		close(s.shutdownPending[i].shutdown)
		s.shutdownPending = s.shutdownPending[:i]
	}
	return
}

func (s *pState) renderByPriority(tw int, pSyncer, aSyncer *widthSyncer) []<-chan *bFrame {
	bff := make([]<-chan *bFrame, 0, s.bHeap.Len())
	for s.bHeap.Len() > 0 {
		b := heap.Pop(s.bHeap).(*Bar)
		defer heap.Push(s.bHeap, b)
		bff = append(bff, b.render(s.debugOut, tw, pSyncer, aSyncer))
	}
	return bff
}

func calcMax(slice []int) int {
	if len(slice) == 0 {
		return 0
	}

	max := slice[0]
	for i := 1; i < len(slice); i++ {
		if slice[i] > max {
			max = slice[i]
		}
	}
	return max
}
<commit_msg>some comments for future self<commit_after>package mpb

import (
	"container/heap"
	"io"
	"io/ioutil"
	"os"
	"sync"
	"time"

	"github.com/vbauerster/mpb/cwriter"
)

const (
	// default RefreshRate
	prr = 120 * time.Millisecond
	// default width
	pwidth = 80
	// default format
	pformat = "[=>-]"
)

// Progress represents the container that renders Progress bars
type Progress struct {
	wg           *sync.WaitGroup
	uwg          *sync.WaitGroup
	operateState chan func(*pState)
	done         chan struct{}
}

type (
	// progress state, which may contain several bars
	pState struct {
		bHeap           *priorityQueue
		shutdownPending []*Bar
		heapUpdated     bool
		zeroWait        bool
		idCounter       int
		width           int
		format          string
		rr              time.Duration
		cw              *cwriter.Writer
		ticker          *time.Ticker

		// following are provided by user
		uwg              *sync.WaitGroup
		cancel           <-chan struct{}
		shutdownNotifier chan struct{}
		interceptors     []func(io.Writer)
		waitBars         map[*Bar]*Bar
		debugOut         io.Writer
	}
	widthSyncer struct {
		// Public for easy testing
		Accumulator []chan int
		Distributor []chan int
	}
)

// New creates new Progress instance, which orchestrates bars rendering process.
// Accepts mpb.ProgressOption funcs for customization.
func New(options ...ProgressOption) *Progress {
	pq := make(priorityQueue, 0)
	heap.Init(&pq)
	s := &pState{
		bHeap:    &pq,
		width:    pwidth,
		format:   pformat,
		cw:       cwriter.New(os.Stdout),
		rr:       prr,
		ticker:   time.NewTicker(prr),
		waitBars: make(map[*Bar]*Bar),
		debugOut: ioutil.Discard,
	}

	for _, opt := range options {
		if opt != nil {
			opt(s)
		}
	}

	p := &Progress{
		uwg:          s.uwg,
		wg:           new(sync.WaitGroup),
		operateState: make(chan func(*pState)),
		done:         make(chan struct{}),
	}
	go p.serve(s)
	return p
}

// AddBar creates a new progress bar and adds to the container.
func (p *Progress) AddBar(total int64, options ...BarOption) *Bar {
	p.wg.Add(1)
	result := make(chan *Bar, 1)
	select {
	case p.operateState <- func(s *pState) {
		options = append(options, barWidth(s.width), barFormat(s.format))
		b := newBar(p.wg, s.idCounter, total, s.cancel, options...)
		if b.runningBar != nil {
			s.waitBars[b.runningBar] = b
		} else {
			heap.Push(s.bHeap, b)
			s.heapUpdated = true
		}
		s.idCounter++
		result <- b
	}:
		return <-result
	case <-p.done:
		p.wg.Done()
		return nil
	}
}

// Abort is only effective while bar progress is running,
// it means remove bar now without waiting for its completion.
// If bar is already completed, there is nothing to abort.
// If you need to remove bar after completion, use BarRemoveOnComplete BarOption.
func (p *Progress) Abort(b *Bar) {
	select {
	case p.operateState <- func(s *pState) {
		if b.index < 0 {
			return
		}
		s.heapUpdated = heap.Remove(s.bHeap, b.index) != nil
		s.shutdownPending = append(s.shutdownPending, b)
	}:
	case <-p.done:
	}
}

// UpdateBarPriority provides a way to change bar's order position.
// Zero is highest priority, i.e. bar will be on top.
func (p *Progress) UpdateBarPriority(b *Bar, priority int) {
	select {
	case p.operateState <- func(s *pState) { s.bHeap.update(b, priority) }:
	case <-p.done:
	}
}

// BarCount returns bars count
func (p *Progress) BarCount() int {
	result := make(chan int, 1)
	select {
	case p.operateState <- func(s *pState) { result <- s.bHeap.Len() }:
		return <-result
	case <-p.done:
		return 0
	}
}

// Wait first waits for user provided *sync.WaitGroup, if any,
// then waits far all bars to complete and finally shutdowns master goroutine.
// After this method has been called, there is no way to reuse *Progress instance.
func (p *Progress) Wait() {
	if p.uwg != nil {
		p.uwg.Wait()
	}

	p.wg.Wait()

	select {
	case p.operateState <- func(s *pState) { s.zeroWait = true }:
		<-p.done
	case <-p.done:
	}
}

func newWidthSyncer(timeout <-chan struct{}, numBars, numColumn int) *widthSyncer {
	ws := &widthSyncer{
		Accumulator: make([]chan int, numColumn),
		Distributor: make([]chan int, numColumn),
	}
	for i := 0; i < numColumn; i++ {
		ws.Accumulator[i] = make(chan int, numBars)
		ws.Distributor[i] = make(chan int, numBars)
	}
	for i := 0; i < numColumn; i++ {
		go func(accumulator <-chan int, distributor chan<- int) {
			defer close(distributor)
			widths := make([]int, 0, numBars)
		loop:
			for {
				select {
				case w := <-accumulator:
					widths = append(widths, w)
					if len(widths) == numBars {
						break loop
					}
				case <-timeout:
					if len(widths) == 0 {
						return
					}
					break loop
				}
			}
			maxWidth := calcMax(widths)
			for i := 0; i < len(widths); i++ {
				distributor <- maxWidth
			}
		}(ws.Accumulator[i], ws.Distributor[i])
	}
	return ws
}

func (s *pState) writeAndFlush(tw, numP, numA int) (err error) {
	timeout := make(chan struct{})
	pSyncer := newWidthSyncer(timeout, s.bHeap.Len(), numP)
	aSyncer := newWidthSyncer(timeout, s.bHeap.Len(), numA)
	time.AfterFunc(s.rr-s.rr/12, func() {
		close(timeout)
	})

	for _, ch := range s.renderByPriority(tw, pSyncer, aSyncer) {
		bf := <-ch
		_, err = s.cw.ReadFrom(bf.reader)
		if !bf.bar.completed && bf.toComplete {
			// shutdown at next flush, in other words decrement underlying WaitGroup
			// only after the bar with completed state has been flushed.
			// this ensures no bar ends up with less than 100% rendered.
			defer func() {
				s.shutdownPending = append(s.shutdownPending, bf.bar)
			}()
			if bf.bar.removeOnComplete {
				s.heapUpdated = heap.Remove(s.bHeap, bf.bar.index) != nil
			}
			if replacementBar, ok := s.waitBars[bf.bar]; ok {
				heap.Push(s.bHeap, replacementBar)
				s.heapUpdated = true
				delete(s.waitBars, bf.bar)
			}
			bf.bar.completed = true
		}
	}

	for _, interceptor := range s.interceptors {
		interceptor(s.cw)
	}

	if e := s.cw.Flush(); err == nil {
		err = e
	}

	for i := len(s.shutdownPending) - 1; i >= 0; i-- {
		close(s.shutdownPending[i].shutdown)
		s.shutdownPending = s.shutdownPending[:i]
	}
	return
}

func (s *pState) renderByPriority(tw int, pSyncer, aSyncer *widthSyncer) []<-chan *bFrame {
	bff := make([]<-chan *bFrame, 0, s.bHeap.Len())
	for s.bHeap.Len() > 0 {
		b := heap.Pop(s.bHeap).(*Bar)
		defer heap.Push(s.bHeap, b)
		bff = append(bff, b.render(s.debugOut, tw, pSyncer, aSyncer))
	}
	return bff
}

func calcMax(slice []int) int {
	if len(slice) == 0 {
		return 0
	}

	max := slice[0]
	for i := 1; i < len(slice); i++ {
		if slice[i] > max {
			max = slice[i]
		}
	}
	return max
}
<|endoftext|>


<commit_before>package mpb

import (
	"bytes"
	"container/heap"
	"context"
	"fmt"
	"io"
	"math"
	"os"
	"sync"
	"time"

	"github.com/vbauerster/mpb/v8/cwriter"
)

const (
	prr = 150 * time.Millisecond // default RefreshRate
)

// DoneError represents an error when `*mpb.Progress` is done but its functionality is requested.
var DoneError = fmt.Errorf("%T instance can't be reused after it's done!", (*Progress)(nil))

// Progress represents a container that renders one or more progress bars.
type Progress struct {
	ctx          context.Context
	uwg          *sync.WaitGroup
	cwg          *sync.WaitGroup
	bwg          *sync.WaitGroup
	operateState chan func(*pState)
	interceptIo  chan func(io.Writer)
	done         chan struct{}
	refreshCh    chan time.Time
	once         sync.Once
}

// pState holds bars in its priorityQueue, it gets passed to (*Progress).serve monitor goroutine.
type pState struct {
	bHeap       priorityQueue
	heapUpdated bool
	pMatrix     map[int][]chan int
	aMatrix     map[int][]chan int

	// for reuse purposes
	rows []io.Reader
	pool []*Bar

	// following are provided/overrided by user
	idCount          int
	reqWidth         int
	popPriority      int
	popCompleted     bool
	outputDiscarded  bool
	rr               time.Duration
	uwg              *sync.WaitGroup
	externalRefresh  <-chan interface{}
	renderDelay      <-chan struct{}
	shutdownNotifier chan struct{}
	queueBars        map[*Bar]*Bar
	output           io.Writer
	debugOut         io.Writer
}

// New creates new Progress container instance. It's not possible to
// reuse instance after (*Progress).Wait method has been called.
func New(options ...ContainerOption) *Progress {
	return NewWithContext(context.Background(), options...)
}

// NewWithContext creates new Progress container instance with provided
// context. It's not possible to reuse instance after (*Progress).Wait
// method has been called.
func NewWithContext(ctx context.Context, options ...ContainerOption) *Progress {
	s := &pState{
		bHeap:       priorityQueue{},
		rows:        make([]io.Reader, 0, 64),
		pool:        make([]*Bar, 0, 64),
		rr:          prr,
		queueBars:   make(map[*Bar]*Bar),
		output:      os.Stdout,
		popPriority: math.MinInt32,
	}

	for _, opt := range options {
		if opt != nil {
			opt(s)
		}
	}

	p := &Progress{
		ctx:          ctx,
		uwg:          s.uwg,
		cwg:          new(sync.WaitGroup),
		bwg:          new(sync.WaitGroup),
		operateState: make(chan func(*pState)),
		interceptIo:  make(chan func(io.Writer)),
		done:         make(chan struct{}),
	}

	p.cwg.Add(1)
	go p.serve(s, cwriter.New(s.output))
	return p
}

// AddBar creates a bar with default bar filler.
func (p *Progress) AddBar(total int64, options ...BarOption) *Bar {
	return p.New(total, BarStyle(), options...)
}

// AddSpinner creates a bar with default spinner filler.
func (p *Progress) AddSpinner(total int64, options ...BarOption) *Bar {
	return p.New(total, SpinnerStyle(), options...)
}

// New creates a bar with provided BarFillerBuilder.
func (p *Progress) New(total int64, builder BarFillerBuilder, options ...BarOption) *Bar {
	return p.Add(total, builder.Build(), options...)
}

// Add creates a bar which renders itself by provided filler.
// If `total <= 0` triggering complete event by increment methods is disabled.
// Panics if *Progress instance is done, i.e. called after (*Progress).Wait().
func (p *Progress) Add(total int64, filler BarFiller, options ...BarOption) *Bar {
	if filler == nil {
		filler = NopStyle().Build()
	}
	p.bwg.Add(1)
	result := make(chan *Bar)
	select {
	case p.operateState <- func(ps *pState) {
		bs := ps.makeBarState(total, filler, options...)
		bar := newBar(p, bs)
		if bs.wait.bar != nil {
			ps.queueBars[bs.wait.bar] = bar
		} else {
			heap.Push(&ps.bHeap, bar)
			ps.heapUpdated = true
		}
		ps.idCount++
		result <- bar
	}:
		bar := <-result
		return bar
	case <-p.done:
		p.bwg.Done()
		panic(DoneError)
	}
}

func (p *Progress) traverseBars(cb func(b *Bar) bool) {
	sync := make(chan struct{})
	select {
	case p.operateState <- func(s *pState) {
		for i := 0; i < s.bHeap.Len(); i++ {
			bar := s.bHeap[i]
			if !cb(bar) {
				break
			}
		}
		close(sync)
	}:
		<-sync
	case <-p.done:
	}
}

// UpdateBarPriority same as *Bar.SetPriority(int).
func (p *Progress) UpdateBarPriority(b *Bar, priority int) {
	select {
	case p.operateState <- func(s *pState) {
		if b.index < 0 {
			return
		}
		b.priority = priority
		heap.Fix(&s.bHeap, b.index)
	}:
	case <-p.done:
	}
}

// BarCount returns bars count.
func (p *Progress) BarCount() int {
	result := make(chan int)
	select {
	case p.operateState <- func(s *pState) { result <- s.bHeap.Len() }:
		return <-result
	case <-p.done:
		return 0
	}
}

// Write is implementation of io.Writer.
// Writing to `*mpb.Progress` will print lines above a running bar.
// Writes aren't flushed immediatly, but at next refresh cycle.
// If Write is called after `*mpb.Progress` is done, `mpb.DoneError`
// is returned.
func (p *Progress) Write(b []byte) (int, error) {
	type result struct {
		n   int
		err error
	}
	ch := make(chan *result)
	select {
	case p.interceptIo <- func(w io.Writer) {
		n, err := w.Write(b)
		ch <- &result{n, err}
	}:
		res := <-ch
		return res.n, res.err
	case <-p.done:
		return 0, DoneError
	}
}

// Wait waits for all bars to complete and finally shutdowns container.
// After this method has been called, there is no way to reuse *Progress
// instance.
func (p *Progress) Wait() {
	// wait for user wg, if any
	if p.uwg != nil {
		p.uwg.Wait()
	}

	// wait for bars to quit, if any
	p.bwg.Wait()

	p.once.Do(p.shutdown)

	// wait for container to quit
	p.cwg.Wait()
}

func (p *Progress) shutdown() {
	close(p.done)
}

func (p *Progress) serve(s *pState, cw *cwriter.Writer) {
	defer p.cwg.Done()

	render := func(out io.Writer) {
		err := s.render(cw)
		for err != nil {
			if out != nil {
				_, err = fmt.Fprintln(out, err)
			} else {
				panic(err)
			}
			out = nil
		}
	}

	p.refreshCh = s.newTicker(p.done)

	for {
		select {
		case op := <-p.operateState:
			op(s)
		case fn := <-p.interceptIo:
			fn(cw)
		case <-p.refreshCh:
			render(s.debugOut)
		case <-s.shutdownNotifier:
			for s.heapUpdated {
				render(s.debugOut)
			}
			return
		}
	}
}

func (s *pState) render(cw *cwriter.Writer) error {
	if s.heapUpdated {
		s.updateSyncMatrix()
		s.heapUpdated = false
	}
	syncWidth(s.pMatrix)
	syncWidth(s.aMatrix)

	width, height, err := cw.GetTermSize()
	if err != nil {
		width = s.reqWidth
		height = s.bHeap.Len()
	}
	for i := 0; i < s.bHeap.Len(); i++ {
		bar := s.bHeap[i]
		go bar.render(width)
	}

	return s.flush(cw, height)
}

func (s *pState) flush(cw *cwriter.Writer, height int) error {
	var wg sync.WaitGroup
	var popCount int

	for s.bHeap.Len() > 0 {
		var usedRows int
		b := heap.Pop(&s.bHeap).(*Bar)
		frame := <-b.frameCh
		if frame.recovered {
			s.heapUpdated = true
			continue
		}
		for i := len(frame.rows) - 1; i >= 0; i-- {
			if row := frame.rows[i]; len(s.rows) < height {
				s.rows = append(s.rows, row)
				usedRows++
			} else {
				wg.Add(1)
				go func() {
					_, _ = io.Copy(io.Discard, row)
					wg.Done()
				}()
			}
		}
		if frame.shutdown != 0 {
			b.Wait() // waiting for b.done, so it's safe to read b.bs
			drop := b.bs.dropOnComplete
			if qb, ok := s.queueBars[b]; ok {
				delete(s.queueBars, b)
				qb.priority = b.priority
				s.pool = append(s.pool, qb)
				drop = true
			} else if s.popCompleted && !b.bs.noPop {
				if frame.shutdown > 1 {
					popCount += usedRows
					drop = true
				} else {
					s.popPriority++
					b.priority = s.popPriority
				}
			}
			if drop {
				s.heapUpdated = true
				continue
			}
		}
		s.pool = append(s.pool, b)
	}

	wg.Add(1)
	go func() {
		for _, b := range s.pool {
			heap.Push(&s.bHeap, b)
		}
		wg.Done()
	}()

	readRows := len(s.rows)
	for i := readRows - 1; i >= 0; i-- {
		_, err := cw.ReadFrom(s.rows[i])
		if err != nil {
			if s.debugOut != nil {
				fmt.Fprintf(s.debugOut, "cw.ReadFrom: %s\n", err.Error())
			}
			readRows--
		}
	}

	err := cw.Flush(readRows - popCount)
	wg.Wait()
	s.rows = s.rows[:0]
	s.pool = s.pool[:0]
	return err
}

func (s *pState) newTicker(done <-chan struct{}) chan time.Time {
	ch := make(chan time.Time)
	if s.shutdownNotifier == nil {
		s.shutdownNotifier = make(chan struct{})
	}
	go func() {
		if s.renderDelay != nil {
			<-s.renderDelay
		}
		var internalRefresh <-chan time.Time
		if !s.outputDiscarded {
			if s.externalRefresh == nil {
				ticker := time.NewTicker(s.rr)
				defer ticker.Stop()
				internalRefresh = ticker.C
			}
		} else {
			s.externalRefresh = nil
		}
		for {
			select {
			case t := <-internalRefresh:
				ch <- t
			case x := <-s.externalRefresh:
				if t, ok := x.(time.Time); ok {
					ch <- t
				} else {
					ch <- time.Now()
				}
			case <-done:
				close(s.shutdownNotifier)
				return
			}
		}
	}()
	return ch
}

func (s *pState) updateSyncMatrix() {
	s.pMatrix = make(map[int][]chan int)
	s.aMatrix = make(map[int][]chan int)
	for i := 0; i < s.bHeap.Len(); i++ {
		bar := s.bHeap[i]
		table := bar.wSyncTable()
		pRow, aRow := table[0], table[1]

		for i, ch := range pRow {
			s.pMatrix[i] = append(s.pMatrix[i], ch)
		}

		for i, ch := range aRow {
			s.aMatrix[i] = append(s.aMatrix[i], ch)
		}
	}
}

func (s *pState) makeBarState(total int64, filler BarFiller, options ...BarOption) *bState {
	bs := &bState{
		id:       s.idCount,
		priority: s.idCount,
		reqWidth: s.reqWidth,
		total:    total,
		filler:   filler,
		debugOut: s.debugOut,
	}

	if total > 0 {
		bs.triggerComplete = true
	}

	for _, opt := range options {
		if opt != nil {
			opt(bs)
		}
	}

	if bs.middleware != nil {
		bs.filler = bs.middleware(filler)
		bs.middleware = nil
	}

	for i := 0; i < len(bs.buffers); i++ {
		bs.buffers[i] = bytes.NewBuffer(make([]byte, 0, 512))
	}

	bs.subscribeDecorators()

	return bs
}

func syncWidth(matrix map[int][]chan int) {
	for _, column := range matrix {
		go maxWidthDistributor(column)
	}
}

func maxWidthDistributor(column []chan int) {
	var maxWidth int
	for _, ch := range column {
		if w := <-ch; w > maxWidth {
			maxWidth = w
		}
	}
	for _, ch := range column {
		ch <- maxWidth
	}
}
<commit_msg>*Progress.Add is redundant, removing<commit_after>package mpb

import (
	"bytes"
	"container/heap"
	"context"
	"fmt"
	"io"
	"math"
	"os"
	"sync"
	"time"

	"github.com/vbauerster/mpb/v8/cwriter"
)

const (
	prr = 150 * time.Millisecond // default RefreshRate
)

// DoneError represents an error when `*mpb.Progress` is done but its functionality is requested.
var DoneError = fmt.Errorf("%T instance can't be reused after it's done!", (*Progress)(nil))

// Progress represents a container that renders one or more progress bars.
type Progress struct {
	ctx          context.Context
	uwg          *sync.WaitGroup
	cwg          *sync.WaitGroup
	bwg          *sync.WaitGroup
	operateState chan func(*pState)
	interceptIo  chan func(io.Writer)
	done         chan struct{}
	refreshCh    chan time.Time
	once         sync.Once
}

// pState holds bars in its priorityQueue, it gets passed to (*Progress).serve monitor goroutine.
type pState struct {
	bHeap       priorityQueue
	heapUpdated bool
	pMatrix     map[int][]chan int
	aMatrix     map[int][]chan int

	// for reuse purposes
	rows []io.Reader
	pool []*Bar

	// following are provided/overrided by user
	idCount          int
	reqWidth         int
	popPriority      int
	popCompleted     bool
	outputDiscarded  bool
	rr               time.Duration
	uwg              *sync.WaitGroup
	externalRefresh  <-chan interface{}
	renderDelay      <-chan struct{}
	shutdownNotifier chan struct{}
	queueBars        map[*Bar]*Bar
	output           io.Writer
	debugOut         io.Writer
}

// New creates new Progress container instance. It's not possible to
// reuse instance after (*Progress).Wait method has been called.
func New(options ...ContainerOption) *Progress {
	return NewWithContext(context.Background(), options...)
}

// NewWithContext creates new Progress container instance with provided
// context. It's not possible to reuse instance after (*Progress).Wait
// method has been called.
func NewWithContext(ctx context.Context, options ...ContainerOption) *Progress {
	s := &pState{
		bHeap:       priorityQueue{},
		rows:        make([]io.Reader, 0, 64),
		pool:        make([]*Bar, 0, 64),
		rr:          prr,
		queueBars:   make(map[*Bar]*Bar),
		output:      os.Stdout,
		popPriority: math.MinInt32,
	}

	for _, opt := range options {
		if opt != nil {
			opt(s)
		}
	}

	p := &Progress{
		ctx:          ctx,
		uwg:          s.uwg,
		cwg:          new(sync.WaitGroup),
		bwg:          new(sync.WaitGroup),
		operateState: make(chan func(*pState)),
		interceptIo:  make(chan func(io.Writer)),
		done:         make(chan struct{}),
	}

	p.cwg.Add(1)
	go p.serve(s, cwriter.New(s.output))
	return p
}

// AddBar creates a bar with default bar filler.
func (p *Progress) AddBar(total int64, options ...BarOption) *Bar {
	return p.New(total, BarStyle(), options...)
}

// AddSpinner creates a bar with default spinner filler.
func (p *Progress) AddSpinner(total int64, options ...BarOption) *Bar {
	return p.New(total, SpinnerStyle(), options...)
}

// New creates a bar which renders itself by `BarFiller`, which is obtained
// via provided `BarFillerBuilder`. If `total <= 0` triggering complete event
// by increment methods is disabled. Panics if *Progress instance is done,
// i.e. called after (*Progress).Wait().
func (p *Progress) New(total int64, builder BarFillerBuilder, options ...BarOption) *Bar {
	filler := builder.Build()
	if filler == nil {
		filler = NopStyle().Build()
	}
	p.bwg.Add(1)
	result := make(chan *Bar)
	select {
	case p.operateState <- func(ps *pState) {
		bs := ps.makeBarState(total, filler, options...)
		bar := newBar(p, bs)
		if bs.wait.bar != nil {
			ps.queueBars[bs.wait.bar] = bar
		} else {
			heap.Push(&ps.bHeap, bar)
			ps.heapUpdated = true
		}
		ps.idCount++
		result <- bar
	}:
		bar := <-result
		return bar
	case <-p.done:
		p.bwg.Done()
		panic(DoneError)
	}
}

func (p *Progress) traverseBars(cb func(b *Bar) bool) {
	sync := make(chan struct{})
	select {
	case p.operateState <- func(s *pState) {
		for i := 0; i < s.bHeap.Len(); i++ {
			bar := s.bHeap[i]
			if !cb(bar) {
				break
			}
		}
		close(sync)
	}:
		<-sync
	case <-p.done:
	}
}

// UpdateBarPriority same as *Bar.SetPriority(int).
func (p *Progress) UpdateBarPriority(b *Bar, priority int) {
	select {
	case p.operateState <- func(s *pState) {
		if b.index < 0 {
			return
		}
		b.priority = priority
		heap.Fix(&s.bHeap, b.index)
	}:
	case <-p.done:
	}
}

// BarCount returns bars count.
func (p *Progress) BarCount() int {
	result := make(chan int)
	select {
	case p.operateState <- func(s *pState) { result <- s.bHeap.Len() }:
		return <-result
	case <-p.done:
		return 0
	}
}

// Write is implementation of io.Writer.
// Writing to `*mpb.Progress` will print lines above a running bar.
// Writes aren't flushed immediatly, but at next refresh cycle.
// If Write is called after `*mpb.Progress` is done, `mpb.DoneError`
// is returned.
func (p *Progress) Write(b []byte) (int, error) {
	type result struct {
		n   int
		err error
	}
	ch := make(chan *result)
	select {
	case p.interceptIo <- func(w io.Writer) {
		n, err := w.Write(b)
		ch <- &result{n, err}
	}:
		res := <-ch
		return res.n, res.err
	case <-p.done:
		return 0, DoneError
	}
}

// Wait waits for all bars to complete and finally shutdowns container.
// After this method has been called, there is no way to reuse *Progress
// instance.
func (p *Progress) Wait() {
	// wait for user wg, if any
	if p.uwg != nil {
		p.uwg.Wait()
	}

	// wait for bars to quit, if any
	p.bwg.Wait()

	p.once.Do(p.shutdown)

	// wait for container to quit
	p.cwg.Wait()
}

func (p *Progress) shutdown() {
	close(p.done)
}

func (p *Progress) serve(s *pState, cw *cwriter.Writer) {
	defer p.cwg.Done()

	render := func(out io.Writer) {
		err := s.render(cw)
		for err != nil {
			if out != nil {
				_, err = fmt.Fprintln(out, err)
			} else {
				panic(err)
			}
			out = nil
		}
	}

	p.refreshCh = s.newTicker(p.done)

	for {
		select {
		case op := <-p.operateState:
			op(s)
		case fn := <-p.interceptIo:
			fn(cw)
		case <-p.refreshCh:
			render(s.debugOut)
		case <-s.shutdownNotifier:
			for s.heapUpdated {
				render(s.debugOut)
			}
			return
		}
	}
}

func (s *pState) render(cw *cwriter.Writer) error {
	if s.heapUpdated {
		s.updateSyncMatrix()
		s.heapUpdated = false
	}
	syncWidth(s.pMatrix)
	syncWidth(s.aMatrix)

	width, height, err := cw.GetTermSize()
	if err != nil {
		width = s.reqWidth
		height = s.bHeap.Len()
	}
	for i := 0; i < s.bHeap.Len(); i++ {
		bar := s.bHeap[i]
		go bar.render(width)
	}

	return s.flush(cw, height)
}

func (s *pState) flush(cw *cwriter.Writer, height int) error {
	var wg sync.WaitGroup
	var popCount int

	for s.bHeap.Len() > 0 {
		var usedRows int
		b := heap.Pop(&s.bHeap).(*Bar)
		frame := <-b.frameCh
		if frame.recovered {
			s.heapUpdated = true
			continue
		}
		for i := len(frame.rows) - 1; i >= 0; i-- {
			if row := frame.rows[i]; len(s.rows) < height {
				s.rows = append(s.rows, row)
				usedRows++
			} else {
				wg.Add(1)
				go func() {
					_, _ = io.Copy(io.Discard, row)
					wg.Done()
				}()
			}
		}
		if frame.shutdown != 0 {
			b.Wait() // waiting for b.done, so it's safe to read b.bs
			drop := b.bs.dropOnComplete
			if qb, ok := s.queueBars[b]; ok {
				delete(s.queueBars, b)
				qb.priority = b.priority
				s.pool = append(s.pool, qb)
				drop = true
			} else if s.popCompleted && !b.bs.noPop {
				if frame.shutdown > 1 {
					popCount += usedRows
					drop = true
				} else {
					s.popPriority++
					b.priority = s.popPriority
				}
			}
			if drop {
				s.heapUpdated = true
				continue
			}
		}
		s.pool = append(s.pool, b)
	}

	wg.Add(1)
	go func() {
		for _, b := range s.pool {
			heap.Push(&s.bHeap, b)
		}
		wg.Done()
	}()

	readRows := len(s.rows)
	for i := readRows - 1; i >= 0; i-- {
		_, err := cw.ReadFrom(s.rows[i])
		if err != nil {
			if s.debugOut != nil {
				fmt.Fprintf(s.debugOut, "cw.ReadFrom: %s\n", err.Error())
			}
			readRows--
		}
	}

	err := cw.Flush(readRows - popCount)
	wg.Wait()
	s.rows = s.rows[:0]
	s.pool = s.pool[:0]
	return err
}

func (s *pState) newTicker(done <-chan struct{}) chan time.Time {
	ch := make(chan time.Time)
	if s.shutdownNotifier == nil {
		s.shutdownNotifier = make(chan struct{})
	}
	go func() {
		if s.renderDelay != nil {
			<-s.renderDelay
		}
		var internalRefresh <-chan time.Time
		if !s.outputDiscarded {
			if s.externalRefresh == nil {
				ticker := time.NewTicker(s.rr)
				defer ticker.Stop()
				internalRefresh = ticker.C
			}
		} else {
			s.externalRefresh = nil
		}
		for {
			select {
			case t := <-internalRefresh:
				ch <- t
			case x := <-s.externalRefresh:
				if t, ok := x.(time.Time); ok {
					ch <- t
				} else {
					ch <- time.Now()
				}
			case <-done:
				close(s.shutdownNotifier)
				return
			}
		}
	}()
	return ch
}

func (s *pState) updateSyncMatrix() {
	s.pMatrix = make(map[int][]chan int)
	s.aMatrix = make(map[int][]chan int)
	for i := 0; i < s.bHeap.Len(); i++ {
		bar := s.bHeap[i]
		table := bar.wSyncTable()
		pRow, aRow := table[0], table[1]

		for i, ch := range pRow {
			s.pMatrix[i] = append(s.pMatrix[i], ch)
		}

		for i, ch := range aRow {
			s.aMatrix[i] = append(s.aMatrix[i], ch)
		}
	}
}

func (s *pState) makeBarState(total int64, filler BarFiller, options ...BarOption) *bState {
	bs := &bState{
		id:       s.idCount,
		priority: s.idCount,
		reqWidth: s.reqWidth,
		total:    total,
		filler:   filler,
		debugOut: s.debugOut,
	}

	if total > 0 {
		bs.triggerComplete = true
	}

	for _, opt := range options {
		if opt != nil {
			opt(bs)
		}
	}

	if bs.middleware != nil {
		bs.filler = bs.middleware(filler)
		bs.middleware = nil
	}

	for i := 0; i < len(bs.buffers); i++ {
		bs.buffers[i] = bytes.NewBuffer(make([]byte, 0, 512))
	}

	bs.subscribeDecorators()

	return bs
}

func syncWidth(matrix map[int][]chan int) {
	for _, column := range matrix {
		go maxWidthDistributor(column)
	}
}

func maxWidthDistributor(column []chan int) {
	var maxWidth int
	for _, ch := range column {
		if w := <-ch; w > maxWidth {
			maxWidth = w
		}
	}
	for _, ch := range column {
		ch <- maxWidth
	}
}
<|endoftext|>