<commit_before>// Package trillian_client provides some useful utilities for
// interacting with Trillian.
package trillian_client

import (
	"context"
	"fmt"
	"log"

	"github.com/google/trillian"
	"google.golang.org/grpc"
)

const chunk = 10

// A type that is passed to TrillianClient.Scan(). Leaf() is called on
// it for each leaf in the log.
type LogScanner interface {
	Leaf(leaf *trillian.LogLeaf) error
}

// A Trillian client. Create a new one with trillian_client.New().
type TrillianClient interface {
	Scan(logID int64, s LogScanner) error
	Close()
}

type trillianClient struct {
	g  *grpc.ClientConn
	tc trillian.TrillianLogClient
}

// New creates and connects new TrillianClient, given the URL of the
// Trillian server.
func New(logAddr string) TrillianClient {
	g, err := grpc.Dial(logAddr, grpc.WithInsecure())
	if err != nil {
		log.Fatalf("Failed to dial Trillian Log: %v", err)
	}

	tc := trillian.NewTrillianLogClient(g)

	return &trillianClient{g, tc}
}

func (t *trillianClient) Scan(logID int64, s LogScanner) error {
	ctx := context.Background()

	rr := &trillian.GetLatestSignedLogRootRequest{LogId: logID}
	lr, err := t.tc.GetLatestSignedLogRoot(ctx, rr)
	if err != nil {
		log.Fatalf("Can't get log root: %v", err)
	}

	ts := lr.SignedLogRoot.TreeSize
	for n := int64(0); n < ts; {
		g := &trillian.GetLeavesByRangeRequest{LogId: logID, StartIndex: n, Count: chunk}
		r, err := t.tc.GetLeavesByRange(ctx, g)
		if err != nil {
			return fmt.Errorf("Can't get leaf %d: %v", n, err)
		}

		// Deal with server skew, if tree size has reduced.
		// Don't allow increases so this terminates eventually.
		rts := r.SignedLogRoot.TreeSize
		if rts < ts {
			ts = rts
		}

		if n < ts && len(r.Leaves) == 0 {
			return fmt.Errorf("No progress at leaf %d", n)
		}

		for m := 0; m < len(r.Leaves) && n < ts; n++ {
			if r.Leaves[m] == nil {
				return fmt.Errorf("Can't get leaf %d (no error)", n)
			}
			if r.Leaves[m].LeafIndex != n {
				return fmt.Errorf("Got index %d expected %d", r.Leaves[n].LeafIndex, n)
			}
			err := s.Leaf(r.Leaves[m])
			if err != nil {
				return err
			}
			m++
		}
	}
	return nil
}

func (t *trillianClient) Close() {
	t.g.Close()
}
<commit_msg>Remove deprecated slr (#53)<commit_after>// Package trillian_client provides some useful utilities for
// interacting with Trillian.
package trillian_client

import (
	"context"
	"fmt"
	"log"

	"github.com/google/trillian"
	"github.com/google/trillian/types"
	"google.golang.org/grpc"
)

const chunk = 10

// A type that is passed to TrillianClient.Scan(). Leaf() is called on
// it for each leaf in the log.
type LogScanner interface {
	Leaf(leaf *trillian.LogLeaf) error
}

// A Trillian client. Create a new one with trillian_client.New().
type TrillianClient interface {
	Scan(logID int64, s LogScanner) error
	Close()
}

type trillianClient struct {
	g  *grpc.ClientConn
	tc trillian.TrillianLogClient
}

// New creates and connects new TrillianClient, given the URL of the
// Trillian server.
func New(logAddr string) TrillianClient {
	g, err := grpc.Dial(logAddr, grpc.WithInsecure())
	if err != nil {
		log.Fatalf("Failed to dial Trillian Log: %v", err)
	}

	tc := trillian.NewTrillianLogClient(g)

	return &trillianClient{g, tc}
}

func (t *trillianClient) Scan(logID int64, s LogScanner) error {
	ctx := context.Background()

	rr := &trillian.GetLatestSignedLogRootRequest{LogId: logID}
	lr, err := t.tc.GetLatestSignedLogRoot(ctx, rr)
	if err != nil {
		log.Fatalf("Can't get log root: %v", err)
	}

	var root types.LogRootV1
	// TODO(Martin2112): Verify root signature.
	if err := root.UnmarshalBinary(lr.SignedLogRoot.LogRoot); err != nil {
		return fmt.Errorf("Root failed to unmarshal: %v", err)
	}

	ts := root.TreeSize
	for n := uint64(0); n < ts; {
		g := &trillian.GetLeavesByRangeRequest{LogId: logID, StartIndex: int64(n), Count: chunk}
		r, err := t.tc.GetLeavesByRange(ctx, g)
		if err != nil {
			return fmt.Errorf("Can't get leaf %d: %v", n, err)
		}

		// Deal with server skew, if tree size has reduced.
		// Don't allow increases so this terminates eventually.
		rts := root.TreeSize
		if rts < ts {
			ts = rts
		}

		if n < ts && len(r.Leaves) == 0 {
			return fmt.Errorf("No progress at leaf %d", n)
		}

		for m := 0; m < len(r.Leaves) && n < ts; n++ {
			if r.Leaves[m] == nil {
				return fmt.Errorf("Can't get leaf %d (no error)", n)
			}
			if uint64(r.Leaves[m].LeafIndex) != n {
				return fmt.Errorf("Got index %d expected %d", r.Leaves[n].LeafIndex, n)
			}
			err := s.Leaf(r.Leaves[m])
			if err != nil {
				return err
			}
			m++
		}
	}
	return nil
}

func (t *trillianClient) Close() {
	t.g.Close()
}
<|endoftext|>


<commit_before>package timeseries

import (
	"container/list"
	"time"
)

func Abs(a int) int {
	if a < 0 {
		return -a
	}
	return a
}

func minTime(a, b time.Time) time.Time {
	if a.Before(b) {
		return a
	}
	return b
}

func maxTime(a, b time.Time) time.Time {
	if a.After(b) {
		return a
	}
	return b
}

type Primitive interface {
	Add(other Primitive)
	CopyFrom(other Primitive)
	Reset()
}

type Integer int

func NewInteger() Primitive                 { i := Integer(0); return &i }
func (i *Integer) Value() int               { return int(*i) }
func (i *Integer) Add(other Primitive)      { *i += *(other.(*Integer)) }
func (i *Integer) CopyFrom(other Primitive) { *i = *(other.(*Integer)) }
func (i *Integer) Reset()                   { *i = 0 }

const (
	ResolutionOneSecond  = 1 * time.Second
	ResolutionTenSeconds = 10 * time.Second
	ResolutionOneMinute  = 1 * time.Minute
	ResolutionTenMinutes = 10 * time.Minute
	ResolutionOneHour    = 1 * time.Hour
	ResolutionSixHours   = 6 * time.Hour
	ResolutionOneDay     = 24 * time.Hour
	ResolutionOneWeek    = 7 * 24 * time.Hour
	ResolutionFourWeeks  = 4 * 7 * 24 * time.Hour
)

type dataStream struct {
	primitiveFunc func() Primitive
	buckets       *list.List
	resolution    time.Duration
	beginTime     time.Time
	endTime       time.Time
}

func (ds *dataStream) initialize(p func() Primitive, resolution time.Duration) {
	ds.primitiveFunc = p
	ds.resolution = resolution
	ds.buckets = list.New()
}

func (ds *dataStream) NumBuckets() int {
	return ds.buckets.Len()
}

func (ds *dataStream) reset() {
	ds.beginTime = time.Time{}
	ds.endTime = time.Time{}

	var next *list.Element
	for e := ds.buckets.Front(); e != nil; e = next {
		next = e.Next()
		ds.buckets.Remove(e)
	}
}

type TimeSeries struct {
	primitiveFunc func() Primitive
	dataStreams   []*dataStream
	total         Primitive
}

func NewTimeSeries(p func() Primitive, resolutions []time.Duration) *TimeSeries {
	timeSeries := new(TimeSeries)
	timeSeries.initialize(p, resolutions)
	return timeSeries
}

func (ts *TimeSeries) initialize(p func() Primitive, resolutions []time.Duration) {
	ts.primitiveFunc = p
	ts.total = ts.primitiveFunc()
	ts.dataStreams = make([]*dataStream, len(resolutions))
	for i := range resolutions {
		ts.dataStreams[i] = new(dataStream)
		ts.dataStreams[i].initialize(p, resolutions[i])
	}
	ts.reset()
}

func (ts *TimeSeries) reset() {
	ts.total.Reset()
	for i := range ts.dataStreams {
		ts.dataStreams[i].reset()
	}
}

func (ts *TimeSeries) Add(d Primitive, t time.Time) {
	for _, ds := range ts.dataStreams {
		isFirstAdd := ds.buckets.Len() == 0
		if isFirstAdd {
			ds.beginTime = t
			ds.endTime = t

			first := ds.primitiveFunc()
			ds.buckets.PushBack(first)
		}

		bucketIdxFromEnd := int(t.Sub(ds.endTime) / ds.resolution)
		for i := 0; i < bucketIdxFromEnd; i++ {
			p := ds.primitiveFunc()
			ds.buckets.PushBack(p)
		}

		lastBucket := ds.buckets.Back().Value.(Primitive)
		lastBucket.Add(d)

		// update begin and end time
		ds.beginTime = minTime(ds.beginTime, t)
		ds.endTime = maxTime(ds.endTime, t)
	}
	ts.total.Add(d)
}

func (ts *TimeSeries) Total() Primitive {
	return ts.total
}

func (ts *TimeSeries) Range(resolutionIdx int, fromTime, toTime time.Time) []Primitive {
	ds := ts.dataStreams[resolutionIdx]

	beginBucketIdx := int(fromTime.Sub(ds.beginTime) / ds.resolution)
	endBucketIdx := int(toTime.Sub(ds.beginTime) / ds.resolution)
	filteredBuckets := list.New()

	iterIdx := 0
	for e := ds.buckets.Front(); e != nil; e = e.Next() {
		if beginBucketIdx <= iterIdx && iterIdx <= endBucketIdx {
			filteredBuckets.PushBack(e.Value)
		}
		iterIdx++
	}

	filtered := make([]Primitive, filteredBuckets.Len())
	insertIdx := 0
	for e := filteredBuckets.Front(); e != nil; e = e.Next() {
		filtered[insertIdx] = e.Value.(Primitive)
		insertIdx++
	}

	return filtered
}
<commit_msg>bug fix in bucket gap filling<commit_after>package timeseries

import (
	"container/list"
	"time"
)

func Abs(a int) int {
	if a < 0 {
		return -a
	}
	return a
}

func minTime(a, b time.Time) time.Time {
	if a.Before(b) {
		return a
	}
	return b
}

func maxTime(a, b time.Time) time.Time {
	if a.After(b) {
		return a
	}
	return b
}

type Primitive interface {
	Add(other Primitive)
	CopyFrom(other Primitive)
	Reset()
}

type Integer int

func NewInteger() Primitive                 { i := Integer(0); return &i }
func (i *Integer) Value() int               { return int(*i) }
func (i *Integer) Add(other Primitive)      { *i += *(other.(*Integer)) }
func (i *Integer) CopyFrom(other Primitive) { *i = *(other.(*Integer)) }
func (i *Integer) Reset()                   { *i = 0 }

const (
	ResolutionOneSecond  = 1 * time.Second
	ResolutionTenSeconds = 10 * time.Second
	ResolutionOneMinute  = 1 * time.Minute
	ResolutionTenMinutes = 10 * time.Minute
	ResolutionOneHour    = 1 * time.Hour
	ResolutionSixHours   = 6 * time.Hour
	ResolutionOneDay     = 24 * time.Hour
	ResolutionOneWeek    = 7 * 24 * time.Hour
	ResolutionFourWeeks  = 4 * 7 * 24 * time.Hour
)

type dataStream struct {
	primitiveFunc func() Primitive
	buckets       *list.List
	resolution    time.Duration
	beginTime     time.Time
	endTime       time.Time
}

func (ds *dataStream) initialize(p func() Primitive, resolution time.Duration) {
	ds.primitiveFunc = p
	ds.resolution = resolution
	ds.buckets = list.New()
}

func (ds *dataStream) NumBuckets() int {
	return ds.buckets.Len()
}

func (ds *dataStream) reset() {
	ds.beginTime = time.Time{}
	ds.endTime = time.Time{}

	var next *list.Element
	for e := ds.buckets.Front(); e != nil; e = next {
		next = e.Next()
		ds.buckets.Remove(e)
	}
}

type TimeSeries struct {
	primitiveFunc func() Primitive
	dataStreams   []*dataStream
	total         Primitive
}

func NewTimeSeries(p func() Primitive, resolutions []time.Duration) *TimeSeries {
	timeSeries := new(TimeSeries)
	timeSeries.initialize(p, resolutions)
	return timeSeries
}

func (ts *TimeSeries) initialize(p func() Primitive, resolutions []time.Duration) {
	ts.primitiveFunc = p
	ts.total = ts.primitiveFunc()
	ts.dataStreams = make([]*dataStream, len(resolutions))
	for i := range resolutions {
		ts.dataStreams[i] = new(dataStream)
		ts.dataStreams[i].initialize(p, resolutions[i])
	}
	ts.reset()
}

func (ts *TimeSeries) reset() {
	ts.total.Reset()
	for i := range ts.dataStreams {
		ts.dataStreams[i].reset()
	}
}

func (ts *TimeSeries) Add(d Primitive, t time.Time) {
	for _, ds := range ts.dataStreams {
		isFirstAdd := ds.buckets.Len() == 0
		if isFirstAdd {
			ds.beginTime = t
			ds.endTime = t

			first := ds.primitiveFunc()
			ds.buckets.PushBack(first)
		}

		bucketIdxAtEnd := int(ds.endTime.Sub(ds.beginTime) / ds.resolution)
		bucketIdxFromBegin := int(t.Sub(ds.beginTime) / ds.resolution)
		for i := 0; i < bucketIdxFromBegin-bucketIdxAtEnd; i++ {
			p := ds.primitiveFunc()
			ds.buckets.PushBack(p)
		}

		lastBucket := ds.buckets.Back().Value.(Primitive)
		lastBucket.Add(d)

		// update begin and end time
		ds.beginTime = minTime(ds.beginTime, t)
		ds.endTime = maxTime(ds.endTime, t)
	}
	ts.total.Add(d)
}

func (ts *TimeSeries) Total() Primitive {
	return ts.total
}

func (ts *TimeSeries) Range(resolutionIdx int, fromTime, toTime time.Time) []Primitive {
	ds := ts.dataStreams[resolutionIdx]

	beginBucketIdx := int(fromTime.Sub(ds.beginTime) / ds.resolution)
	endBucketIdx := int(toTime.Sub(ds.beginTime) / ds.resolution)
	filteredBuckets := list.New()

	iterIdx := 0
	for e := ds.buckets.Front(); e != nil; e = e.Next() {
		if beginBucketIdx <= iterIdx && iterIdx <= endBucketIdx {
			filteredBuckets.PushBack(e.Value)
		}
		iterIdx++
	}

	filtered := make([]Primitive, filteredBuckets.Len())
	insertIdx := 0
	for e := filteredBuckets.Front(); e != nil; e = e.Next() {
		filtered[insertIdx] = e.Value.(Primitive)
		insertIdx++
	}

	return filtered
}
<|endoftext|>


<commit_before>/*
 This package implements the leftpad function, inspired by the NPM (JS)
 package of the same name.

 Two functions are defined:

 import "leftpad"

 // pad with spaces
 str, err := LeftPad(s, n)

 // pad with specified character
 str, err := func LeftPadStr(s, n, c)

*/
package leftpad

import (
	"errors"
	"fmt"
	"strings"
)

// ErrInvalidChar indicates the pad character is too long.
var ErrInvalidChar = errors.New("Invalid character")

func doLeftPad(s string, n int, c string) (string, error) {
	if n < 0 {
		return "", errors.New(fmt.Sprintf("Invalid length %d", n))
	}

	if len(c) != 1 {
		return "", ErrInvalidChar
	}

	toAdd := n - len(s)
	if toAdd <= 0 {
		return s, nil
	}

	return strings.Repeat(c, toAdd) + s, nil
}

// LeftPad left-pads s with spaces, to length n.
// If n is smaller than s, LeftPad is a no-op.
func LeftPad(s string, n int) (string, error) {
	return doLeftPad(s, n, " ")
}

// LeftPadStr left-pads s with the char c, to length n.
// If n is smaller than s, LeftPadStr is a no-op.
func LeftPadStr(s string, n int, c string) (string, error) {
	return doLeftPad(s, n, c)
}
<commit_msg>Paddng takes a rune, rather than a string<commit_after>/*
 This package implements the leftpad function, inspired by the NPM (JS)
 package of the same name.

 Two functions are defined:

 import "leftpad"

 // pad with spaces
 str, err := LeftPad(s, n)

 // pad with specified character
 str, err := func LeftPadStr(s, n, c)

*/
package leftpad

import (
	"errors"
	"fmt"
	"strings"
)

func doLeftPad(s string, n int, r rune) (string, error) {
	if n < 0 {
		return "", errors.New(fmt.Sprintf("Invalid length %d", n))
	}

	toAdd := n - len(s)
	if toAdd <= 0 {
		return s, nil
	}

	return strings.Repeat(string(r), toAdd) + s, nil
}

// LeftPad left-pads s with spaces, to length n.
// If n is smaller than s, LeftPad is a no-op.
func LeftPad(s string, n int) (string, error) {
	return doLeftPad(s, n, ' ')
}

// LeftPadStr left-pads s with the rune r, to length n.
// If n is smaller than s, LeftPadStr is a no-op.
func LeftPadStr(s string, n int, r rune) (string, error) {
	return doLeftPad(s, n, r)
}
<|endoftext|>


<commit_before>// Copyright 2022 Google LLC
//
// 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 lemming provides reference device to be used with ondatra.
package lemming

import (
	"log"
	"net"
	"os"
	"sync"

	fgnmi "github.com/openconfig/lemming/gnmi"
	"github.com/openconfig/lemming/gnmi/gnmit"
	"github.com/openconfig/lemming/gnmi/testagentlocal"
	fgnoi "github.com/openconfig/lemming/gnoi"
	fgnsi "github.com/openconfig/lemming/gnsi"
	fgribi "github.com/openconfig/lemming/gribi"
	fp4rt "github.com/openconfig/lemming/p4rt"
	"github.com/openconfig/lemming/sysrib"
	"google.golang.org/grpc"
	"google.golang.org/grpc/reflection"
	"k8s.io/klog/v2"

	zpb "github.com/openconfig/lemming/proto/sysrib"
)

// Device is the reference device implementation.
type Device struct {
	s           *grpc.Server
	lis         net.Listener
	stop        func()
	gnmiServer  *fgnmi.Server
	gnoiServer  *fgnoi.Server
	gribiServer *fgribi.Server
	gnsiServer  *fgnsi.Server
	p4rtServer  *fp4rt.Server
	// Stores the error if the server fails will be returned on call to stop.
	mu      sync.Mutex
	err     error
	stopped chan struct{}
}

// registerTestTask registers a test gothread that reads from the central
// datastore.
//
// Note: This should only be used for testing lemming, since interface paths
// should be owned by the dataplane module.
func registerTestTask(gnmiServer *gnmit.GNMIServer, targetName string) error {
	return gnmiServer.RegisterTask(testagentlocal.InterfaceTask(targetName))
}

// startSysrib starts the sysrib gRPC service at a unix domain socket. This
// should be started prior to routing services to allow them to connect to
// sysrib during their initialization.
func startSysrib() {
	if err := os.RemoveAll(sysrib.SockAddr); err != nil {
		log.Fatal(err)
	}

	lis, err := net.Listen("unix", sysrib.SockAddr)
	if err != nil {
		log.Fatalf("listen error: %v", err)
	}

	grpcServer := grpc.NewServer()
	s, err := sysrib.NewServer(nil)
	if err != nil {
		log.Fatalf("error while creating sysrib server: %v", err)
	}
	s.AddInterface("eth0", 0, true, "192.0.0.0/8", "DEFAULT")
	zpb.RegisterSysribServer(grpcServer, s)

	go func() {
		grpcServer.Serve(lis)
	}()
}

// New returns a new initialized device.
func New(lis net.Listener, targetName string, opts ...grpc.ServerOption) (*Device, error) {
	startSysrib()

	s := grpc.NewServer(opts...)

	gnmiServer, err := fgnmi.New(s, targetName)
	if err != nil {
		return nil, err
	}
	if err := registerTestTask(gnmiServer.GNMIServer, targetName); err != nil {
		return nil, err
	}

	gribiServer, err := fgribi.New(s)
	if err != nil {
		return nil, err
	}

	d := &Device{
		lis:         lis,
		s:           s,
		gnmiServer:  gnmiServer,
		gnoiServer:  fgnoi.New(s),
		gribiServer: gribiServer,
		gnsiServer:  fgnsi.New(s),
		p4rtServer:  fp4rt.New(s),
	}
	reflection.Register(s)
	d.startServer()
	return d, nil
}

// Addr returns the currently configured ip:port for the listening services.
func (d *Device) Addr() string {
	return d.lis.Addr().String()
}

// Stop stops the listening services.
// If error is not nil, it will contain why the server failed.
func (d *Device) Stop() error {
	klog.Info("Stopping server")
	select {
	case <-d.stopped:
		klog.Info("Server already stopped: ", d.err)
	default:
		d.stop()
	}
	d.mu.Lock()
	defer d.mu.Unlock()
	return d.err
}

// GNMI returns the gNMI server implementation.
func (d *Device) GNMI() *fgnmi.Server {
	return d.gnmiServer
}

// GNSI returns the gNSI server implementation.
func (d *Device) GNSI() *fgnsi.Server {
	return d.gnsiServer
}

func (d *Device) startServer() {
	d.stopped = make(chan struct{})
	go func() {
		err := d.s.Serve(d.lis)
		d.mu.Lock()
		defer d.mu.Unlock()
		d.err = err
		klog.Infof("Server stopped: %v", err)
		close(d.stopped)
	}()
	d.stop = func() {
		d.s.Stop()
		<-d.stopped
	}
}
<commit_msg>Remove stale testing line<commit_after>// Copyright 2022 Google LLC
//
// 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 lemming provides reference device to be used with ondatra.
package lemming

import (
	"log"
	"net"
	"os"
	"sync"

	fgnmi "github.com/openconfig/lemming/gnmi"
	"github.com/openconfig/lemming/gnmi/gnmit"
	"github.com/openconfig/lemming/gnmi/testagentlocal"
	fgnoi "github.com/openconfig/lemming/gnoi"
	fgnsi "github.com/openconfig/lemming/gnsi"
	fgribi "github.com/openconfig/lemming/gribi"
	fp4rt "github.com/openconfig/lemming/p4rt"
	"github.com/openconfig/lemming/sysrib"
	"google.golang.org/grpc"
	"google.golang.org/grpc/reflection"
	"k8s.io/klog/v2"

	zpb "github.com/openconfig/lemming/proto/sysrib"
)

// Device is the reference device implementation.
type Device struct {
	s           *grpc.Server
	lis         net.Listener
	stop        func()
	gnmiServer  *fgnmi.Server
	gnoiServer  *fgnoi.Server
	gribiServer *fgribi.Server
	gnsiServer  *fgnsi.Server
	p4rtServer  *fp4rt.Server
	// Stores the error if the server fails will be returned on call to stop.
	mu      sync.Mutex
	err     error
	stopped chan struct{}
}

// registerTestTask registers a test gothread that reads from the central
// datastore.
//
// Note: This should only be used for testing lemming, since interface paths
// should be owned by the dataplane module.
func registerTestTask(gnmiServer *gnmit.GNMIServer, targetName string) error {
	return gnmiServer.RegisterTask(testagentlocal.InterfaceTask(targetName))
}

// startSysrib starts the sysrib gRPC service at a unix domain socket. This
// should be started prior to routing services to allow them to connect to
// sysrib during their initialization.
func startSysrib() {
	if err := os.RemoveAll(sysrib.SockAddr); err != nil {
		log.Fatal(err)
	}

	lis, err := net.Listen("unix", sysrib.SockAddr)
	if err != nil {
		log.Fatalf("listen error: %v", err)
	}

	grpcServer := grpc.NewServer()
	s, err := sysrib.NewServer(nil)
	if err != nil {
		log.Fatalf("error while creating sysrib server: %v", err)
	}
	zpb.RegisterSysribServer(grpcServer, s)

	go func() {
		grpcServer.Serve(lis)
	}()
}

// New returns a new initialized device.
func New(lis net.Listener, targetName string, opts ...grpc.ServerOption) (*Device, error) {
	startSysrib()

	s := grpc.NewServer(opts...)

	gnmiServer, err := fgnmi.New(s, targetName)
	if err != nil {
		return nil, err
	}
	if err := registerTestTask(gnmiServer.GNMIServer, targetName); err != nil {
		return nil, err
	}

	gribiServer, err := fgribi.New(s)
	if err != nil {
		return nil, err
	}

	d := &Device{
		lis:         lis,
		s:           s,
		gnmiServer:  gnmiServer,
		gnoiServer:  fgnoi.New(s),
		gribiServer: gribiServer,
		gnsiServer:  fgnsi.New(s),
		p4rtServer:  fp4rt.New(s),
	}
	reflection.Register(s)
	d.startServer()
	return d, nil
}

// Addr returns the currently configured ip:port for the listening services.
func (d *Device) Addr() string {
	return d.lis.Addr().String()
}

// Stop stops the listening services.
// If error is not nil, it will contain why the server failed.
func (d *Device) Stop() error {
	klog.Info("Stopping server")
	select {
	case <-d.stopped:
		klog.Info("Server already stopped: ", d.err)
	default:
		d.stop()
	}
	d.mu.Lock()
	defer d.mu.Unlock()
	return d.err
}

// GNMI returns the gNMI server implementation.
func (d *Device) GNMI() *fgnmi.Server {
	return d.gnmiServer
}

// GNSI returns the gNSI server implementation.
func (d *Device) GNSI() *fgnsi.Server {
	return d.gnsiServer
}

func (d *Device) startServer() {
	d.stopped = make(chan struct{})
	go func() {
		err := d.s.Serve(d.lis)
		d.mu.Lock()
		defer d.mu.Unlock()
		d.err = err
		klog.Infof("Server stopped: %v", err)
		close(d.stopped)
	}()
	d.stop = func() {
		d.s.Stop()
		<-d.stopped
	}
}
<|endoftext|>


<commit_before>// Copyright (c) 2015-2021 MinIO, Inc.
//
// This file is part of MinIO Object Storage stack
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

package cmd

import (
	"github.com/minio/cli"
	"github.com/minio/mc/pkg/probe"
)

var adminUserSvcAcctRemoveCmd = cli.Command{
	Name:         "rm",
	Aliases:      []string{"remove"},
	Usage:        "remove a service account",
	Action:       mainAdminUserSvcAcctRemove,
	OnUsageError: onUsageError,
	Before:       setGlobalsFromContext,
	Flags:        globalFlags,
	CustomHelpTemplate: `NAME:
  {{.HelpName}} - {{.Usage}}

USAGE:
  {{.HelpName}} ALIAS SERVICE-ACCOUNT

FLAGS:
  {{range .VisibleFlags}}{{.}}
  {{end}}
EXAMPLES:
  1. Remove a service account 'J123C4ZXEQN8RK6ND35I' from MinIO server.
     {{.Prompt}} {{.HelpName}} myminio/ J123C4ZXEQN8RK6ND35I
`,
}

// checkAdminUserSvcAcctRemoveSyntax - validate all the passed arguments
func checkAdminUserSvcAcctRemoveSyntax(ctx *cli.Context) {
	if len(ctx.Args()) != 2 {
		cli.ShowCommandHelpAndExit(ctx, "rm", 1)
	}
}

// mainAdminUserSvcAcctRemove is the handle for "mc admin user svcacct rm" command.
func mainAdminUserSvcAcctRemove(ctx *cli.Context) error {
	checkAdminUserSvcAcctRemoveSyntax(ctx)

	// Get the alias parameter from cli
	args := ctx.Args()
	aliasedURL := args.Get(0)
	svcAccount := args.Get(1)

	// Create a new MinIO Admin Client
	client, err := newAdminClient(aliasedURL)
	fatalIf(err, "Unable to initialize admin connection.")

	e := client.DeleteServiceAccount(globalContext, svcAccount)
	fatalIf(probe.NewError(e).Trace(args...), "Unable to remove a new service account")

	printMsg(svcAcctMessage{
		op:        "ls",
		AccessKey: svcAccount,
	})

	return nil
}
<commit_msg>svcacct: Show successful message after deletion (#4016)<commit_after>// Copyright (c) 2015-2021 MinIO, Inc.
//
// This file is part of MinIO Object Storage stack
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

package cmd

import (
	"github.com/fatih/color"
	"github.com/minio/cli"
	"github.com/minio/mc/pkg/probe"
	"github.com/minio/pkg/console"
)

var adminUserSvcAcctRemoveCmd = cli.Command{
	Name:         "rm",
	Aliases:      []string{"remove"},
	Usage:        "remove a service account",
	Action:       mainAdminUserSvcAcctRemove,
	OnUsageError: onUsageError,
	Before:       setGlobalsFromContext,
	Flags:        globalFlags,
	CustomHelpTemplate: `NAME:
  {{.HelpName}} - {{.Usage}}

USAGE:
  {{.HelpName}} ALIAS SERVICE-ACCOUNT

FLAGS:
  {{range .VisibleFlags}}{{.}}
  {{end}}
EXAMPLES:
  1. Remove a service account 'J123C4ZXEQN8RK6ND35I' from MinIO server.
     {{.Prompt}} {{.HelpName}} myminio/ J123C4ZXEQN8RK6ND35I
`,
}

// checkAdminUserSvcAcctRemoveSyntax - validate all the passed arguments
func checkAdminUserSvcAcctRemoveSyntax(ctx *cli.Context) {
	if len(ctx.Args()) != 2 {
		cli.ShowCommandHelpAndExit(ctx, "rm", 1)
	}
}

// mainAdminUserSvcAcctRemove is the handle for "mc admin user svcacct rm" command.
func mainAdminUserSvcAcctRemove(ctx *cli.Context) error {
	console.SetColor("SVCMessage", color.New(color.FgGreen))

	checkAdminUserSvcAcctRemoveSyntax(ctx)

	// Get the alias parameter from cli
	args := ctx.Args()
	aliasedURL := args.Get(0)
	svcAccount := args.Get(1)

	// Create a new MinIO Admin Client
	client, err := newAdminClient(aliasedURL)
	fatalIf(err, "Unable to initialize admin connection.")

	e := client.DeleteServiceAccount(globalContext, svcAccount)
	fatalIf(probe.NewError(e).Trace(args...), "Unable to remove a new service account")

	printMsg(svcAcctMessage{
		op:        "rm",
		AccessKey: svcAccount,
	})

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/workfit/tester/assert"
	"testing"
)

func TestEnumParent(t *testing.T) {
	tests := []struct {
		description string
		strValues   map[string]string
		//If nil, expect no change from strValues
		expectedValues  map[string]string
		expectedParents map[string]string
	}{
		{
			"Single layer",
			map[string]string{
				"Phase":         "",
				"PhaseAnother":  "Another",
				"PhaseOverride": "Heyo",
			},
			nil,
			map[string]string{
				"Phase":         "Phase",
				"PhaseAnother":  "Phase",
				"PhaseOverride": "Phase",
			},
		},
		{
			"No Tree",
			map[string]string{
				"ColorBlue":  "Blue",
				"ColorGreen": "Green",
				"ColorRed":   "Red",
			},
			nil,
			nil,
		},
		{
			"Two layers",
			map[string]string{
				"Color":          "",
				"ColorBlue":      "Blue",
				"ColorBlue_One":  "Blue > One",
				"ColorBlue_Two":  "Blue > Two",
				"ColorGreen":     "Green",
				"ColorGreen_One": "Green > One",
			},
			map[string]string{
				"Color":          "",
				"ColorBlue":      "Blue",
				"ColorBlue_One":  "One",
				"ColorBlue_Two":  "Two",
				"ColorGreen":     "Green",
				"ColorGreen_One": "One",
			},
			map[string]string{
				"Color":          "Color",
				"ColorBlue":      "Color",
				"ColorBlue_One":  "ColorBlue",
				"ColorBlue_Two":  "ColorBlue",
				"ColorGreen":     "Color",
				"ColorGreen_One": "ColorGreen",
			},
		},
		{
			"Three layers",
			map[string]string{
				"Color":             "",
				"ColorBlue":         "Blue",
				"ColorBlue_One":     "Blue > One",
				"ColorBlue_Two":     "Blue > Two",
				"ColorBlue_One_One": "Blue > One > One",
				"ColorBlue_One_Two": "Blue > One > Two",
			},
			map[string]string{
				"Color":             "",
				"ColorBlue":         "Blue",
				"ColorBlue_One":     "One",
				"ColorBlue_Two":     "Two",
				"ColorBlue_One_One": "One",
				"ColorBlue_One_Two": "Two",
			},
			map[string]string{
				"Color":             "Color",
				"ColorBlue":         "Color",
				"ColorBlue_One":     "ColorBlue",
				"ColorBlue_Two":     "ColorBlue",
				"ColorBlue_One_One": "ColorBlue_One",
				"ColorBlue_One_Two": "ColorBlue_One",
			},
		},
		{
			"Single implied layer",
			map[string]string{
				"Color":         "",
				"ColorBlue_One": "Blue > One",
			},
			map[string]string{
				"Color":                "",
				"-9223372036854775808": "Blue",
				"ColorBlue_One":        "One",
			},
			map[string]string{
				"Color":                "Color",
				"-9223372036854775808": "Color",
				"ColorBlue_One":        "-9223372036854775808",
			},
		},
		{
			"Two implied layers",
			map[string]string{
				"Color":            "",
				"ColorGreen_One_A": "Green > One > A",
			},
			map[string]string{
				"Color":                "",
				"-9223372036854775808": "Green",
				"-9223372036854775807": "One",
				"ColorGreen_One_A":     "A",
			},
			map[string]string{
				"Color":                "Color",
				"-9223372036854775808": "Color",
				"-9223372036854775807": "-9223372036854775808",
				"ColorGreen_One_A":     "-9223372036854775807",
			},
		},
		{
			"Single word implied nesting",
			map[string]string{
				"Color":        "",
				"ColorBlue":    "Blue",
				"ColorBlueOne": "Blue One",
				"ColorBlueTwo": "Blue Two",
			},
			map[string]string{
				"Color":        "",
				"ColorBlue":    "Blue",
				"ColorBlueOne": "One",
				"ColorBlueTwo": "Two",
			},
			map[string]string{
				"Color":        "Color",
				"ColorBlue":    "Color",
				"ColorBlueOne": "ColorBlue",
				"ColorBlueTwo": "ColorBlue",
			},
		},
		{
			"Multi-Word implied nesting",
			map[string]string{
				"Color":             "",
				"ColorBlueGreen":    "Blue Green",
				"ColorBlueGreenOne": "Blue Green One",
				"ColorBlueGreenTwo": "Blue Green Two",
			},
			map[string]string{
				"Color":             "",
				"ColorBlueGreen":    "Blue Green",
				"ColorBlueGreenOne": "One",
				"ColorBlueGreenTwo": "Two",
			},
			map[string]string{
				"Color":             "Color",
				"ColorBlueGreen":    "Color",
				"ColorBlueGreenOne": "ColorBlueGreen",
				"ColorBlueGreenTwo": "ColorBlueGreen",
			},
		},
		{
			"Implied node with implied nesting",
			map[string]string{
				"Color":             "",
				"ColorBlueGreenOne": "Blue Green One",
				"ColorBlueGreenTwo": "Blue Green Two",
			},
			map[string]string{
				"Color":                "",
				"-9223372036854775808": "Blue Green",
				"ColorBlueGreenOne":    "One",
				"ColorBlueGreenTwo":    "Two",
			},
			map[string]string{
				"Color":                "Color",
				"-9223372036854775808": "Color",
				"ColorBlueGreenOne":    "-9223372036854775808",
				"ColorBlueGreenTwo":    "-9223372036854775808",
			},
		},
		{
			"Multiple implied layers",
			map[string]string{
				"Color":              "",
				"ColorBlueGreen":     "Blue Green",
				"ColorBlueGreenOne":  "Blue Green One",
				"ColorBlueGreenOneA": "Blue Green One A",
				"ColorBlueGreenOneB": "Blue Green One B",
				"ColorBlueGreenTwo":  "Blue Green Two",
			},
			map[string]string{
				"Color":              "",
				"ColorBlueGreen":     "Blue Green",
				"ColorBlueGreenOne":  "One",
				"ColorBlueGreenOneA": "A",
				"ColorBlueGreenOneB": "B",
				"ColorBlueGreenTwo":  "Two",
			},
			map[string]string{
				"Color":              "Color",
				"ColorBlueGreen":     "Color",
				"ColorBlueGreenOne":  "ColorBlueGreen",
				"ColorBlueGreenOneA": "ColorBlueGreenOne",
				"ColorBlueGreenOneB": "ColorBlueGreenOne",
				"ColorBlueGreenTwo":  "ColorBlueGreen",
			},
		},

		{
			"Multiple implied layers with implied node",
			map[string]string{
				"Color":              "",
				"ColorBlueGreen":     "Blue Green",
				"ColorBlueGreenOneA": "Blue Green One A",
				"ColorBlueGreenOneB": "Blue Green One B",
				"ColorBlueGreenTwo":  "Blue Green Two",
			},
			map[string]string{
				"Color":                "",
				"ColorBlueGreen":       "Blue Green",
				"-9223372036854775808": "One",
				"ColorBlueGreenOneA":   "A",
				"ColorBlueGreenOneB":   "B",
				"ColorBlueGreenTwo":    "Two",
			},
			map[string]string{
				"Color":                "Color",
				"ColorBlueGreen":       "Color",
				"-9223372036854775808": "ColorBlueGreen",
				"ColorBlueGreenOneA":   "-9223372036854775808",
				"ColorBlueGreenOneB":   "-9223372036854775808",
				"ColorBlueGreenTwo":    "ColorBlueGreen",
			},
		},
		{
			"Mix implicit and explicit layers",
			map[string]string{
				"Color":               "",
				"ColorBlueGreen":      "Blue Green",
				"ColorBlueGreenOne":   "Blue Green One",
				"ColorBlueGreenOne_A": "Blue Green One > A",
				"ColorBlueGreenOne_B": "Blue Green One > B",
				"ColorBlueGreenTwo":   "Blue Green Two",
			},
			map[string]string{
				"Color":               "",
				"ColorBlueGreen":      "Blue Green",
				"ColorBlueGreenOne":   "One",
				"ColorBlueGreenOne_A": "A",
				"ColorBlueGreenOne_B": "B",
				"ColorBlueGreenTwo":   "Two",
			},
			map[string]string{
				"Color":               "Color",
				"ColorBlueGreen":      "Color",
				"ColorBlueGreenOne":   "ColorBlueGreen",
				"ColorBlueGreenOne_A": "ColorBlueGreenOne",
				"ColorBlueGreenOne_B": "ColorBlueGreenOne",
				"ColorBlueGreenTwo":   "ColorBlueGreen",
			},
		},
	}

	for i, test := range tests {

		e := newEnum("test", transformNone)

		for key, val := range test.strValues {
			e.AddTransformKey(key, true, val, transformNone)
		}

		err := e.Process()
		assert.For(t, i).ThatActual(err).IsNil()

		actualValues := e.ValueMap()
		actualParents := e.Parents()
		if test.expectedValues == nil {
			//Expect no change from test.strValues
			assert.For(t, i).ThatActual(actualValues).Equals(test.strValues).ThenDiffOnFail()
		} else {
			assert.For(t, i).ThatActual(actualValues).Equals(test.expectedValues).ThenDiffOnFail()
		}
		assert.For(t, i).ThatActual(actualParents).Equals(test.expectedParents).ThenDiffOnFail()
	}
}
<commit_msg>Add a test of two implied nodes (one of which is a multi-word node) in a row. Part of #628.<commit_after>package main

import (
	"github.com/workfit/tester/assert"
	"testing"
)

func TestEnumParent(t *testing.T) {
	tests := []struct {
		description string
		strValues   map[string]string
		//If nil, expect no change from strValues
		expectedValues  map[string]string
		expectedParents map[string]string
	}{
		{
			"Single layer",
			map[string]string{
				"Phase":         "",
				"PhaseAnother":  "Another",
				"PhaseOverride": "Heyo",
			},
			nil,
			map[string]string{
				"Phase":         "Phase",
				"PhaseAnother":  "Phase",
				"PhaseOverride": "Phase",
			},
		},
		{
			"No Tree",
			map[string]string{
				"ColorBlue":  "Blue",
				"ColorGreen": "Green",
				"ColorRed":   "Red",
			},
			nil,
			nil,
		},
		{
			"Two layers",
			map[string]string{
				"Color":          "",
				"ColorBlue":      "Blue",
				"ColorBlue_One":  "Blue > One",
				"ColorBlue_Two":  "Blue > Two",
				"ColorGreen":     "Green",
				"ColorGreen_One": "Green > One",
			},
			map[string]string{
				"Color":          "",
				"ColorBlue":      "Blue",
				"ColorBlue_One":  "One",
				"ColorBlue_Two":  "Two",
				"ColorGreen":     "Green",
				"ColorGreen_One": "One",
			},
			map[string]string{
				"Color":          "Color",
				"ColorBlue":      "Color",
				"ColorBlue_One":  "ColorBlue",
				"ColorBlue_Two":  "ColorBlue",
				"ColorGreen":     "Color",
				"ColorGreen_One": "ColorGreen",
			},
		},
		{
			"Three layers",
			map[string]string{
				"Color":             "",
				"ColorBlue":         "Blue",
				"ColorBlue_One":     "Blue > One",
				"ColorBlue_Two":     "Blue > Two",
				"ColorBlue_One_One": "Blue > One > One",
				"ColorBlue_One_Two": "Blue > One > Two",
			},
			map[string]string{
				"Color":             "",
				"ColorBlue":         "Blue",
				"ColorBlue_One":     "One",
				"ColorBlue_Two":     "Two",
				"ColorBlue_One_One": "One",
				"ColorBlue_One_Two": "Two",
			},
			map[string]string{
				"Color":             "Color",
				"ColorBlue":         "Color",
				"ColorBlue_One":     "ColorBlue",
				"ColorBlue_Two":     "ColorBlue",
				"ColorBlue_One_One": "ColorBlue_One",
				"ColorBlue_One_Two": "ColorBlue_One",
			},
		},
		{
			"Single implied layer",
			map[string]string{
				"Color":         "",
				"ColorBlue_One": "Blue > One",
			},
			map[string]string{
				"Color":                "",
				"-9223372036854775808": "Blue",
				"ColorBlue_One":        "One",
			},
			map[string]string{
				"Color":                "Color",
				"-9223372036854775808": "Color",
				"ColorBlue_One":        "-9223372036854775808",
			},
		},
		{
			"Two implied layers",
			map[string]string{
				"Color":            "",
				"ColorGreen_One_A": "Green > One > A",
			},
			map[string]string{
				"Color":                "",
				"-9223372036854775808": "Green",
				"-9223372036854775807": "One",
				"ColorGreen_One_A":     "A",
			},
			map[string]string{
				"Color":                "Color",
				"-9223372036854775808": "Color",
				"-9223372036854775807": "-9223372036854775808",
				"ColorGreen_One_A":     "-9223372036854775807",
			},
		},
		{
			"Single word implied nesting",
			map[string]string{
				"Color":        "",
				"ColorBlue":    "Blue",
				"ColorBlueOne": "Blue One",
				"ColorBlueTwo": "Blue Two",
			},
			map[string]string{
				"Color":        "",
				"ColorBlue":    "Blue",
				"ColorBlueOne": "One",
				"ColorBlueTwo": "Two",
			},
			map[string]string{
				"Color":        "Color",
				"ColorBlue":    "Color",
				"ColorBlueOne": "ColorBlue",
				"ColorBlueTwo": "ColorBlue",
			},
		},
		{
			"Multi-Word implied nesting",
			map[string]string{
				"Color":             "",
				"ColorBlueGreen":    "Blue Green",
				"ColorBlueGreenOne": "Blue Green One",
				"ColorBlueGreenTwo": "Blue Green Two",
			},
			map[string]string{
				"Color":             "",
				"ColorBlueGreen":    "Blue Green",
				"ColorBlueGreenOne": "One",
				"ColorBlueGreenTwo": "Two",
			},
			map[string]string{
				"Color":             "Color",
				"ColorBlueGreen":    "Color",
				"ColorBlueGreenOne": "ColorBlueGreen",
				"ColorBlueGreenTwo": "ColorBlueGreen",
			},
		},
		{
			"Implied node with implied nesting",
			map[string]string{
				"Color":             "",
				"ColorBlueGreenOne": "Blue Green One",
				"ColorBlueGreenTwo": "Blue Green Two",
			},
			map[string]string{
				"Color":                "",
				"-9223372036854775808": "Blue Green",
				"ColorBlueGreenOne":    "One",
				"ColorBlueGreenTwo":    "Two",
			},
			map[string]string{
				"Color":                "Color",
				"-9223372036854775808": "Color",
				"ColorBlueGreenOne":    "-9223372036854775808",
				"ColorBlueGreenTwo":    "-9223372036854775808",
			},
		},
		{
			"Multiple implied layers",
			map[string]string{
				"Color":              "",
				"ColorBlueGreen":     "Blue Green",
				"ColorBlueGreenOne":  "Blue Green One",
				"ColorBlueGreenOneA": "Blue Green One A",
				"ColorBlueGreenOneB": "Blue Green One B",
				"ColorBlueGreenTwo":  "Blue Green Two",
			},
			map[string]string{
				"Color":              "",
				"ColorBlueGreen":     "Blue Green",
				"ColorBlueGreenOne":  "One",
				"ColorBlueGreenOneA": "A",
				"ColorBlueGreenOneB": "B",
				"ColorBlueGreenTwo":  "Two",
			},
			map[string]string{
				"Color":              "Color",
				"ColorBlueGreen":     "Color",
				"ColorBlueGreenOne":  "ColorBlueGreen",
				"ColorBlueGreenOneA": "ColorBlueGreenOne",
				"ColorBlueGreenOneB": "ColorBlueGreenOne",
				"ColorBlueGreenTwo":  "ColorBlueGreen",
			},
		},

		{
			"Multiple implied layers with implied node",
			map[string]string{
				"Color":              "",
				"ColorBlueGreen":     "Blue Green",
				"ColorBlueGreenOneA": "Blue Green One A",
				"ColorBlueGreenOneB": "Blue Green One B",
				"ColorBlueGreenTwo":  "Blue Green Two",
			},
			map[string]string{
				"Color":                "",
				"ColorBlueGreen":       "Blue Green",
				"-9223372036854775808": "One",
				"ColorBlueGreenOneA":   "A",
				"ColorBlueGreenOneB":   "B",
				"ColorBlueGreenTwo":    "Two",
			},
			map[string]string{
				"Color":                "Color",
				"ColorBlueGreen":       "Color",
				"-9223372036854775808": "ColorBlueGreen",
				"ColorBlueGreenOneA":   "-9223372036854775808",
				"ColorBlueGreenOneB":   "-9223372036854775808",
				"ColorBlueGreenTwo":    "ColorBlueGreen",
			},
		},
		{
			"Mix implicit and explicit layers",
			map[string]string{
				"Color":               "",
				"ColorBlueGreen":      "Blue Green",
				"ColorBlueGreenOne":   "Blue Green One",
				"ColorBlueGreenOne_A": "Blue Green One > A",
				"ColorBlueGreenOne_B": "Blue Green One > B",
				"ColorBlueGreenTwo":   "Blue Green Two",
			},
			map[string]string{
				"Color":               "",
				"ColorBlueGreen":      "Blue Green",
				"ColorBlueGreenOne":   "One",
				"ColorBlueGreenOne_A": "A",
				"ColorBlueGreenOne_B": "B",
				"ColorBlueGreenTwo":   "Two",
			},
			map[string]string{
				"Color":               "Color",
				"ColorBlueGreen":      "Color",
				"ColorBlueGreenOne":   "ColorBlueGreen",
				"ColorBlueGreenOne_A": "ColorBlueGreenOne",
				"ColorBlueGreenOne_B": "ColorBlueGreenOne",
				"ColorBlueGreenTwo":   "ColorBlueGreen",
			},
		},
		{
			"Multiple implied layers in a row",
			map[string]string{
				"Color":              "",
				"ColorBlueGreenOneA": "Blue Green One A",
				"ColorBlueGreenOneB": "Blue Green One B",
				"ColorBlueGreenTwo":  "Blue Green Two",
			},
			map[string]string{
				"Color":                "",
				"-9223372036854775808": "Blue Green",
				"-9223372036854775807": "One",
				"ColorBlueGreenOneA":   "A",
				"ColorBlueGreenOneB":   "B",
				"ColorBlueGreenTwo":    "Two",
			},
			map[string]string{
				"Color":                "Color",
				"-9223372036854775808": "Color",
				"-9223372036854775807": "-9223372036854775808",
				"ColorBlueGreenOneA":   "-9223372036854775807",
				"ColorBlueGreenOneB":   "-9223372036854775807",
				"ColorBlueGreenTwo":    "-9223372036854775808",
			},
		},
	}

	for i, test := range tests {

		e := newEnum("test", transformNone)

		for key, val := range test.strValues {
			e.AddTransformKey(key, true, val, transformNone)
		}

		err := e.Process()
		assert.For(t, i).ThatActual(err).IsNil()

		actualValues := e.ValueMap()
		actualParents := e.Parents()
		if test.expectedValues == nil {
			//Expect no change from test.strValues
			assert.For(t, i).ThatActual(actualValues).Equals(test.strValues).ThenDiffOnFail()
		} else {
			assert.For(t, i).ThatActual(actualValues).Equals(test.expectedValues).ThenDiffOnFail()
		}
		assert.For(t, i).ThatActual(actualParents).Equals(test.expectedParents).ThenDiffOnFail()
	}
}
<|endoftext|>


<commit_before>package main

// importer.go implements a data fetching service capable of pulling objects from remote object
// stores and writing to a local directory. It utilizes the minio-go client sdk for s3 remotes,
// https for public remotes, and "file" for local files. The main use-case for this importer is
// to copy VM images to a "golden" namespace for consumption by kubevirt.
// This process expects several environmental variables:
//    ImporterEndpoint       Endpoint url minus scheme, bucket/object and port, eg. s3.amazon.com.
//			      Access and secret keys are optional. If omitted no creds are passed
//			      to the object store client.
//    ImporterAccessKeyID  Optional. Access key is the user ID that uniquely identifies your
//			      account.
//    ImporterSecretKey     Optional. Secret key is the password to your account.

import (
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"strconv"
	"time"

	"github.com/pkg/errors"

	v1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/resource"
	"k8s.io/klog/v2"

	cdiv1 "kubevirt.io/containerized-data-importer-api/pkg/apis/core/v1beta1"
	"kubevirt.io/containerized-data-importer/pkg/common"
	"kubevirt.io/containerized-data-importer/pkg/controller"
	"kubevirt.io/containerized-data-importer/pkg/image"
	"kubevirt.io/containerized-data-importer/pkg/importer"
	"kubevirt.io/containerized-data-importer/pkg/util"
	prometheusutil "kubevirt.io/containerized-data-importer/pkg/util/prometheus"
)

func init() {
	klog.InitFlags(nil)
	flag.Parse()
}

func waitForReadyFile() {
	readyFile, _ := util.ParseEnvVar(common.ImporterReadyFile, false)
	if readyFile == "" {
		return
	}
	for {
		if _, err := os.Stat(readyFile); err == nil {
			break
		}
		time.Sleep(time.Second)
	}
}

func touchDoneFile() {
	doneFile, _ := util.ParseEnvVar(common.ImporterDoneFile, false)
	if doneFile == "" {
		return
	}
	f, err := os.OpenFile(doneFile, os.O_CREATE|os.O_EXCL, 0666)
	if err != nil {
		klog.Errorf("Failed creating file %s: %+v", doneFile, err)
	}
	f.Close()
}

func main() {
	defer klog.Flush()

	certsDirectory, err := ioutil.TempDir("", "certsdir")
	if err != nil {
		panic(err)
	}
	defer os.RemoveAll(certsDirectory)
	prometheusutil.StartPrometheusEndpoint(certsDirectory)
	klog.V(1).Infoln("Starting importer")

	source, _ := util.ParseEnvVar(common.ImporterSource, false)
	contentType, _ := util.ParseEnvVar(common.ImporterContentType, false)
	imageSize, _ := util.ParseEnvVar(common.ImporterImageSize, false)
	filesystemOverhead, _ := strconv.ParseFloat(os.Getenv(common.FilesystemOverheadVar), 64)
	preallocation, err := strconv.ParseBool(os.Getenv(common.Preallocation))
	var preallocationApplied bool
	var ds importer.DataSourceInterface

	volumeMode := v1.PersistentVolumeBlock
	if _, err := os.Stat(common.WriteBlockPath); os.IsNotExist(err) {
		volumeMode = v1.PersistentVolumeFilesystem
	} else {
		preallocation = true
	}

	// With writeback cache mode it's possible that the process will exit before all writes have been commited to storage.
	// To guarantee that our write was commited to storage, we make a fsync syscall and ensure success.
	// Also might be a good idea to sync any chmod's we might have done.
	defer func() {
		dataFile := getImporterDestPath(contentType, volumeMode)
		file, err := os.Open(dataFile)
		if err != nil {
			klog.Errorf("could not get file descriptor for fsync call: %+v", err)
			os.Exit(1)
		}
		if err := file.Sync(); err != nil {
			klog.Errorf("could not fsync following qemu-img writing: %+v", err)
			os.Exit(1)
		}
		klog.V(3).Infof("Successfully completed fsync(%s) syscall, commited to disk\n", dataFile)
		file.Close()
	}()

	//Registry import currently support kubevirt content type only
	if contentType != string(cdiv1.DataVolumeKubeVirt) && (source == controller.SourceRegistry || source == controller.SourceImageio) {
		klog.Errorf("Unsupported content type %s when importing from %s", contentType, source)
		os.Exit(1)
	}

	availableDestSpace, err := util.GetAvailableSpaceByVolumeMode(volumeMode)
	if err != nil {
		klog.Errorf("%+v", err)
		os.Exit(1)
	}
	if source == controller.SourceNone {
		if contentType == string(cdiv1.DataVolumeKubeVirt) {
			createBlankImage(imageSize, availableDestSpace, preallocation, volumeMode, filesystemOverhead)
			preallocationApplied = preallocation
		} else {
			errorEmptyDiskWithContentTypeArchive()
		}
	} else {
		waitForReadyFile()
		klog.V(1).Infoln("begin import process")

		ds = newDataSource(source, contentType, volumeMode)
		defer ds.Close()

		processor := newDataProcessor(contentType, volumeMode, ds, imageSize, filesystemOverhead, preallocation)
		err = processor.ProcessData()

		if err != nil {
			klog.Errorf("%+v", err)
			if err == importer.ErrRequiresScratchSpace {
				ds.Close()
				os.Exit(common.ScratchSpaceNeededExitCode)
			}
			err = util.WriteTerminationMessage(fmt.Sprintf("Unable to process data: %+v", err.Error()))
			if err != nil {
				klog.Errorf("%+v", err)
			}
			ds.Close()
			os.Exit(1)
		}
		touchDoneFile()
		preallocationApplied = processor.PreallocationApplied()
	}
	message := "Import Complete"
	if preallocationApplied {
		message += ", " + common.PreallocationApplied
	}
	err = util.WriteTerminationMessage(message)
	if err != nil {
		klog.Errorf("%+v", err)
		if ds != nil {
			ds.Close()
		}
		os.Exit(1)
	}
	klog.V(1).Infoln(message)
}

func newDataProcessor(contentType string, volumeMode v1.PersistentVolumeMode, ds importer.DataSourceInterface, imageSize string, filesystemOverhead float64, preallocation bool) *importer.DataProcessor {
	dest := getImporterDestPath(contentType, volumeMode)
	processor := importer.NewDataProcessor(ds, dest, common.ImporterDataDir, common.ScratchDataDir, imageSize, filesystemOverhead, preallocation)
	return processor
}

func getImporterDestPath(contentType string, volumeMode v1.PersistentVolumeMode) string {
	dest := common.ImporterWritePath

	if contentType == string(cdiv1.DataVolumeArchive) {
		dest = common.ImporterVolumePath
	}
	if volumeMode == v1.PersistentVolumeBlock {
		dest = common.WriteBlockPath
	}

	return dest
}

func newDataSource(source string, contentType string, volumeMode v1.PersistentVolumeMode) importer.DataSourceInterface {
	ep, _ := util.ParseEnvVar(common.ImporterEndpoint, false)
	acc, _ := util.ParseEnvVar(common.ImporterAccessKeyID, false)
	sec, _ := util.ParseEnvVar(common.ImporterSecretKey, false)
	diskID, _ := util.ParseEnvVar(common.ImporterDiskID, false)
	uuid, _ := util.ParseEnvVar(common.ImporterUUID, false)
	backingFile, _ := util.ParseEnvVar(common.ImporterBackingFile, false)
	certDir, _ := util.ParseEnvVar(common.ImporterCertDirVar, false)
	insecureTLS, _ := strconv.ParseBool(os.Getenv(common.InsecureTLSVar))
	thumbprint, _ := util.ParseEnvVar(common.ImporterThumbprint, false)

	currentCheckpoint, _ := util.ParseEnvVar(common.ImporterCurrentCheckpoint, false)
	previousCheckpoint, _ := util.ParseEnvVar(common.ImporterPreviousCheckpoint, false)
	finalCheckpoint, _ := util.ParseEnvVar(common.ImporterFinalCheckpoint, false)

	switch source {
	case controller.SourceHTTP:
		ds, err := importer.NewHTTPDataSource(ep, acc, sec, certDir, cdiv1.DataVolumeContentType(contentType))
		if err != nil {
			errorCannotConnectDataSource(err, "http")
		}
		return ds
	case controller.SourceImageio:
		ds, err := importer.NewImageioDataSource(ep, acc, sec, certDir, diskID, currentCheckpoint, previousCheckpoint)
		if err != nil {
			errorCannotConnectDataSource(err, "imageio")
		}
		return ds
	case controller.SourceRegistry:
		ds := importer.NewRegistryDataSource(ep, acc, sec, certDir, insecureTLS)
		return ds
	case controller.SourceS3:
		ds, err := importer.NewS3DataSource(ep, acc, sec, certDir)
		if err != nil {
			errorCannotConnectDataSource(err, "s3")
		}
		return ds
	case controller.SourceVDDK:
		ds, err := importer.NewVDDKDataSource(ep, acc, sec, thumbprint, uuid, backingFile, currentCheckpoint, previousCheckpoint, finalCheckpoint, volumeMode)
		if err != nil {
			errorCannotConnectDataSource(err, "vddk")
		}
		return ds
	default:
		klog.Errorf("Unknown source type %s\n", source)
		err := util.WriteTerminationMessage(fmt.Sprintf("Unknown data source: %s", source))
		if err != nil {
			klog.Errorf("%+v", err)
		}
		os.Exit(1)
	}

	return nil
}

func createBlankImage(imageSize string, availableDestSpace int64, preallocation bool, volumeMode v1.PersistentVolumeMode, filesystemOverhead float64) {
	requestImageSizeQuantity := resource.MustParse(imageSize)
	minSizeQuantity := util.MinQuantity(resource.NewScaledQuantity(availableDestSpace, 0), &requestImageSizeQuantity)

	if minSizeQuantity.Cmp(requestImageSizeQuantity) != 0 {
		// Available dest space is smaller than the size we want to create
		klog.Warningf("Available space less than requested size, creating blank image sized to available space: %s.\n", minSizeQuantity.String())
	}

	var err error
	if volumeMode == v1.PersistentVolumeFilesystem {
		quantityWithFSOverhead := util.GetUsableSpace(filesystemOverhead, minSizeQuantity.Value())
		klog.Infof("Space adjusted for filesystem overhead: %d.\n", quantityWithFSOverhead)
		err = image.CreateBlankImage(common.ImporterWritePath, *resource.NewScaledQuantity(quantityWithFSOverhead, 0), preallocation)
	} else if volumeMode == v1.PersistentVolumeBlock && preallocation {
		klog.V(1).Info("Preallocating blank block volume")
		err = image.PreallocateBlankBlock(common.WriteBlockPath, minSizeQuantity)
	}

	if err != nil {
		klog.Errorf("%+v", err)
		message := fmt.Sprintf("Unable to create blank image: %+v", err)
		err = util.WriteTerminationMessage(message)
		if err != nil {
			klog.Errorf("%+v", err)
		}
		os.Exit(1)
	}
}

func errorCannotConnectDataSource(err error, dsName string) {
	klog.Errorf("%+v", err)
	err = util.WriteTerminationMessage(fmt.Sprintf("Unable to connect to %s data source: %+v", dsName, err))
	if err != nil {
		klog.Errorf("%+v", err)
	}
	os.Exit(1)
}

func errorEmptyDiskWithContentTypeArchive() {
	klog.Errorf("%+v", errors.New("Cannot create empty disk with content type archive"))
	err := util.WriteTerminationMessage("Cannot create empty disk with content type archive")
	if err != nil {
		klog.Errorf("%+v", err)
	}
	os.Exit(1)
}
<commit_msg>Simplify importer code, cleanup and extract function (#2141)<commit_after>package main

// importer.go implements a data fetching service capable of pulling objects from remote object
// stores and writing to a local directory. It utilizes the minio-go client sdk for s3 remotes,
// https for public remotes, and "file" for local files. The main use-case for this importer is
// to copy VM images to a "golden" namespace for consumption by kubevirt.
// This process expects several environmental variables:
//    ImporterEndpoint       Endpoint url minus scheme, bucket/object and port, eg. s3.amazon.com.
//			      Access and secret keys are optional. If omitted no creds are passed
//			      to the object store client.
//    ImporterAccessKeyID  Optional. Access key is the user ID that uniquely identifies your
//			      account.
//    ImporterSecretKey     Optional. Secret key is the password to your account.

import (
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"strconv"
	"time"

	"github.com/pkg/errors"

	v1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/resource"
	"k8s.io/klog/v2"

	cdiv1 "kubevirt.io/containerized-data-importer-api/pkg/apis/core/v1beta1"
	"kubevirt.io/containerized-data-importer/pkg/common"
	"kubevirt.io/containerized-data-importer/pkg/controller"
	"kubevirt.io/containerized-data-importer/pkg/image"
	"kubevirt.io/containerized-data-importer/pkg/importer"
	"kubevirt.io/containerized-data-importer/pkg/util"
	prometheusutil "kubevirt.io/containerized-data-importer/pkg/util/prometheus"
)

func init() {
	klog.InitFlags(nil)
	flag.Parse()
}

func waitForReadyFile() {
	readyFile, _ := util.ParseEnvVar(common.ImporterReadyFile, false)
	if readyFile == "" {
		return
	}
	for {
		if _, err := os.Stat(readyFile); err == nil {
			break
		}
		time.Sleep(time.Second)
	}
}

func touchDoneFile() {
	doneFile, _ := util.ParseEnvVar(common.ImporterDoneFile, false)
	if doneFile == "" {
		return
	}
	f, err := os.OpenFile(doneFile, os.O_CREATE|os.O_EXCL, 0666)
	if err != nil {
		klog.Errorf("Failed creating file %s: %+v", doneFile, err)
	}
	f.Close()
}

func main() {
	defer klog.Flush()

	certsDirectory, err := ioutil.TempDir("", "certsdir")
	if err != nil {
		panic(err)
	}
	defer os.RemoveAll(certsDirectory)
	prometheusutil.StartPrometheusEndpoint(certsDirectory)
	klog.V(1).Infoln("Starting importer")

	source, _ := util.ParseEnvVar(common.ImporterSource, false)
	contentType, _ := util.ParseEnvVar(common.ImporterContentType, false)
	imageSize, _ := util.ParseEnvVar(common.ImporterImageSize, false)
	filesystemOverhead, _ := strconv.ParseFloat(os.Getenv(common.FilesystemOverheadVar), 64)
	preallocation, err := strconv.ParseBool(os.Getenv(common.Preallocation))

	volumeMode := v1.PersistentVolumeBlock
	if _, err := os.Stat(common.WriteBlockPath); os.IsNotExist(err) {
		volumeMode = v1.PersistentVolumeFilesystem
	} else {
		preallocation = true
	}

	// With writeback cache mode it's possible that the process will exit before all writes have been commited to storage.
	// To guarantee that our write was commited to storage, we make a fsync syscall and ensure success.
	// Also might be a good idea to sync any chmod's we might have done.
	defer fsyncDataFile(contentType, volumeMode)

	//Registry import currently support kubevirt content type only
	if contentType != string(cdiv1.DataVolumeKubeVirt) && (source == controller.SourceRegistry || source == controller.SourceImageio) {
		klog.Errorf("Unsupported content type %s when importing from %s", contentType, source)
		os.Exit(1)
	}

	availableDestSpace, err := util.GetAvailableSpaceByVolumeMode(volumeMode)
	if err != nil {
		klog.Errorf("%+v", err)
		os.Exit(1)
	}
	if source == controller.SourceNone {
		err := handleEmptyImage(contentType, imageSize, availableDestSpace, preallocation, volumeMode, filesystemOverhead)
		if err != nil {
			klog.Errorf("%+v", err)
			os.Exit(1)
		}

	} else {
		exitCode := handleImport(source, contentType, volumeMode, imageSize, filesystemOverhead, preallocation)
		if exitCode != 0 {
			os.Exit(exitCode)
		}
	}
}

func handleEmptyImage(contentType string, imageSize string, availableDestSpace int64, preallocation bool, volumeMode v1.PersistentVolumeMode, filesystemOverhead float64) error {
	var preallocationApplied bool

	if contentType == string(cdiv1.DataVolumeKubeVirt) {
		createBlankImage(imageSize, availableDestSpace, preallocation, volumeMode, filesystemOverhead)
		preallocationApplied = preallocation
	} else {
		errorEmptyDiskWithContentTypeArchive()
	}

	err := importCompleteTerminationMessage(preallocationApplied)
	return err
}

func handleImport(
	source string,
	contentType string,
	volumeMode v1.PersistentVolumeMode,
	imageSize string,
	filesystemOverhead float64,
	preallocation bool) int {
	klog.V(1).Infoln("begin import process")

	ds := newDataSource(source, contentType, volumeMode)
	defer ds.Close()

	processor := newDataProcessor(contentType, volumeMode, ds, imageSize, filesystemOverhead, preallocation)
	waitForReadyFile()
	err := processor.ProcessData()

	if err != nil {
		klog.Errorf("%+v", err)
		if err == importer.ErrRequiresScratchSpace {
			return common.ScratchSpaceNeededExitCode
		}
		err = util.WriteTerminationMessage(fmt.Sprintf("Unable to process data: %+v", err.Error()))
		if err != nil {
			klog.Errorf("%+v", err)
		}

		return 1
	}
	touchDoneFile()
	// due to the way some data sources can add additional information to termination message
	// after finished (ds.close() ) termination message has to be written first, before the
	// the ds is closed
	// TODO: think about making communication explicit, probably DS interface should be extended
	err = importCompleteTerminationMessage(processor.PreallocationApplied())
	if err != nil {
		klog.Errorf("%+v", err)
		return 1
	}

	return 0
}

func importCompleteTerminationMessage(preallocationApplied bool) error {
	message := "Import Complete"
	if preallocationApplied {
		message += ", " + common.PreallocationApplied
	}
	err := util.WriteTerminationMessage(message)
	if err != nil {
		return err
	}

	klog.V(1).Infoln(message)
	return nil
}

func newDataProcessor(contentType string, volumeMode v1.PersistentVolumeMode, ds importer.DataSourceInterface, imageSize string, filesystemOverhead float64, preallocation bool) *importer.DataProcessor {
	dest := getImporterDestPath(contentType, volumeMode)
	processor := importer.NewDataProcessor(ds, dest, common.ImporterDataDir, common.ScratchDataDir, imageSize, filesystemOverhead, preallocation)
	return processor
}

func getImporterDestPath(contentType string, volumeMode v1.PersistentVolumeMode) string {
	dest := common.ImporterWritePath

	if contentType == string(cdiv1.DataVolumeArchive) {
		dest = common.ImporterVolumePath
	}
	if volumeMode == v1.PersistentVolumeBlock {
		dest = common.WriteBlockPath
	}

	return dest
}

func newDataSource(source string, contentType string, volumeMode v1.PersistentVolumeMode) importer.DataSourceInterface {
	ep, _ := util.ParseEnvVar(common.ImporterEndpoint, false)
	acc, _ := util.ParseEnvVar(common.ImporterAccessKeyID, false)
	sec, _ := util.ParseEnvVar(common.ImporterSecretKey, false)
	diskID, _ := util.ParseEnvVar(common.ImporterDiskID, false)
	uuid, _ := util.ParseEnvVar(common.ImporterUUID, false)
	backingFile, _ := util.ParseEnvVar(common.ImporterBackingFile, false)
	certDir, _ := util.ParseEnvVar(common.ImporterCertDirVar, false)
	insecureTLS, _ := strconv.ParseBool(os.Getenv(common.InsecureTLSVar))
	thumbprint, _ := util.ParseEnvVar(common.ImporterThumbprint, false)

	currentCheckpoint, _ := util.ParseEnvVar(common.ImporterCurrentCheckpoint, false)
	previousCheckpoint, _ := util.ParseEnvVar(common.ImporterPreviousCheckpoint, false)
	finalCheckpoint, _ := util.ParseEnvVar(common.ImporterFinalCheckpoint, false)

	switch source {
	case controller.SourceHTTP:
		ds, err := importer.NewHTTPDataSource(ep, acc, sec, certDir, cdiv1.DataVolumeContentType(contentType))
		if err != nil {
			errorCannotConnectDataSource(err, "http")
		}
		return ds
	case controller.SourceImageio:
		ds, err := importer.NewImageioDataSource(ep, acc, sec, certDir, diskID, currentCheckpoint, previousCheckpoint)
		if err != nil {
			errorCannotConnectDataSource(err, "imageio")
		}
		return ds
	case controller.SourceRegistry:
		ds := importer.NewRegistryDataSource(ep, acc, sec, certDir, insecureTLS)
		return ds
	case controller.SourceS3:
		ds, err := importer.NewS3DataSource(ep, acc, sec, certDir)
		if err != nil {
			errorCannotConnectDataSource(err, "s3")
		}
		return ds
	case controller.SourceVDDK:
		ds, err := importer.NewVDDKDataSource(ep, acc, sec, thumbprint, uuid, backingFile, currentCheckpoint, previousCheckpoint, finalCheckpoint, volumeMode)
		if err != nil {
			errorCannotConnectDataSource(err, "vddk")
		}
		return ds
	default:
		klog.Errorf("Unknown source type %s\n", source)
		err := util.WriteTerminationMessage(fmt.Sprintf("Unknown data source: %s", source))
		if err != nil {
			klog.Errorf("%+v", err)
		}
		os.Exit(1)
	}

	return nil
}

func createBlankImage(imageSize string, availableDestSpace int64, preallocation bool, volumeMode v1.PersistentVolumeMode, filesystemOverhead float64) {
	requestImageSizeQuantity := resource.MustParse(imageSize)
	minSizeQuantity := util.MinQuantity(resource.NewScaledQuantity(availableDestSpace, 0), &requestImageSizeQuantity)

	if minSizeQuantity.Cmp(requestImageSizeQuantity) != 0 {
		// Available dest space is smaller than the size we want to create
		klog.Warningf("Available space less than requested size, creating blank image sized to available space: %s.\n", minSizeQuantity.String())
	}

	var err error
	if volumeMode == v1.PersistentVolumeFilesystem {
		quantityWithFSOverhead := util.GetUsableSpace(filesystemOverhead, minSizeQuantity.Value())
		klog.Infof("Space adjusted for filesystem overhead: %d.\n", quantityWithFSOverhead)
		err = image.CreateBlankImage(common.ImporterWritePath, *resource.NewScaledQuantity(quantityWithFSOverhead, 0), preallocation)
	} else if volumeMode == v1.PersistentVolumeBlock && preallocation {
		klog.V(1).Info("Preallocating blank block volume")
		err = image.PreallocateBlankBlock(common.WriteBlockPath, minSizeQuantity)
	}

	if err != nil {
		klog.Errorf("%+v", err)
		message := fmt.Sprintf("Unable to create blank image: %+v", err)
		err = util.WriteTerminationMessage(message)
		if err != nil {
			klog.Errorf("%+v", err)
		}
		os.Exit(1)
	}
}

func errorCannotConnectDataSource(err error, dsName string) {
	klog.Errorf("%+v", err)
	err = util.WriteTerminationMessage(fmt.Sprintf("Unable to connect to %s data source: %+v", dsName, err))
	if err != nil {
		klog.Errorf("%+v", err)
	}
	os.Exit(1)
}

func errorEmptyDiskWithContentTypeArchive() {
	klog.Errorf("%+v", errors.New("Cannot create empty disk with content type archive"))
	err := util.WriteTerminationMessage("Cannot create empty disk with content type archive")
	if err != nil {
		klog.Errorf("%+v", err)
	}
	os.Exit(1)
}

func fsyncDataFile(contentType string, volumeMode v1.PersistentVolumeMode) {
	dataFile := getImporterDestPath(contentType, volumeMode)
	file, err := os.Open(dataFile)
	if err != nil {
		klog.Errorf("could not get file descriptor for fsync call: %+v", err)
		os.Exit(1)
	}
	if err := file.Sync(); err != nil {
		klog.Errorf("could not fsync following qemu-img writing: %+v", err)
		os.Exit(1)
	}
	klog.V(3).Infof("Successfully completed fsync(%s) syscall, commited to disk\n", dataFile)
	file.Close()
}
<|endoftext|>


<commit_before>// Example bag-of-words shows how to compute a bag-of-words vector given a
// fragment library and a PDB file.
package main

import (
	"flag"
	"fmt"
	"os"
	"path"

	"github.com/BurntSushi/bcbgo/fragbag"
	"github.com/BurntSushi/bcbgo/pdb"
)

func main() {
	if flag.NArg() < 2 {
		usage()
	}

	// Initialize the fragment library whatever is provided. If the library
	// isn't valid or doesn't exist, exit with an error.
	lib, err := fragbag.NewLibrary(flag.Arg(0))
	if err != nil {
		fmt.Println(err)
		return
	}
	fmt.Printf("Using library %s.\n", lib)

	for _, pdbfile := range flag.Args()[1:] {
		entry, err := pdb.New(pdbfile)
		if err != nil {
			fmt.Println(err)
			return
		}
		fmt.Printf("Computing the bag-of-words vector for %s.\n", entry.Name())
		fmt.Println(lib.NewBowPDB(entry))
		fmt.Println("----------------------------------------------")
	}
}

func init() {
	flag.Usage = usage
	flag.Parse()
}

func usage() {
	fmt.Fprintf(os.Stderr,
		"Usage: %s frag-lib-directory pdb-file [ pdb-file ... ]\n",
		path.Base(os.Args[0]))
	flag.PrintDefaults()
	fmt.Fprintf(os.Stderr,
		"\nex. './%s ../../../data/fraglibs/centers400_11 "+
			"../../../data/samples/1ctf.pdb'\n",
		path.Base(os.Args[0]))
	os.Exit(1)
}
<commit_msg>Starting command to create PDB index.<commit_after>// Example bag-of-words shows how to compute a bag-of-words vector given a
// fragment library and a PDB file.
package main

import (
	"flag"
	"fmt"
	"os"
	"path"
	"runtime"

	"github.com/BurntSushi/bcbgo/fragbag"
	"github.com/BurntSushi/bcbgo/pdb"
)

var (
	flagGoMaxProcs = runtime.NumCPU()
)

func main() {
	if flag.NArg() < 2 {
		usage()
	}

	// Initialize the fragment library whatever is provided. If the library
	// isn't valid or doesn't exist, exit with an error.
	lib, err := fragbag.NewLibrary(flag.Arg(0))
	if err != nil {
		fmt.Println(err)
		return
	}
	fmt.Printf("Using library %s.\n", lib)

	pdbFiles := flag.Args()[1:]
	entries := make([]*pdb.Entry, len(pdbFiles))
	for i, pdbfile := range flag.Args()[1:] {
		entry, err := pdb.New(pdbfile)
		if err != nil {
			fmt.Println(err)
			return
		}
		entries[i] = entry
	}
	bows := lib.NewBowsPDBList(entries...)
	for i, entry := range entries {
		fmt.Printf("Computing the bag-of-words vector for %s.\n", entry.Name())
		fmt.Println(bows[i])
		// fmt.Println(lib.NewBowPDB(entry))
		fmt.Println("----------------------------------------------")
	}
}

func init() {
	flag.IntVar(&flagGoMaxProcs, "p", flagGoMaxProcs,
		"The maximum number of CPUs that can be executing simultaneously.")
	flag.Usage = usage
	flag.Parse()

	runtime.GOMAXPROCS(flagGoMaxProcs)
}

func usage() {
	fmt.Fprintf(os.Stderr,
		"Usage: %s frag-lib-directory pdb-file [ pdb-file ... ]\n",
		path.Base(os.Args[0]))
	flag.PrintDefaults()
	fmt.Fprintf(os.Stderr,
		"\nex. './%s ../../../data/fraglibs/centers400_11 "+
			"../../../data/samples/1ctf.pdb'\n",
		path.Base(os.Args[0]))
	os.Exit(1)
}
<|endoftext|>


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

package main

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
)

const gomobileHash = "5d9a33257ab559d10fa8c96087aed99bf8b6d868"

func runCommand(command string, args []string, env []string) error {
	if buildX || buildN {
		for _, e := range env {
			fmt.Printf("%s ", e)
		}
		fmt.Print(command)
		for _, arg := range args {
			fmt.Printf(" %s", arg)
		}
		fmt.Println()
	}

	if buildN {
		return nil
	}

	cmd := exec.Command(command, args...)
	if len(env) > 0 {
		cmd.Env = append(os.Environ(), env...)
	}
	if out, err := cmd.CombinedOutput(); err != nil {
		return fmt.Errorf("%s %v failed: %v\n%v", command, args, string(out), err)
	}
	return nil
}

func removeAll(path string) error {
	if buildX || buildN {
		fmt.Printf("rm -rf %s\n", path)
	}
	if buildN {
		return nil
	}
	return os.RemoveAll(path)
}

func runGo(args ...string) error {
	// TODO: Remove this after Ebiten drops the support of Go 1.15 and older.
	// GO111MODULE is on by default as of Go 1.16.
	env := []string{
		"GO111MODULE=on",
	}
	return runCommand("go", args, env)
}

// exe adds the .exe extension to the given filename.
// Without .exe, the executable won't be found by exec.LookPath on Windows (#1096).
func exe(filename string) string {
	if runtime.GOOS == "windows" {
		return filename + ".exe"
	}
	return filename
}

func prepareGomobileCommands() (string, error) {
	tmp, err := ioutil.TempDir("", "ebitenmobile-")
	if err != nil {
		return "", err
	}

	newpath := filepath.Join(tmp, "bin")
	if path := os.Getenv("PATH"); path != "" {
		newpath += string(filepath.ListSeparator) + path
	}
	if buildX || buildN {
		fmt.Printf("PATH=%s\n", newpath)
	}
	if !buildN {
		if err := os.Setenv("PATH", newpath); err != nil {
			return tmp, err
		}
	}

	pwd, err := os.Getwd()
	if err != nil {
		return tmp, err
	}

	// cd
	if buildX {
		fmt.Printf("cd %s\n", tmp)
	}
	if err := os.Chdir(tmp); err != nil {
		return tmp, err
	}
	defer func() {
		os.Chdir(pwd)
	}()

	const (
		modname   = "ebitenmobiletemporary"
		buildtags = "//go:build tools" +
			"\n// +build tools"
	)
	if err := runGo("mod", "init", modname); err != nil {
		return tmp, err
	}
	if err := ioutil.WriteFile("tools.go", []byte(fmt.Sprintf(`%s

package %s

import (
	_ "golang.org/x/mobile/cmd/gobind"
	_ "golang.org/x/mobile/cmd/gomobile"
)
`, buildtags, modname)), 0644); err != nil {
		return tmp, err
	}

	// To record gomobile to go.sum for Go 1.16 and later, go-get gomobile instaed of golang.org/x/mobile (#1487).
	// This also records gobind as gomobile depends on gobind indirectly.
	// Using `...` doesn't work on Windows since mobile/internal/mobileinit cannot be compiled on Windows w/o Cgo (#1493).
	if err := runGo("get", "golang.org/x/mobile/cmd/gomobile@"+gomobileHash); err != nil {
		return tmp, err
	}
	if localgm := os.Getenv("EBITENMOBILE_GOMOBILE"); localgm != "" {
		if !filepath.IsAbs(localgm) {
			localgm = filepath.Join(pwd, localgm)
		}
		if err := runGo("mod", "edit", "-replace=golang.org/x/mobile="+localgm); err != nil {
			return tmp, err
		}
	}
	if err := runGo("mod", "tidy"); err != nil {
		return tmp, err
	}
	if err := runGo("build", "-o", exe(filepath.Join("bin", "gomobile")), "golang.org/x/mobile/cmd/gomobile"); err != nil {
		return tmp, err
	}
	if err := runGo("build", "-o", exe(filepath.Join("bin", "gobind-original")), "golang.org/x/mobile/cmd/gobind"); err != nil {
		return tmp, err
	}

	if err := os.Mkdir("src", 0755); err != nil {
		return tmp, err
	}
	if err := ioutil.WriteFile(filepath.Join("src", "gobind.go"), gobindsrc, 0644); err != nil {
		return tmp, err
	}

	if err := runGo("build", "-o", exe(filepath.Join("bin", "gobind")), "-tags", "ebitenmobilegobind", filepath.Join("src", "gobind.go")); err != nil {
		return tmp, err
	}

	if err := runCommand("gomobile", []string{"init"}, nil); err != nil {
		return tmp, err
	}

	return tmp, nil
}
<commit_msg>cmd/ebitenmobile: update gomobile version<commit_after>// Copyright 2019 The Ebiten Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
)

const gomobileHash = "4e6c2922fdeed32d3596616518aaee7b0d79ce55"

func runCommand(command string, args []string, env []string) error {
	if buildX || buildN {
		for _, e := range env {
			fmt.Printf("%s ", e)
		}
		fmt.Print(command)
		for _, arg := range args {
			fmt.Printf(" %s", arg)
		}
		fmt.Println()
	}

	if buildN {
		return nil
	}

	cmd := exec.Command(command, args...)
	if len(env) > 0 {
		cmd.Env = append(os.Environ(), env...)
	}
	if out, err := cmd.CombinedOutput(); err != nil {
		return fmt.Errorf("%s %v failed: %v\n%v", command, args, string(out), err)
	}
	return nil
}

func removeAll(path string) error {
	if buildX || buildN {
		fmt.Printf("rm -rf %s\n", path)
	}
	if buildN {
		return nil
	}
	return os.RemoveAll(path)
}

func runGo(args ...string) error {
	// TODO: Remove this after Ebiten drops the support of Go 1.15 and older.
	// GO111MODULE is on by default as of Go 1.16.
	env := []string{
		"GO111MODULE=on",
	}
	return runCommand("go", args, env)
}

// exe adds the .exe extension to the given filename.
// Without .exe, the executable won't be found by exec.LookPath on Windows (#1096).
func exe(filename string) string {
	if runtime.GOOS == "windows" {
		return filename + ".exe"
	}
	return filename
}

func prepareGomobileCommands() (string, error) {
	tmp, err := ioutil.TempDir("", "ebitenmobile-")
	if err != nil {
		return "", err
	}

	newpath := filepath.Join(tmp, "bin")
	if path := os.Getenv("PATH"); path != "" {
		newpath += string(filepath.ListSeparator) + path
	}
	if buildX || buildN {
		fmt.Printf("PATH=%s\n", newpath)
	}
	if !buildN {
		if err := os.Setenv("PATH", newpath); err != nil {
			return tmp, err
		}
	}

	pwd, err := os.Getwd()
	if err != nil {
		return tmp, err
	}

	// cd
	if buildX {
		fmt.Printf("cd %s\n", tmp)
	}
	if err := os.Chdir(tmp); err != nil {
		return tmp, err
	}
	defer func() {
		os.Chdir(pwd)
	}()

	const (
		modname   = "ebitenmobiletemporary"
		buildtags = "//go:build tools" +
			"\n// +build tools"
	)
	if err := runGo("mod", "init", modname); err != nil {
		return tmp, err
	}
	if err := ioutil.WriteFile("tools.go", []byte(fmt.Sprintf(`%s

package %s

import (
	_ "golang.org/x/mobile/cmd/gobind"
	_ "golang.org/x/mobile/cmd/gomobile"
)
`, buildtags, modname)), 0644); err != nil {
		return tmp, err
	}

	// To record gomobile to go.sum for Go 1.16 and later, go-get gomobile instaed of golang.org/x/mobile (#1487).
	// This also records gobind as gomobile depends on gobind indirectly.
	// Using `...` doesn't work on Windows since mobile/internal/mobileinit cannot be compiled on Windows w/o Cgo (#1493).
	if err := runGo("get", "golang.org/x/mobile/cmd/gomobile@"+gomobileHash); err != nil {
		return tmp, err
	}
	if localgm := os.Getenv("EBITENMOBILE_GOMOBILE"); localgm != "" {
		if !filepath.IsAbs(localgm) {
			localgm = filepath.Join(pwd, localgm)
		}
		if err := runGo("mod", "edit", "-replace=golang.org/x/mobile="+localgm); err != nil {
			return tmp, err
		}
	}
	if err := runGo("mod", "tidy"); err != nil {
		return tmp, err
	}
	if err := runGo("build", "-o", exe(filepath.Join("bin", "gomobile")), "golang.org/x/mobile/cmd/gomobile"); err != nil {
		return tmp, err
	}
	if err := runGo("build", "-o", exe(filepath.Join("bin", "gobind-original")), "golang.org/x/mobile/cmd/gobind"); err != nil {
		return tmp, err
	}

	if err := os.Mkdir("src", 0755); err != nil {
		return tmp, err
	}
	if err := ioutil.WriteFile(filepath.Join("src", "gobind.go"), gobindsrc, 0644); err != nil {
		return tmp, err
	}

	if err := runGo("build", "-o", exe(filepath.Join("bin", "gobind")), "-tags", "ebitenmobilegobind", filepath.Join("src", "gobind.go")); err != nil {
		return tmp, err
	}

	if err := runCommand("gomobile", []string{"init"}, nil); err != nil {
		return tmp, err
	}

	return tmp, nil
}
<|endoftext|>


<commit_before>package disk

import (
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"strconv"
	"sync"
	"time"

	"github.com/funkygao/golib/sync2"
	log "github.com/funkygao/log4go"
)

// queue is a bounded, disk-backed, append-only type that combines queue and
// log semantics.
// key/value byte slices can be appended and read back in order through
// cursor.
//
// Internally, the queue writes key/value byte slices to multiple segment files so
// that disk space can be reclaimed. When a segment file is larger than
// the max segment size, a new file is created.   Segments are removed
// after cursor has advanced past the last entry.  The first
// segment is the head, and the last segment is the tail.  Reads are from
// the head segment and writes tail segment.
//
// queues can have a max size configured such that when the size of all
// segments on disk exceeds the size, write will fail.
//
// ┌─────┐
// │head │
// ├─────┘
// │
// ▼
// ┌─────────────────┐ ┌─────────────────┐┌─────────────────┐
// │segment 1 - 10MB │ │segment 2 - 10MB ││segment 3 - 10MB │
// └─────────────────┘ └─────────────────┘└─────────────────┘
//                          ▲                               ▲
//                          │                               │
//                          │                               │
//                       ┌───────┐                     ┌─────┐
//                       │cursor │                     │tail │
//                       └───────┘                     └─────┘
type queue struct {
	mu sync.RWMutex
	wg sync.WaitGroup

	baseDir      string // the slot this queue is using
	dir          string // Directory to create segments
	clusterTopic clusterTopic

	// The maximum size in bytes of a segment file before a new one should be created
	maxSegmentSize int64

	// The maximum size allowed in bytes of all segments before writes will return an error
	// -1 means unlimited
	maxSize int64

	inflights         sync2.AtomicInt64
	appendN, deliverN sync2.AtomicInt64

	purgeInterval time.Duration
	maxAge        time.Duration

	cursor     *cursor
	index      *index
	head, tail *segment
	segments   segments

	quit          chan struct{}
	emptyInflight sync2.AtomicInt32
}

// newQueue create a queue that will store segments in dir and that will
// consume more than maxSize on disk.
func newQueue(baseDir string, ct clusterTopic, maxSize int64, purgeInterval, maxAge time.Duration) *queue {
	q := &queue{
		clusterTopic:   ct,
		baseDir:        baseDir,
		dir:            ct.TopicDir(baseDir),
		maxSegmentSize: defaultSegmentSize,
		maxSize:        maxSize,
		purgeInterval:  purgeInterval,
		maxAge:         maxAge,
		segments:       segments{},
	}
	q.cursor = newCursor(q)
	q.index = newIndex(q)
	return q
}

// Open opens the queue for reading and writing
func (q *queue) Open() error {
	q.mu.Lock()
	defer q.mu.Unlock()

	if err := mkdirIfNotExist(q.dir); err != nil {
		return err
	}

	var (
		minId            uint64 = 0
		moveCursorToHead bool   = false
	)
	if err := q.cursor.open(); err != nil {
		// cursor file might not exist or json file corrupts
		log.Warn("queue[%s] cursor: %s, move to head", q.ident(), err)
		moveCursorToHead = true
	} else {
		// load segments from cursor checkpoint
		minId = q.cursor.pos.SegmentID
	}

	segments, err := q.loadSegments(minId)
	if err != nil {
		return err
	}
	q.segments = segments

	if len(q.segments) == 0 {
		// create the 1st segment
		if _, err = q.addSegment(); err != nil {
			return err
		}
	}

	q.head = q.segments[0]
	q.tail = q.segments[len(q.segments)-1]

	// cursor open must be placed below queue open
	if err = q.cursor.initPosition(moveCursorToHead); err != nil {
		return err
	}

	if q.cursor.seg != q.tail || q.cursor.pos.Offset != q.tail.DiskUsage() {
		q.emptyInflight.Set(0)
	}

	return nil
}

func (q *queue) Start() {
	q.quit = make(chan struct{})

	q.wg.Add(1)
	go q.housekeeping()

	q.wg.Add(1)
	go q.pump()
}

// Close stops the queue for reading and writing
func (q *queue) Close() error {
	close(q.quit)

	q.mu.Lock()
	defer q.mu.Unlock()

	for _, s := range q.segments {
		if err := s.Close(); err != nil {
			return err
		}
	}

	q.head = nil
	q.tail = nil
	q.segments = nil

	q.wg.Wait()
	if err := q.cursor.dump(); err != nil {
		return err
	}
	q.cursor = nil
	return nil
}

func (q *queue) Inflights() int64 {
	return q.inflights.Get()
}

func (q *queue) AppendN() int64 {
	return q.appendN.Get()
}

func (q *queue) DeliverN() int64 {
	return q.deliverN.Get()
}

// Remove removes all underlying file-based resources for the queue.
// It is an error to call this on an open queue.
func (q *queue) Remove() (err error) {
	q.mu.Lock()
	defer q.mu.Unlock()

	if q.head != nil || q.tail != nil || q.segments != nil {
		return ErrQueueOpen
	}

	if err = os.RemoveAll(q.dir); err == nil {
		q.emptyInflight.Set(1)
	}
	return
}

// Purge garbage collects the segments that are behind cursor.
func (q *queue) Purge() error {
	q.mu.Lock()
	defer q.mu.Unlock()

	if len(q.segments) <= 1 {
		// head, curror, tail are in the same segment
		return nil
	}

	for {
		if q.cursor.pos.SegmentID > q.head.id &&
			q.head.LastModified().Add(q.maxAge).Unix() < time.Now().Unix() {
			q.trimHead()
		} else {
			return nil
		}

	}
}

// LastModified returns the last time the queue was modified.
func (q *queue) LastModified() time.Time {
	q.mu.RLock()
	defer q.mu.RUnlock()

	return q.tail.LastModified()
}

// Append appends a block to the end of the queue
func (q *queue) Append(b *block) error {
	q.mu.Lock()
	defer q.mu.Unlock()

	if q.tail == nil {
		return ErrQueueNotOpen
	}

	if q.maxSize > 0 && q.diskUsage()+b.size() > q.maxSize {
		return ErrQueueFull
	}

	// Append the entry to the tail, if the segment is full,
	// try to create new segment and retry the append
	if err := q.tail.Append(b); err == ErrSegmentFull {
		segment, err := q.addSegment()
		if err != nil {
			return err
		}

		q.tail = segment
		err = q.tail.Append(b)
		if err == nil {
			q.emptyInflight.Set(0)
			q.inflights.Add(1)
			q.appendN.Add(1)
		}
		return err
	} else if err != nil {
		return err
	}

	q.emptyInflight.Set(0)
	q.appendN.Add(1)
	q.inflights.Add(1)
	return nil
}

func (q *queue) Rollback(b *block) (err error) {
	c := q.cursor
	if err = c.advanceOffset(-b.size()); err != nil {
		return
	}

	// rollback needn't consider cross segment case
	q.emptyInflight.Set(0)
	return c.seg.Seek(c.pos.Offset)
}

func (q *queue) Next(b *block) (err error) {
	q.mu.RLock()
	defer q.mu.RUnlock()

	c := q.cursor
	if c == nil {
		return ErrQueueNotOpen
	}
	err = c.seg.ReadOne(b)
	switch err {
	case nil:
		q.emptyInflight.Set(0)
		return c.advanceOffset(b.size())

	case io.EOF:
		// cursor might have:
		// 1. reached end of the current segment: will advance to next segment
		// 2. reached end of tail
		if ok := c.advanceSegment(); !ok {
			q.emptyInflight.Set(1)
			return ErrEOQ
		}

		// advanced to next segment, read one block
		err = c.seg.ReadOne(b)
		switch err {
		case nil:
			// bingo!
			return c.advanceOffset(b.size())

		case io.EOF:
			// tail is empty
			return ErrEOQ

		default:
			return
		}

	default:
		return
	}
}

func (q *queue) EmptyInflight() bool {
	return q.emptyInflight.Get() == 1
}

// diskUsage returns the total size on disk used by the queue
func (q *queue) diskUsage() int64 {
	var size int64
	for _, s := range q.segments {
		size += s.DiskUsage()
	}
	return size
}

// loadSegments loads all in-range segments on disk
// FIXME manage q.inflights counter while loading segments
func (q *queue) loadSegments(minId uint64) (segments, error) {
	segments := []*segment{}

	files, err := ioutil.ReadDir(q.dir)
	if err != nil {
		return segments, err
	}

	for _, segment := range files {
		if segment.IsDir() || segment.Name() == cursorFile {
			continue
		}

		// segment file names are all numeric
		id, err := strconv.ParseUint(segment.Name(), 10, 64)
		if err != nil {
			log.Error("queue[%s] segment:%s %s", q.ident(), segment.Name(), err)
			continue
		}
		if id < minId {
			log.Debug("queue[%s] skip stale segment:%s", q.ident(), segment.Name())
			continue
		}

		segment, err := newSegment(id, filepath.Join(q.dir, segment.Name()), q.maxSegmentSize)
		if err != nil {
			return segments, err
		}

		segments = append(segments, segment)
	}
	return segments, nil
}

// addSegment creates a new empty segment file
// caller is responsible for the lock
func (q *queue) addSegment() (*segment, error) {
	nextID, err := q.nextSegmentID()
	if err != nil {
		return nil, err
	}

	path := filepath.Join(q.dir, fmt.Sprintf("%020d", nextID))
	segment, err := newSegment(nextID, path, q.maxSegmentSize)
	if err != nil {
		return nil, err
	}

	q.segments = append(q.segments, segment)
	return segment, nil
}

// nextSegmentID returns the next segment ID that is free
func (q *queue) nextSegmentID() (uint64, error) {
	segments, err := ioutil.ReadDir(q.dir)
	if err != nil {
		return 0, err
	}

	var maxID uint64
	for _, segment := range segments {
		if segment.IsDir() || segment.Name() == cursorFile {
			continue
		}

		// Segments file names are all numeric
		segmentID, err := strconv.ParseUint(segment.Name(), 10, 64)
		if err != nil {
			log.Warn("unexpected segment file: %s", filepath.Join(q.dir, segment.Name()))
			continue
		}

		if segmentID > maxID {
			maxID = segmentID
		}
	}

	return maxID + 1, nil
}

func (q *queue) ident() string {
	return q.dir
}

func (q *queue) trimHead() (err error) {
	if len(q.segments) <= 1 {
		return ErrHeadIsTail
	}

	q.segments = q.segments[1:]

	if err = q.head.Remove(); err != nil {
		return
	}

	q.head = q.segments[0]
	return
}

// TODO skipCursorSegment skip the current corrupted cursor segment and
// advance to next segment.
// if tail corrupts, add new segment.
func (q *queue) skipCursorSegment() {

}
<commit_msg>easy RWLock management<commit_after>package disk

import (
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"strconv"
	"sync"
	"time"

	"github.com/funkygao/golib/sync2"
	log "github.com/funkygao/log4go"
)

// queue is a bounded, disk-backed, append-only type that combines queue and
// log semantics.
// key/value byte slices can be appended and read back in order through
// cursor.
//
// Internally, the queue writes key/value byte slices to multiple segment files so
// that disk space can be reclaimed. When a segment file is larger than
// the max segment size, a new file is created.   Segments are removed
// after cursor has advanced past the last entry.  The first
// segment is the head, and the last segment is the tail.  Reads are from
// the head segment and writes tail segment.
//
// queues can have a max size configured such that when the size of all
// segments on disk exceeds the size, write will fail.
//
// ┌─────┐
// │head │
// ├─────┘
// │
// ▼
// ┌─────────────────┐ ┌─────────────────┐┌─────────────────┐
// │segment 1 - 10MB │ │segment 2 - 10MB ││segment 3 - 10MB │
// └─────────────────┘ └─────────────────┘└─────────────────┘
//                          ▲                               ▲
//                          │                               │
//                          │                               │
//                       ┌───────┐                     ┌─────┐
//                       │cursor │                     │tail │
//                       └───────┘                     └─────┘
type queue struct {
	mu sync.RWMutex
	wg sync.WaitGroup

	baseDir      string // the slot this queue is using
	dir          string // Directory to create segments
	clusterTopic clusterTopic

	// The maximum size in bytes of a segment file before a new one should be created
	maxSegmentSize int64

	// The maximum size allowed in bytes of all segments before writes will return an error
	// -1 means unlimited
	maxSize int64

	inflights         sync2.AtomicInt64
	appendN, deliverN sync2.AtomicInt64

	purgeInterval time.Duration
	maxAge        time.Duration

	cursor     *cursor
	index      *index
	head, tail *segment
	segments   segments

	quit          chan struct{}
	emptyInflight sync2.AtomicInt32
}

// newQueue create a queue that will store segments in dir and that will
// consume more than maxSize on disk.
func newQueue(baseDir string, ct clusterTopic, maxSize int64, purgeInterval, maxAge time.Duration) *queue {
	q := &queue{
		clusterTopic:   ct,
		baseDir:        baseDir,
		dir:            ct.TopicDir(baseDir),
		maxSegmentSize: defaultSegmentSize,
		maxSize:        maxSize,
		purgeInterval:  purgeInterval,
		maxAge:         maxAge,
		segments:       segments{},
	}
	q.cursor = newCursor(q)
	q.index = newIndex(q)
	return q
}

// Open opens the queue for reading and writing
func (q *queue) Open() error {
	q.mu.Lock()
	defer q.mu.Unlock()

	if err := mkdirIfNotExist(q.dir); err != nil {
		return err
	}

	var (
		minId            uint64 = 0
		moveCursorToHead bool   = false
	)
	if err := q.cursor.open(); err != nil {
		// cursor file might not exist or json file corrupts
		log.Warn("queue[%s] cursor: %s, move to head", q.ident(), err)
		moveCursorToHead = true
	} else {
		// load segments from cursor checkpoint
		minId = q.cursor.pos.SegmentID
	}

	segments, err := q.loadSegments(minId)
	if err != nil {
		return err
	}
	q.segments = segments

	if len(q.segments) == 0 {
		// create the 1st segment
		if _, err = q.addSegment(); err != nil {
			return err
		}
	}

	q.head = q.segments[0]
	q.tail = q.segments[len(q.segments)-1]

	// cursor open must be placed below queue open
	if err = q.cursor.initPosition(moveCursorToHead); err != nil {
		return err
	}

	if q.cursor.seg != q.tail || q.cursor.pos.Offset != q.tail.DiskUsage() {
		q.emptyInflight.Set(0)
	}

	return nil
}

func (q *queue) Start() {
	q.quit = make(chan struct{})

	q.wg.Add(1)
	go q.housekeeping()

	q.wg.Add(1)
	go q.pump()
}

// Close stops the queue for reading and writing
func (q *queue) Close() error {
	close(q.quit)

	q.mu.Lock()
	defer q.mu.Unlock()

	for _, s := range q.segments {
		if err := s.Close(); err != nil {
			return err
		}
	}

	q.head = nil
	q.tail = nil
	q.segments = nil

	q.wg.Wait()
	if err := q.cursor.dump(); err != nil {
		return err
	}
	q.cursor = nil
	return nil
}

func (q *queue) Inflights() int64 {
	return q.inflights.Get()
}

func (q *queue) AppendN() int64 {
	return q.appendN.Get()
}

func (q *queue) DeliverN() int64 {
	return q.deliverN.Get()
}

// Remove removes all underlying file-based resources for the queue.
// It is an error to call this on an open queue.
func (q *queue) Remove() (err error) {
	q.mu.Lock()
	defer q.mu.Unlock()

	if q.head != nil || q.tail != nil || q.segments != nil {
		return ErrQueueOpen
	}

	if err = os.RemoveAll(q.dir); err == nil {
		q.emptyInflight.Set(1)
	}
	return
}

// Purge garbage collects the segments that are behind cursor.
func (q *queue) Purge() error {
	q.mu.Lock()
	defer q.mu.Unlock()

	if len(q.segments) <= 1 {
		// head, curror, tail are in the same segment
		return nil
	}

	for {
		if q.cursor.pos.SegmentID > q.head.id &&
			q.head.LastModified().Add(q.maxAge).Unix() < time.Now().Unix() {
			q.trimHead()
		} else {
			return nil
		}

	}
}

// LastModified returns the last time the queue was modified.
func (q *queue) LastModified() time.Time {
	q.mu.RLock()
	defer q.mu.RUnlock()

	return q.tail.LastModified()
}

// Append appends a block to the end of the queue
func (q *queue) Append(b *block) error {
	q.mu.Lock()
	defer q.mu.Unlock()

	if q.tail == nil {
		return ErrQueueNotOpen
	}

	if q.maxSize > 0 && q.diskUsage()+b.size() > q.maxSize {
		return ErrQueueFull
	}

	// Append the entry to the tail, if the segment is full,
	// try to create new segment and retry the append
	if err := q.tail.Append(b); err == ErrSegmentFull {
		segment, err := q.addSegment()
		if err != nil {
			return err
		}

		q.tail = segment
		err = q.tail.Append(b)
		if err == nil {
			q.emptyInflight.Set(0)
			q.inflights.Add(1)
			q.appendN.Add(1)
		}
		return err
	} else if err != nil {
		return err
	}

	q.emptyInflight.Set(0)
	q.appendN.Add(1)
	q.inflights.Add(1)
	return nil
}

func (q *queue) Rollback(b *block) (err error) {
	c := q.cursor
	if err = c.advanceOffset(-b.size()); err != nil {
		return
	}

	// rollback needn't consider cross segment case
	q.emptyInflight.Set(0)
	return c.seg.Seek(c.pos.Offset)
}

func (q *queue) Next(b *block) (err error) {
	q.mu.RLock()
	c := q.cursor
	q.mu.RUnlock()

	if c == nil {
		return ErrQueueNotOpen
	}

	err = c.seg.ReadOne(b)
	switch err {
	case nil:
		q.emptyInflight.Set(0)
		return c.advanceOffset(b.size())

	case io.EOF:
		// cursor might have:
		// 1. reached end of the current segment: will advance to next segment
		// 2. reached end of tail
		if ok := c.advanceSegment(); !ok {
			q.emptyInflight.Set(1)
			return ErrEOQ
		}

		// advanced to next segment, read one block
		err = c.seg.ReadOne(b)
		switch err {
		case nil:
			// bingo!
			return c.advanceOffset(b.size())

		case io.EOF:
			// tail is empty
			return ErrEOQ

		default:
			return
		}

	default:
		return
	}
}

func (q *queue) EmptyInflight() bool {
	return q.emptyInflight.Get() == 1
}

// diskUsage returns the total size on disk used by the queue
func (q *queue) diskUsage() int64 {
	var size int64
	for _, s := range q.segments {
		size += s.DiskUsage()
	}
	return size
}

// loadSegments loads all in-range segments on disk
// FIXME manage q.inflights counter while loading segments
func (q *queue) loadSegments(minId uint64) (segments, error) {
	segments := []*segment{}

	files, err := ioutil.ReadDir(q.dir)
	if err != nil {
		return segments, err
	}

	for _, segment := range files {
		if segment.IsDir() || segment.Name() == cursorFile {
			continue
		}

		// segment file names are all numeric
		id, err := strconv.ParseUint(segment.Name(), 10, 64)
		if err != nil {
			log.Error("queue[%s] segment:%s %s", q.ident(), segment.Name(), err)
			continue
		}
		if id < minId {
			log.Debug("queue[%s] skip stale segment:%s", q.ident(), segment.Name())
			continue
		}

		segment, err := newSegment(id, filepath.Join(q.dir, segment.Name()), q.maxSegmentSize)
		if err != nil {
			return segments, err
		}

		segments = append(segments, segment)
	}
	return segments, nil
}

// addSegment creates a new empty segment file
// caller is responsible for the lock
func (q *queue) addSegment() (*segment, error) {
	nextID, err := q.nextSegmentID()
	if err != nil {
		return nil, err
	}

	path := filepath.Join(q.dir, fmt.Sprintf("%020d", nextID))
	segment, err := newSegment(nextID, path, q.maxSegmentSize)
	if err != nil {
		return nil, err
	}

	q.segments = append(q.segments, segment)
	return segment, nil
}

// nextSegmentID returns the next segment ID that is free
func (q *queue) nextSegmentID() (uint64, error) {
	segments, err := ioutil.ReadDir(q.dir)
	if err != nil {
		return 0, err
	}

	var maxID uint64
	for _, segment := range segments {
		if segment.IsDir() || segment.Name() == cursorFile {
			continue
		}

		// Segments file names are all numeric
		segmentID, err := strconv.ParseUint(segment.Name(), 10, 64)
		if err != nil {
			log.Warn("unexpected segment file: %s", filepath.Join(q.dir, segment.Name()))
			continue
		}

		if segmentID > maxID {
			maxID = segmentID
		}
	}

	return maxID + 1, nil
}

func (q *queue) ident() string {
	return q.dir
}

func (q *queue) trimHead() (err error) {
	if len(q.segments) <= 1 {
		return ErrHeadIsTail
	}

	q.segments = q.segments[1:]

	if err = q.head.Remove(); err != nil {
		return
	}

	q.head = q.segments[0]
	return
}

// TODO skipCursorSegment skip the current corrupted cursor segment and
// advance to next segment.
// if tail corrupts, add new segment: Append will write to new segment and
// reader read from the new segment.
func (q *queue) skipCursorSegment() {

}
<|endoftext|>