<commit_before>package llbsolver

import (
	"context"
	"fmt"
	"strings"
	"sync"
	"time"

	"github.com/containerd/containerd/platforms"
	"github.com/mitchellh/hashstructure/v2"
	"github.com/moby/buildkit/cache/remotecache"
	"github.com/moby/buildkit/client"
	"github.com/moby/buildkit/client/llb"
	"github.com/moby/buildkit/exporter/containerimage/exptypes"
	"github.com/moby/buildkit/frontend"
	gw "github.com/moby/buildkit/frontend/gateway/client"
	"github.com/moby/buildkit/identity"
	"github.com/moby/buildkit/session"
	"github.com/moby/buildkit/solver"
	"github.com/moby/buildkit/solver/errdefs"
	llberrdefs "github.com/moby/buildkit/solver/llbsolver/errdefs"
	"github.com/moby/buildkit/solver/pb"
	"github.com/moby/buildkit/util/bklog"
	"github.com/moby/buildkit/util/buildinfo"
	"github.com/moby/buildkit/util/flightcontrol"
	"github.com/moby/buildkit/util/progress"
	"github.com/moby/buildkit/worker"
	digest "github.com/opencontainers/go-digest"
	"github.com/pkg/errors"
)

type llbBridge struct {
	builder                   solver.Builder
	frontends                 map[string]frontend.Frontend
	resolveWorker             func() (worker.Worker, error)
	eachWorker                func(func(worker.Worker) error) error
	resolveCacheImporterFuncs map[string]remotecache.ResolveCacheImporterFunc
	cms                       map[string]solver.CacheManager
	cmsMu                     sync.Mutex
	sm                        *session.Manager
}

func (b *llbBridge) Warn(ctx context.Context, dgst digest.Digest, msg string, opts frontend.WarnOpts) error {
	return b.builder.InContext(ctx, func(ctx context.Context, g session.Group) error {
		pw, ok, _ := progress.NewFromContext(ctx, progress.WithMetadata("vertex", dgst))
		if !ok {
			return nil
		}
		level := opts.Level
		if level == 0 {
			level = 1
		}
		pw.Write(identity.NewID(), client.VertexWarning{
			Vertex:     dgst,
			Level:      level,
			Short:      []byte(msg),
			SourceInfo: opts.SourceInfo,
			Range:      opts.Range,
			Detail:     opts.Detail,
			URL:        opts.URL,
		})
		return pw.Close()
	})
}

func (b *llbBridge) loadResult(ctx context.Context, def *pb.Definition, cacheImports []gw.CacheOptionsEntry) (solver.CachedResult, solver.BuildSources, error) {
	w, err := b.resolveWorker()
	if err != nil {
		return nil, nil, err
	}
	ent, err := loadEntitlements(b.builder)
	if err != nil {
		return nil, nil, err
	}
	var cms []solver.CacheManager
	for _, im := range cacheImports {
		cmID, err := cmKey(im)
		if err != nil {
			return nil, nil, err
		}
		b.cmsMu.Lock()
		var cm solver.CacheManager
		if prevCm, ok := b.cms[cmID]; !ok {
			func(cmID string, im gw.CacheOptionsEntry) {
				cm = newLazyCacheManager(cmID, func() (solver.CacheManager, error) {
					var cmNew solver.CacheManager
					if err := inBuilderContext(context.TODO(), b.builder, "importing cache manifest from "+cmID, "", func(ctx context.Context, g session.Group) error {
						resolveCI, ok := b.resolveCacheImporterFuncs[im.Type]
						if !ok {
							return errors.Errorf("unknown cache importer: %s", im.Type)
						}
						ci, desc, err := resolveCI(ctx, g, im.Attrs)
						if err != nil {
							return err
						}
						cmNew, err = ci.Resolve(ctx, desc, cmID, w)
						return err
					}); err != nil {
						bklog.G(ctx).Debugf("error while importing cache manifest from cmId=%s: %v", cmID, err)
						return nil, err
					}
					return cmNew, nil
				})
			}(cmID, im)
			b.cms[cmID] = cm
		} else {
			cm = prevCm
		}
		cms = append(cms, cm)
		b.cmsMu.Unlock()
	}
	dpc := &detectPrunedCacheID{}

	edge, err := Load(def, dpc.Load, ValidateEntitlements(ent), WithCacheSources(cms), NormalizeRuntimePlatforms(), WithValidateCaps())
	if err != nil {
		return nil, nil, errors.Wrap(err, "failed to load LLB")
	}

	if len(dpc.ids) > 0 {
		ids := make([]string, 0, len(dpc.ids))
		for id := range dpc.ids {
			ids = append(ids, id)
		}
		if err := b.eachWorker(func(w worker.Worker) error {
			return w.PruneCacheMounts(ctx, ids)
		}); err != nil {
			return nil, nil, err
		}
	}

	res, bi, err := b.builder.Build(ctx, edge)
	if err != nil {
		return nil, nil, err
	}
	return res, bi, nil
}

func (b *llbBridge) Solve(ctx context.Context, req frontend.SolveRequest, sid string) (res *frontend.Result, err error) {
	if req.Definition != nil && req.Definition.Def != nil && req.Frontend != "" {
		return nil, errors.New("cannot solve with both Definition and Frontend specified")
	}

	if req.Definition != nil && req.Definition.Def != nil {
		res = &frontend.Result{Ref: newResultProxy(b, req)}
		if req.Evaluate {
			_, err = res.Ref.Result(ctx)
		}
	} else if req.Frontend != "" {
		f, ok := b.frontends[req.Frontend]
		if !ok {
			return nil, errors.Errorf("invalid frontend: %s", req.Frontend)
		}
		res, err = f.Solve(ctx, b, req.FrontendOpt, req.FrontendInputs, sid, b.sm)
		if err != nil {
			return nil, err
		}
	} else {
		return &frontend.Result{}, nil
	}

	if res.Metadata == nil {
		res.Metadata = make(map[string][]byte)
	}

	if len(res.Refs) > 0 {
		for p := range res.Refs {
			dtbi, errm := buildinfo.GetMetadata(res.Metadata, fmt.Sprintf("%s/%s", exptypes.ExporterBuildInfo, p), req.Frontend, req.FrontendOpt)
			if errm != nil {
				return nil, err
			}
			res.Metadata[fmt.Sprintf("%s/%s", exptypes.ExporterBuildInfo, p)] = dtbi
		}
	} else {
		dtbi, errm := buildinfo.GetMetadata(res.Metadata, exptypes.ExporterBuildInfo, req.Frontend, req.FrontendOpt)
		if errm != nil {
			return nil, err
		}
		res.Metadata[exptypes.ExporterBuildInfo] = dtbi
	}

	return
}

type resultProxy struct {
	cb         func(context.Context) (solver.CachedResult, solver.BuildSources, error)
	def        *pb.Definition
	g          flightcontrol.Group
	mu         sync.Mutex
	released   bool
	v          solver.CachedResult
	bsrc       solver.BuildSources
	err        error
	errResults []solver.Result
}

func newResultProxy(b *llbBridge, req frontend.SolveRequest) *resultProxy {
	rp := &resultProxy{
		def: req.Definition,
	}
	rp.cb = func(ctx context.Context) (solver.CachedResult, solver.BuildSources, error) {
		res, bsrc, err := b.loadResult(ctx, req.Definition, req.CacheImports)
		var ee *llberrdefs.ExecError
		if errors.As(err, &ee) {
			ee.EachRef(func(res solver.Result) error {
				rp.errResults = append(rp.errResults, res)
				return nil
			})
			// acquire ownership so ExecError finalizer doesn't attempt to release as well
			ee.OwnerBorrowed = true
		}
		return res, bsrc, err
	}
	return rp
}

func (rp *resultProxy) Definition() *pb.Definition {
	return rp.def
}

func (rp *resultProxy) BuildSources() solver.BuildSources {
	return rp.bsrc
}

func (rp *resultProxy) Release(ctx context.Context) (err error) {
	rp.mu.Lock()
	defer rp.mu.Unlock()
	for _, res := range rp.errResults {
		rerr := res.Release(ctx)
		if rerr != nil {
			err = rerr
		}
	}
	if rp.v != nil {
		if rp.released {
			bklog.G(ctx).Warnf("release of already released result")
		}
		rerr := rp.v.Release(ctx)
		if err != nil {
			return rerr
		}
	}
	rp.released = true
	return
}

func (rp *resultProxy) wrapError(err error) error {
	if err == nil {
		return nil
	}
	var ve *errdefs.VertexError
	if errors.As(err, &ve) {
		if rp.def.Source != nil {
			locs, ok := rp.def.Source.Locations[string(ve.Digest)]
			if ok {
				for _, loc := range locs.Locations {
					err = errdefs.WithSource(err, errdefs.Source{
						Info:   rp.def.Source.Infos[loc.SourceIndex],
						Ranges: loc.Ranges,
					})
				}
			}
		}
	}
	return err
}

func (rp *resultProxy) Result(ctx context.Context) (res solver.CachedResult, err error) {
	defer func() {
		err = rp.wrapError(err)
	}()
	r, err := rp.g.Do(ctx, "result", func(ctx context.Context) (interface{}, error) {
		rp.mu.Lock()
		if rp.released {
			rp.mu.Unlock()
			return nil, errors.Errorf("accessing released result")
		}
		if rp.v != nil || rp.err != nil {
			rp.mu.Unlock()
			return rp.v, rp.err
		}
		rp.mu.Unlock()
		v, bsrc, err := rp.cb(ctx)
		if err != nil {
			select {
			case <-ctx.Done():
				if strings.Contains(err.Error(), context.Canceled.Error()) {
					return v, err
				}
			default:
			}
		}
		rp.mu.Lock()
		if rp.released {
			if v != nil {
				v.Release(context.TODO())
			}
			rp.mu.Unlock()
			return nil, errors.Errorf("evaluating released result")
		}
		rp.v = v
		rp.bsrc = bsrc
		rp.err = err
		rp.mu.Unlock()
		return v, err
	})
	if r != nil {
		return r.(solver.CachedResult), nil
	}
	return nil, err
}

func (b *llbBridge) ResolveImageConfig(ctx context.Context, ref string, opt llb.ResolveImageConfigOpt) (dgst digest.Digest, config []byte, err error) {
	w, err := b.resolveWorker()
	if err != nil {
		return "", nil, err
	}
	if opt.LogName == "" {
		opt.LogName = fmt.Sprintf("resolve image config for %s", ref)
	}
	id := ref // make a deterministic ID for avoiding duplicates
	if platform := opt.Platform; platform == nil {
		id += platforms.Format(platforms.DefaultSpec())
	} else {
		id += platforms.Format(*platform)
	}
	err = inBuilderContext(ctx, b.builder, opt.LogName, id, func(ctx context.Context, g session.Group) error {
		dgst, config, err = w.ResolveImageConfig(ctx, ref, opt, b.sm, g)
		return err
	})
	return dgst, config, err
}

type lazyCacheManager struct {
	id   string
	main solver.CacheManager

	waitCh chan struct{}
	err    error
}

func (lcm *lazyCacheManager) ID() string {
	return lcm.id
}
func (lcm *lazyCacheManager) Query(inp []solver.CacheKeyWithSelector, inputIndex solver.Index, dgst digest.Digest, outputIndex solver.Index) ([]*solver.CacheKey, error) {
	lcm.wait()
	if lcm.main == nil {
		return nil, nil
	}
	return lcm.main.Query(inp, inputIndex, dgst, outputIndex)
}
func (lcm *lazyCacheManager) Records(ck *solver.CacheKey) ([]*solver.CacheRecord, error) {
	lcm.wait()
	if lcm.main == nil {
		return nil, nil
	}
	return lcm.main.Records(ck)
}
func (lcm *lazyCacheManager) Load(ctx context.Context, rec *solver.CacheRecord) (solver.Result, error) {
	if err := lcm.wait(); err != nil {
		return nil, err
	}
	return lcm.main.Load(ctx, rec)
}
func (lcm *lazyCacheManager) Save(key *solver.CacheKey, s solver.Result, createdAt time.Time) (*solver.ExportableCacheKey, error) {
	if err := lcm.wait(); err != nil {
		return nil, err
	}
	return lcm.main.Save(key, s, createdAt)
}

func (lcm *lazyCacheManager) wait() error {
	<-lcm.waitCh
	return lcm.err
}

func newLazyCacheManager(id string, fn func() (solver.CacheManager, error)) solver.CacheManager {
	lcm := &lazyCacheManager{id: id, waitCh: make(chan struct{})}
	go func() {
		defer close(lcm.waitCh)
		cm, err := fn()
		if err != nil {
			lcm.err = err
			return
		}
		lcm.main = cm
	}()
	return lcm
}

func cmKey(im gw.CacheOptionsEntry) (string, error) {
	if im.Type == "registry" && im.Attrs["ref"] != "" {
		return im.Attrs["ref"], nil
	}
	i, err := hashstructure.Hash(im, hashstructure.FormatV2, nil)
	if err != nil {
		return "", err
	}
	return fmt.Sprintf("%s:%d", im.Type, i), nil
}
<commit_msg>fix nil pointer dereference in bridge result<commit_after>package llbsolver

import (
	"context"
	"fmt"
	"strings"
	"sync"
	"time"

	"github.com/containerd/containerd/platforms"
	"github.com/mitchellh/hashstructure/v2"
	"github.com/moby/buildkit/cache/remotecache"
	"github.com/moby/buildkit/client"
	"github.com/moby/buildkit/client/llb"
	"github.com/moby/buildkit/exporter/containerimage/exptypes"
	"github.com/moby/buildkit/frontend"
	gw "github.com/moby/buildkit/frontend/gateway/client"
	"github.com/moby/buildkit/identity"
	"github.com/moby/buildkit/session"
	"github.com/moby/buildkit/solver"
	"github.com/moby/buildkit/solver/errdefs"
	llberrdefs "github.com/moby/buildkit/solver/llbsolver/errdefs"
	"github.com/moby/buildkit/solver/pb"
	"github.com/moby/buildkit/util/bklog"
	"github.com/moby/buildkit/util/buildinfo"
	"github.com/moby/buildkit/util/flightcontrol"
	"github.com/moby/buildkit/util/progress"
	"github.com/moby/buildkit/worker"
	digest "github.com/opencontainers/go-digest"
	"github.com/pkg/errors"
)

type llbBridge struct {
	builder                   solver.Builder
	frontends                 map[string]frontend.Frontend
	resolveWorker             func() (worker.Worker, error)
	eachWorker                func(func(worker.Worker) error) error
	resolveCacheImporterFuncs map[string]remotecache.ResolveCacheImporterFunc
	cms                       map[string]solver.CacheManager
	cmsMu                     sync.Mutex
	sm                        *session.Manager
}

func (b *llbBridge) Warn(ctx context.Context, dgst digest.Digest, msg string, opts frontend.WarnOpts) error {
	return b.builder.InContext(ctx, func(ctx context.Context, g session.Group) error {
		pw, ok, _ := progress.NewFromContext(ctx, progress.WithMetadata("vertex", dgst))
		if !ok {
			return nil
		}
		level := opts.Level
		if level == 0 {
			level = 1
		}
		pw.Write(identity.NewID(), client.VertexWarning{
			Vertex:     dgst,
			Level:      level,
			Short:      []byte(msg),
			SourceInfo: opts.SourceInfo,
			Range:      opts.Range,
			Detail:     opts.Detail,
			URL:        opts.URL,
		})
		return pw.Close()
	})
}

func (b *llbBridge) loadResult(ctx context.Context, def *pb.Definition, cacheImports []gw.CacheOptionsEntry) (solver.CachedResult, solver.BuildSources, error) {
	w, err := b.resolveWorker()
	if err != nil {
		return nil, nil, err
	}
	ent, err := loadEntitlements(b.builder)
	if err != nil {
		return nil, nil, err
	}
	var cms []solver.CacheManager
	for _, im := range cacheImports {
		cmID, err := cmKey(im)
		if err != nil {
			return nil, nil, err
		}
		b.cmsMu.Lock()
		var cm solver.CacheManager
		if prevCm, ok := b.cms[cmID]; !ok {
			func(cmID string, im gw.CacheOptionsEntry) {
				cm = newLazyCacheManager(cmID, func() (solver.CacheManager, error) {
					var cmNew solver.CacheManager
					if err := inBuilderContext(context.TODO(), b.builder, "importing cache manifest from "+cmID, "", func(ctx context.Context, g session.Group) error {
						resolveCI, ok := b.resolveCacheImporterFuncs[im.Type]
						if !ok {
							return errors.Errorf("unknown cache importer: %s", im.Type)
						}
						ci, desc, err := resolveCI(ctx, g, im.Attrs)
						if err != nil {
							return err
						}
						cmNew, err = ci.Resolve(ctx, desc, cmID, w)
						return err
					}); err != nil {
						bklog.G(ctx).Debugf("error while importing cache manifest from cmId=%s: %v", cmID, err)
						return nil, err
					}
					return cmNew, nil
				})
			}(cmID, im)
			b.cms[cmID] = cm
		} else {
			cm = prevCm
		}
		cms = append(cms, cm)
		b.cmsMu.Unlock()
	}
	dpc := &detectPrunedCacheID{}

	edge, err := Load(def, dpc.Load, ValidateEntitlements(ent), WithCacheSources(cms), NormalizeRuntimePlatforms(), WithValidateCaps())
	if err != nil {
		return nil, nil, errors.Wrap(err, "failed to load LLB")
	}

	if len(dpc.ids) > 0 {
		ids := make([]string, 0, len(dpc.ids))
		for id := range dpc.ids {
			ids = append(ids, id)
		}
		if err := b.eachWorker(func(w worker.Worker) error {
			return w.PruneCacheMounts(ctx, ids)
		}); err != nil {
			return nil, nil, err
		}
	}

	res, bi, err := b.builder.Build(ctx, edge)
	if err != nil {
		return nil, nil, err
	}
	return res, bi, nil
}

func (b *llbBridge) Solve(ctx context.Context, req frontend.SolveRequest, sid string) (res *frontend.Result, err error) {
	if req.Definition != nil && req.Definition.Def != nil && req.Frontend != "" {
		return nil, errors.New("cannot solve with both Definition and Frontend specified")
	}

	if req.Definition != nil && req.Definition.Def != nil {
		res = &frontend.Result{Ref: newResultProxy(b, req)}
		if req.Evaluate {
			_, err = res.Ref.Result(ctx)
		}
	} else if req.Frontend != "" {
		f, ok := b.frontends[req.Frontend]
		if !ok {
			return nil, errors.Errorf("invalid frontend: %s", req.Frontend)
		}
		res, err = f.Solve(ctx, b, req.FrontendOpt, req.FrontendInputs, sid, b.sm)
		if err != nil {
			return nil, err
		}
	} else {
		return &frontend.Result{}, nil
	}

	if res.Metadata == nil {
		res.Metadata = make(map[string][]byte)
	}

	if len(res.Refs) > 0 {
		for p := range res.Refs {
			dtbi, err := buildinfo.GetMetadata(res.Metadata, fmt.Sprintf("%s/%s", exptypes.ExporterBuildInfo, p), req.Frontend, req.FrontendOpt)
			if err != nil {
				return nil, err
			}
			res.Metadata[fmt.Sprintf("%s/%s", exptypes.ExporterBuildInfo, p)] = dtbi
		}
	} else {
		dtbi, err := buildinfo.GetMetadata(res.Metadata, exptypes.ExporterBuildInfo, req.Frontend, req.FrontendOpt)
		if err != nil {
			return nil, err
		}
		res.Metadata[exptypes.ExporterBuildInfo] = dtbi
	}

	return
}

type resultProxy struct {
	cb         func(context.Context) (solver.CachedResult, solver.BuildSources, error)
	def        *pb.Definition
	g          flightcontrol.Group
	mu         sync.Mutex
	released   bool
	v          solver.CachedResult
	bsrc       solver.BuildSources
	err        error
	errResults []solver.Result
}

func newResultProxy(b *llbBridge, req frontend.SolveRequest) *resultProxy {
	rp := &resultProxy{
		def: req.Definition,
	}
	rp.cb = func(ctx context.Context) (solver.CachedResult, solver.BuildSources, error) {
		res, bsrc, err := b.loadResult(ctx, req.Definition, req.CacheImports)
		var ee *llberrdefs.ExecError
		if errors.As(err, &ee) {
			ee.EachRef(func(res solver.Result) error {
				rp.errResults = append(rp.errResults, res)
				return nil
			})
			// acquire ownership so ExecError finalizer doesn't attempt to release as well
			ee.OwnerBorrowed = true
		}
		return res, bsrc, err
	}
	return rp
}

func (rp *resultProxy) Definition() *pb.Definition {
	return rp.def
}

func (rp *resultProxy) BuildSources() solver.BuildSources {
	return rp.bsrc
}

func (rp *resultProxy) Release(ctx context.Context) (err error) {
	rp.mu.Lock()
	defer rp.mu.Unlock()
	for _, res := range rp.errResults {
		rerr := res.Release(ctx)
		if rerr != nil {
			err = rerr
		}
	}
	if rp.v != nil {
		if rp.released {
			bklog.G(ctx).Warnf("release of already released result")
		}
		rerr := rp.v.Release(ctx)
		if err != nil {
			return rerr
		}
	}
	rp.released = true
	return
}

func (rp *resultProxy) wrapError(err error) error {
	if err == nil {
		return nil
	}
	var ve *errdefs.VertexError
	if errors.As(err, &ve) {
		if rp.def.Source != nil {
			locs, ok := rp.def.Source.Locations[string(ve.Digest)]
			if ok {
				for _, loc := range locs.Locations {
					err = errdefs.WithSource(err, errdefs.Source{
						Info:   rp.def.Source.Infos[loc.SourceIndex],
						Ranges: loc.Ranges,
					})
				}
			}
		}
	}
	return err
}

func (rp *resultProxy) Result(ctx context.Context) (res solver.CachedResult, err error) {
	defer func() {
		err = rp.wrapError(err)
	}()
	r, err := rp.g.Do(ctx, "result", func(ctx context.Context) (interface{}, error) {
		rp.mu.Lock()
		if rp.released {
			rp.mu.Unlock()
			return nil, errors.Errorf("accessing released result")
		}
		if rp.v != nil || rp.err != nil {
			rp.mu.Unlock()
			return rp.v, rp.err
		}
		rp.mu.Unlock()
		v, bsrc, err := rp.cb(ctx)
		if err != nil {
			select {
			case <-ctx.Done():
				if strings.Contains(err.Error(), context.Canceled.Error()) {
					return v, err
				}
			default:
			}
		}
		rp.mu.Lock()
		if rp.released {
			if v != nil {
				v.Release(context.TODO())
			}
			rp.mu.Unlock()
			return nil, errors.Errorf("evaluating released result")
		}
		rp.v = v
		rp.bsrc = bsrc
		rp.err = err
		rp.mu.Unlock()
		return v, err
	})
	if r != nil {
		return r.(solver.CachedResult), nil
	}
	return nil, err
}

func (b *llbBridge) ResolveImageConfig(ctx context.Context, ref string, opt llb.ResolveImageConfigOpt) (dgst digest.Digest, config []byte, err error) {
	w, err := b.resolveWorker()
	if err != nil {
		return "", nil, err
	}
	if opt.LogName == "" {
		opt.LogName = fmt.Sprintf("resolve image config for %s", ref)
	}
	id := ref // make a deterministic ID for avoiding duplicates
	if platform := opt.Platform; platform == nil {
		id += platforms.Format(platforms.DefaultSpec())
	} else {
		id += platforms.Format(*platform)
	}
	err = inBuilderContext(ctx, b.builder, opt.LogName, id, func(ctx context.Context, g session.Group) error {
		dgst, config, err = w.ResolveImageConfig(ctx, ref, opt, b.sm, g)
		return err
	})
	return dgst, config, err
}

type lazyCacheManager struct {
	id   string
	main solver.CacheManager

	waitCh chan struct{}
	err    error
}

func (lcm *lazyCacheManager) ID() string {
	return lcm.id
}
func (lcm *lazyCacheManager) Query(inp []solver.CacheKeyWithSelector, inputIndex solver.Index, dgst digest.Digest, outputIndex solver.Index) ([]*solver.CacheKey, error) {
	lcm.wait()
	if lcm.main == nil {
		return nil, nil
	}
	return lcm.main.Query(inp, inputIndex, dgst, outputIndex)
}
func (lcm *lazyCacheManager) Records(ck *solver.CacheKey) ([]*solver.CacheRecord, error) {
	lcm.wait()
	if lcm.main == nil {
		return nil, nil
	}
	return lcm.main.Records(ck)
}
func (lcm *lazyCacheManager) Load(ctx context.Context, rec *solver.CacheRecord) (solver.Result, error) {
	if err := lcm.wait(); err != nil {
		return nil, err
	}
	return lcm.main.Load(ctx, rec)
}
func (lcm *lazyCacheManager) Save(key *solver.CacheKey, s solver.Result, createdAt time.Time) (*solver.ExportableCacheKey, error) {
	if err := lcm.wait(); err != nil {
		return nil, err
	}
	return lcm.main.Save(key, s, createdAt)
}

func (lcm *lazyCacheManager) wait() error {
	<-lcm.waitCh
	return lcm.err
}

func newLazyCacheManager(id string, fn func() (solver.CacheManager, error)) solver.CacheManager {
	lcm := &lazyCacheManager{id: id, waitCh: make(chan struct{})}
	go func() {
		defer close(lcm.waitCh)
		cm, err := fn()
		if err != nil {
			lcm.err = err
			return
		}
		lcm.main = cm
	}()
	return lcm
}

func cmKey(im gw.CacheOptionsEntry) (string, error) {
	if im.Type == "registry" && im.Attrs["ref"] != "" {
		return im.Attrs["ref"], nil
	}
	i, err := hashstructure.Hash(im, hashstructure.FormatV2, nil)
	if err != nil {
		return "", err
	}
	return fmt.Sprintf("%s:%d", im.Type, i), nil
}
<|endoftext|>


<commit_before>package main

import (
	"database/sql"
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"strings"

	"github.com/satori/go.uuid"
)

type guidePost struct {
	input     chan treeRequest
	shutdown  chan bool
	conn      *sql.DB
	jbsv_stmt *sql.Stmt
	repo_stmt *sql.Stmt
	name_stmt *sql.Stmt
	node_stmt *sql.Stmt
	serv_stmt *sql.Stmt
	attr_stmt *sql.Stmt
}

func (g *guidePost) run() {
	var err error

	log.Println("Prepare: guide/job-save")
	g.jbsv_stmt, err = g.conn.Prepare(`
INSERT INTO soma.jobs (
	job_id,
	job_status,
	job_result,
	job_type,
	repository_id,
	user_id,
	organizational_team_id,
	job)
SELECT	$1::uuid,
		$2::varchar,
		$3::varchar,
		$4::varchar,
		$5::uuid,
		$6::uuid,
		$7::uuid,
		$8::jsonb;`)
	if err != nil {
		log.Fatal("guide/job-save: ", err)
	}
	defer g.jbsv_stmt.Close()

	log.Println("Prepare: guide/repo-by-bucket")
	g.repo_stmt, err = g.conn.Prepare(`
SELECT	sb.repository_id,
		sr.repository_name
FROM	soma.buckets sb
JOIN    soma.repositories sr
ON		sb.repository_id = sr.repository_id
WHERE	sb.bucket_id = $1::uuid;`)
	if err != nil {
		log.Fatal("guide/repo-by-bucket: ", err)
	}
	defer g.repo_stmt.Close()

	log.Println("Prepare: guide/load-node-details")
	g.node_stmt, err = g.conn.Prepare(`
SELECT    sn.node_asset_id,
	      sn.node_name,
	      sn.organizational_team_id,
	      sn.server_id,
	      sn.node_online,
	      sn.node_deleted
FROM      soma.nodes sn
LEFT JOIN soma.node_bucket_assignment snba
ON        sn.node_id = snba.node_id
WHERE     sn.node_online = 'yes'
AND       sn.node_deleted = 'false'
AND       snba.node_id IS NULL
AND       sn.node_id = $1::uuid;`)
	if err != nil {
		log.Fatal("guide/load-node-details: ", err)
	}
	defer g.node_stmt.Close()

	log.Println("Prepare: guide/repo-by-id")
	g.name_stmt, err = g.conn.Prepare(`
SELECT repository_name
FROM   soma.repositories
WHERE  repository_id = $1::uuid;`)
	if err != nil {
		log.Fatal("guide/repo-by-id: ", err)
	}
	defer g.name_stmt.Close()

	log.Println("Prepare: guide/service-lookup")
	g.serv_stmt, err = g.conn.Prepare(`
SELECT stsp.service_property
FROM   soma.repositories sr
JOIN   soma.team_service_properties stsp
ON     sr.organizational_team_id = stsp.organizational_team_id
WHERE  sr.repository_id = $1::uuid
AND    stsp.service_property = $2::varchar
AND    sr.organizational_team_id = $3::uuid;`)
	if err != nil {
		log.Fatal("guide/service-lookup: ", err)
	}
	defer g.serv_stmt.Close()

	log.Println("Prepare: guide/populate-service-attributes")
	g.attr_stmt, err = g.conn.Prepare(`
SELECT stspv.service_property_attribute,
       stspv.value
FROM   soma.repositories sr
JOIN   soma.team_service_properties stsp
ON     sr.organizational_team_id = stsp.organizational_team_id
JOIN   soma.team_service_property_values stspv
ON     stsp.organizational_team_id = stspv.organizational_team_id
AND    stsp.service_property = stspv.service_property
WHERE  sr.repository_id = $1::uuid
AND    stsp.service_property = $2::varchar
AND    sr.organizational_team_id = $3::uuid;`)
	if err != nil {
		log.Fatal("guide/populate-service-attributes: ", err)
	}
	defer g.attr_stmt.Close()

runloop:
	for {
		select {
		case <-g.shutdown:
			break runloop
		case req := <-g.input:
			g.process(&req)
		}
	}
}

func (g *guidePost) process(q *treeRequest) {
	var (
		res                                sql.Result
		err                                error
		j                                  []byte
		repoId, repoName, keeper, bucketId string
		ndName, ndTeam, ndServer           string
		ndAsset                            int64
		ndOnline, ndDeleted                bool
	)
	result := somaResult{}

	switch q.Action {
	case "add_system_property_to_repository":
		fallthrough
	case "add_custom_property_to_repository":
		fallthrough
	case "add_oncall_property_to_repository":
		fallthrough
	case "add_service_property_to_repository":
		repoId = q.Repository.Repository.Id

	case "create_bucket":
		repoId = q.Bucket.Bucket.Repository
	case "add_system_property_to_bucket":
		fallthrough
	case "add_custom_property_to_bucket":
		fallthrough
	case "add_oncall_property_to_bucket":
		fallthrough
	case "add_service_property_to_bucket":
		bucketId = q.Bucket.Bucket.Id

	case "create_group":
		fallthrough
	case "add_group_to_group":
		fallthrough
	case "add_cluster_to_group":
		fallthrough
	case "add_node_to_group":
		fallthrough
	case "add_system_property_to_group":
		fallthrough
	case "add_custom_property_to_group":
		fallthrough
	case "add_oncall_property_to_group":
		fallthrough
	case "add_service_property_to_group":
		bucketId = q.Group.Group.BucketId

	case "add_node_to_cluster":
		if q.Cluster.Cluster.BucketId != q.Cluster.Cluster.Members[0].Config.BucketId {
			panic("This should not happen.")
		}
		fallthrough
	case "create_cluster":
		fallthrough
	case "add_system_property_to_cluster":
		fallthrough
	case "add_custom_property_to_cluster":
		fallthrough
	case "add_oncall_property_to_cluster":
		fallthrough
	case "add_service_property_to_cluster":
		bucketId = q.Cluster.Cluster.BucketId

	case "assign_node":
		if err = g.node_stmt.QueryRow(q.Node.Node.Id).Scan(
			&ndAsset,
			&ndName,
			&ndTeam,
			&ndServer,
			&ndOnline,
			&ndDeleted,
		); err != nil {
			if err == sql.ErrNoRows {
				result.SetNotFound()
			} else {
				_ = result.SetRequestError(err)
			}
			q.reply <- result
			return
		}

		q.Node.Node.AssetId = uint64(ndAsset)
		q.Node.Node.Name = ndName
		q.Node.Node.Team = ndTeam
		q.Node.Node.Server = ndServer
		q.Node.Node.IsOnline = ndOnline
		q.Node.Node.IsDeleted = ndDeleted
		fallthrough
	case "add_system_property_to_node":
		fallthrough
	case "add_custom_property_to_node":
		fallthrough
	case "add_oncall_property_to_node":
		fallthrough
	case "add_service_property_to_node":
		repoId = q.Node.Node.Config.RepositoryId
		bucketId = q.Node.Node.Config.BucketId

	default:
		log.Printf("R: unimplemented guidepost/%s", q.Action)
		result.SetNotImplemented()
		q.reply <- result
		return
	}

	// lookup repository by bucket
	if bucketId != "" {
		if err = g.repo_stmt.QueryRow(bucketId).Scan(&repoId, &repoName); err != nil {
			if err == sql.ErrNoRows {
				result.SetNotFound()
			} else {
				_ = result.SetRequestError(err)
			}
			q.reply <- result
			return
		}
	}

	// lookup repository name
	if repoName == "" && repoId != "" {
		if err = g.name_stmt.QueryRow(repoId).Scan(&repoName); err != nil {
			if err == sql.ErrNoRows {
				result.SetNotFound()
			} else {
				_ = result.SetRequestError(err)
			}
			q.reply <- result
			return
		}
	}

	// XXX
	if repoName == "" {
		panic(`Have no repository name`)
	}

	// check we have a treekeeper for that repository
	keeper = fmt.Sprintf("repository_%s", repoName)
	if _, ok := handlerMap[keeper].(*treeKeeper); !ok {
		_ = result.SetRequestError(
			fmt.Errorf("No handler for repository %s registered.\n", repoName),
		)
		q.reply <- result
		return
	}

	// check the treekeeper has finished loading
	handler := handlerMap[keeper].(*treeKeeper)
	if !handler.isReady() {
		_ = result.SetRequestError( // TODO should be 503/ServiceUnavailable
			fmt.Errorf("Repository %s not fully loaded yet.\n", repoName),
		)
		q.reply <- result
		return
	}

	// check the treekeeper has not encountered a broken tree
	if handler.isBroken() {
		_ = result.SetRequestError(
			fmt.Errorf("Repository %s is broken.\n", repoName),
		)
		q.reply <- result
		return
	}

	// load authoritative copy of the service attributes from the
	// database. Replaces whatever the client sent in.
	if strings.Contains(q.Action, "add_service_property_to_") {
		var service, attr, val, svName, svTeam string
		var rows *sql.Rows
		attrs := []somaproto.TreeServiceAttribute{}

		switch q.RequestType {
		case "repository":
			svName = (*q.Repository.Repository.Properties)[0].Service.Name
			svTeam = (*q.Repository.Repository.Properties)[0].Service.TeamId
		case "bucket":
			svName = (*q.Bucket.Bucket.Properties)[0].Service.Name
			svTeam = (*q.Bucket.Bucket.Properties)[0].Service.TeamId
		case "group":
			svName = (*q.Group.Group.Properties)[0].Service.Name
			svTeam = (*q.Group.Group.Properties)[0].Service.TeamId
		case "cluster":
			svName = (*q.Cluster.Cluster.Properties)[0].Service.Name
			svTeam = (*q.Cluster.Cluster.Properties)[0].Service.TeamId
		case "node":
			svName = (*q.Node.Node.Properties)[0].Service.Name
			svTeam = (*q.Node.Node.Properties)[0].Service.TeamId
		}

		if err = g.serv_stmt.QueryRow(repoId, svName, svTeam).Scan(&service); err != nil {
			goto inputabort
		}

		if rows, err = g.attr_stmt.Query(repoId, svName, svTeam); err != nil {
			goto inputabort
		}
		defer rows.Close()

	attrloop:
		for rows.Next() {
			if err = rows.Scan(&attr, &val); err != nil {
				break attrloop
			}
			attrs = append(attrs, somaproto.TreeServiceAttribute{
				Attribute: attr,
				Value:     val,
			})
		}

	inputabort:
		if err != nil {
			if err == sql.ErrNoRows {
				result.SetNotFound()
			} else {
				_ = result.SetRequestError(err)
			}
			q.reply <- result
			return
		}
		// not aborted: set the loaded attributes
		switch q.RequestType {
		case "repository":
			(*q.Repository.Repository.Properties)[0].Service.Attributes = attrs
		case "bucket":
			(*q.Bucket.Bucket.Properties)[0].Service.Attributes = attrs
		case "group":
			(*q.Group.Group.Properties)[0].Service.Attributes = attrs
		case "cluster":
			(*q.Cluster.Cluster.Properties)[0].Service.Attributes = attrs
		case "node":
			(*q.Node.Node.Properties)[0].Service.Attributes = attrs
		}
	}

	// store job in database
	log.Printf("R: jobsave/%s", q.Action)
	q.JobId = uuid.NewV4()
	j, _ = json.Marshal(q)
	res, err = g.jbsv_stmt.Exec(
		q.JobId.String(),
		"queued",
		"pending",
		q.Action,
		repoId,
		"00000000-0000-0000-0000-000000000000", // XXX user uuid
		"00000000-0000-0000-0000-000000000000", // XXX team uuid
		string(j),
	)
	if result.SetRequestError(err) {
		q.reply <- result
		return
	}
	rowCnt, _ := res.RowsAffected()
	switch {
	case rowCnt == 0:
		err = errors.New("No rows affected")
	case rowCnt > 1:
		err = fmt.Errorf("Too many rows affected: %d", rowCnt)
	case rowCnt < 0:
		err = fmt.Errorf("Space/Time Continuum broke, rows affected: %d", rowCnt)
	}
	if err != nil {
		switch q.RequestType {
		case "repository":
			result.Append(err, &somaRepositoryResult{})
		case "bucket":
			result.Append(err, &somaBucketResult{})
		case "group":
			result.Append(err, &somaGroupResult{})
		case "cluster":
			result.Append(err, &somaClusterResult{})
		case "node":
			result.Append(err, &somaNodeResult{})
		}
		q.reply <- result
		return
	}

	handler.input <- *q
	result.JobId = q.JobId.String()

	switch q.RequestType {
	case "repository":
		result.Append(nil, &somaRepositoryResult{
			Repository: q.Repository.Repository,
		})
	case "bucket":
		result.Append(nil, &somaBucketResult{
			Bucket: q.Bucket.Bucket,
		})
	case "group":
		result.Append(nil, &somaGroupResult{
			Group: q.Group.Group,
		})
	case "cluster":
		result.Append(nil, &somaClusterResult{
			Cluster: q.Cluster.Cluster,
		})
	case "node":
		result.Append(nil, &somaNodeResult{
			Node: q.Node.Node,
		})
	}
	q.reply <- result
}

// vim: ts=4 sw=4 sts=4 noet fenc=utf-8 ffs=unix
<commit_msg>Add CheckConfig support to GuidePost<commit_after>package main

import (
	"database/sql"
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"strings"

	"github.com/satori/go.uuid"
)

type guidePost struct {
	input     chan treeRequest
	shutdown  chan bool
	conn      *sql.DB
	jbsv_stmt *sql.Stmt
	repo_stmt *sql.Stmt
	name_stmt *sql.Stmt
	node_stmt *sql.Stmt
	serv_stmt *sql.Stmt
	attr_stmt *sql.Stmt
}

func (g *guidePost) run() {
	var err error

	log.Println("Prepare: guide/job-save")
	g.jbsv_stmt, err = g.conn.Prepare(`
INSERT INTO soma.jobs (
	job_id,
	job_status,
	job_result,
	job_type,
	repository_id,
	user_id,
	organizational_team_id,
	job)
SELECT	$1::uuid,
		$2::varchar,
		$3::varchar,
		$4::varchar,
		$5::uuid,
		$6::uuid,
		$7::uuid,
		$8::jsonb;`)
	if err != nil {
		log.Fatal("guide/job-save: ", err)
	}
	defer g.jbsv_stmt.Close()

	log.Println("Prepare: guide/repo-by-bucket")
	g.repo_stmt, err = g.conn.Prepare(`
SELECT	sb.repository_id,
		sr.repository_name
FROM	soma.buckets sb
JOIN    soma.repositories sr
ON		sb.repository_id = sr.repository_id
WHERE	sb.bucket_id = $1::uuid;`)
	if err != nil {
		log.Fatal("guide/repo-by-bucket: ", err)
	}
	defer g.repo_stmt.Close()

	log.Println("Prepare: guide/load-node-details")
	g.node_stmt, err = g.conn.Prepare(`
SELECT    sn.node_asset_id,
	      sn.node_name,
	      sn.organizational_team_id,
	      sn.server_id,
	      sn.node_online,
	      sn.node_deleted
FROM      soma.nodes sn
LEFT JOIN soma.node_bucket_assignment snba
ON        sn.node_id = snba.node_id
WHERE     sn.node_online = 'yes'
AND       sn.node_deleted = 'false'
AND       snba.node_id IS NULL
AND       sn.node_id = $1::uuid;`)
	if err != nil {
		log.Fatal("guide/load-node-details: ", err)
	}
	defer g.node_stmt.Close()

	log.Println("Prepare: guide/repo-by-id")
	g.name_stmt, err = g.conn.Prepare(`
SELECT repository_name
FROM   soma.repositories
WHERE  repository_id = $1::uuid;`)
	if err != nil {
		log.Fatal("guide/repo-by-id: ", err)
	}
	defer g.name_stmt.Close()

	log.Println("Prepare: guide/service-lookup")
	g.serv_stmt, err = g.conn.Prepare(`
SELECT stsp.service_property
FROM   soma.repositories sr
JOIN   soma.team_service_properties stsp
ON     sr.organizational_team_id = stsp.organizational_team_id
WHERE  sr.repository_id = $1::uuid
AND    stsp.service_property = $2::varchar
AND    sr.organizational_team_id = $3::uuid;`)
	if err != nil {
		log.Fatal("guide/service-lookup: ", err)
	}
	defer g.serv_stmt.Close()

	log.Println("Prepare: guide/populate-service-attributes")
	g.attr_stmt, err = g.conn.Prepare(`
SELECT stspv.service_property_attribute,
       stspv.value
FROM   soma.repositories sr
JOIN   soma.team_service_properties stsp
ON     sr.organizational_team_id = stsp.organizational_team_id
JOIN   soma.team_service_property_values stspv
ON     stsp.organizational_team_id = stspv.organizational_team_id
AND    stsp.service_property = stspv.service_property
WHERE  sr.repository_id = $1::uuid
AND    stsp.service_property = $2::varchar
AND    sr.organizational_team_id = $3::uuid;`)
	if err != nil {
		log.Fatal("guide/populate-service-attributes: ", err)
	}
	defer g.attr_stmt.Close()

runloop:
	for {
		select {
		case <-g.shutdown:
			break runloop
		case req := <-g.input:
			g.process(&req)
		}
	}
}

func (g *guidePost) process(q *treeRequest) {
	var (
		res                                sql.Result
		err                                error
		j                                  []byte
		repoId, repoName, keeper, bucketId string
		ndName, ndTeam, ndServer           string
		ndAsset                            int64
		ndOnline, ndDeleted                bool
	)
	result := somaResult{}

	switch q.Action {
	case "add_system_property_to_repository":
		fallthrough
	case "add_custom_property_to_repository":
		fallthrough
	case "add_oncall_property_to_repository":
		fallthrough
	case "add_service_property_to_repository":
		repoId = q.Repository.Repository.Id

	case "create_bucket":
		repoId = q.Bucket.Bucket.Repository
	case "add_system_property_to_bucket":
		fallthrough
	case "add_custom_property_to_bucket":
		fallthrough
	case "add_oncall_property_to_bucket":
		fallthrough
	case "add_service_property_to_bucket":
		bucketId = q.Bucket.Bucket.Id

	case "create_group":
		fallthrough
	case "add_group_to_group":
		fallthrough
	case "add_cluster_to_group":
		fallthrough
	case "add_node_to_group":
		fallthrough
	case "add_system_property_to_group":
		fallthrough
	case "add_custom_property_to_group":
		fallthrough
	case "add_oncall_property_to_group":
		fallthrough
	case "add_service_property_to_group":
		bucketId = q.Group.Group.BucketId

	case "add_node_to_cluster":
		if q.Cluster.Cluster.BucketId != q.Cluster.Cluster.Members[0].Config.BucketId {
			panic("This should not happen.")
		}
		fallthrough
	case "create_cluster":
		fallthrough
	case "add_system_property_to_cluster":
		fallthrough
	case "add_custom_property_to_cluster":
		fallthrough
	case "add_oncall_property_to_cluster":
		fallthrough
	case "add_service_property_to_cluster":
		bucketId = q.Cluster.Cluster.BucketId

	case "add_check_to_repository":
		fallthrough
	case "add_check_to_bucket":
		fallthrough
	case "add_check_to_group":
		fallthrough
	case "add_check_to_cluster":
		fallthrough
	case "add_check_to_node":
		repoId = q.CheckConfig.CheckConfig.RepositoryId

	case "assign_node":
		if err = g.node_stmt.QueryRow(q.Node.Node.Id).Scan(
			&ndAsset,
			&ndName,
			&ndTeam,
			&ndServer,
			&ndOnline,
			&ndDeleted,
		); err != nil {
			if err == sql.ErrNoRows {
				result.SetNotFound()
			} else {
				_ = result.SetRequestError(err)
			}
			q.reply <- result
			return
		}

		q.Node.Node.AssetId = uint64(ndAsset)
		q.Node.Node.Name = ndName
		q.Node.Node.Team = ndTeam
		q.Node.Node.Server = ndServer
		q.Node.Node.IsOnline = ndOnline
		q.Node.Node.IsDeleted = ndDeleted
		fallthrough
	case "add_system_property_to_node":
		fallthrough
	case "add_custom_property_to_node":
		fallthrough
	case "add_oncall_property_to_node":
		fallthrough
	case "add_service_property_to_node":
		repoId = q.Node.Node.Config.RepositoryId
		bucketId = q.Node.Node.Config.BucketId

	default:
		log.Printf("R: unimplemented guidepost/%s", q.Action)
		result.SetNotImplemented()
		q.reply <- result
		return
	}

	// lookup repository by bucket
	if bucketId != "" {
		if err = g.repo_stmt.QueryRow(bucketId).Scan(&repoId, &repoName); err != nil {
			if err == sql.ErrNoRows {
				result.SetNotFound()
			} else {
				_ = result.SetRequestError(err)
			}
			q.reply <- result
			return
		}
	}

	// lookup repository name
	if repoName == "" && repoId != "" {
		if err = g.name_stmt.QueryRow(repoId).Scan(&repoName); err != nil {
			if err == sql.ErrNoRows {
				result.SetNotFound()
			} else {
				_ = result.SetRequestError(err)
			}
			q.reply <- result
			return
		}
	}

	// XXX
	if repoName == "" {
		panic(`Have no repository name`)
	}

	// check we have a treekeeper for that repository
	keeper = fmt.Sprintf("repository_%s", repoName)
	if _, ok := handlerMap[keeper].(*treeKeeper); !ok {
		_ = result.SetRequestError(
			fmt.Errorf("No handler for repository %s registered.\n", repoName),
		)
		q.reply <- result
		return
	}

	// check the treekeeper has finished loading
	handler := handlerMap[keeper].(*treeKeeper)
	if !handler.isReady() {
		_ = result.SetRequestError( // TODO should be 503/ServiceUnavailable
			fmt.Errorf("Repository %s not fully loaded yet.\n", repoName),
		)
		q.reply <- result
		return
	}

	// check the treekeeper has not encountered a broken tree
	if handler.isBroken() {
		_ = result.SetRequestError(
			fmt.Errorf("Repository %s is broken.\n", repoName),
		)
		q.reply <- result
		return
	}

	// load authoritative copy of the service attributes from the
	// database. Replaces whatever the client sent in.
	if strings.Contains(q.Action, "add_service_property_to_") {
		var service, attr, val, svName, svTeam string
		var rows *sql.Rows
		attrs := []somaproto.TreeServiceAttribute{}

		switch q.RequestType {
		case "repository":
			svName = (*q.Repository.Repository.Properties)[0].Service.Name
			svTeam = (*q.Repository.Repository.Properties)[0].Service.TeamId
		case "bucket":
			svName = (*q.Bucket.Bucket.Properties)[0].Service.Name
			svTeam = (*q.Bucket.Bucket.Properties)[0].Service.TeamId
		case "group":
			svName = (*q.Group.Group.Properties)[0].Service.Name
			svTeam = (*q.Group.Group.Properties)[0].Service.TeamId
		case "cluster":
			svName = (*q.Cluster.Cluster.Properties)[0].Service.Name
			svTeam = (*q.Cluster.Cluster.Properties)[0].Service.TeamId
		case "node":
			svName = (*q.Node.Node.Properties)[0].Service.Name
			svTeam = (*q.Node.Node.Properties)[0].Service.TeamId
		}

		if err = g.serv_stmt.QueryRow(repoId, svName, svTeam).Scan(&service); err != nil {
			goto inputabort
		}

		if rows, err = g.attr_stmt.Query(repoId, svName, svTeam); err != nil {
			goto inputabort
		}
		defer rows.Close()

	attrloop:
		for rows.Next() {
			if err = rows.Scan(&attr, &val); err != nil {
				break attrloop
			}
			attrs = append(attrs, somaproto.TreeServiceAttribute{
				Attribute: attr,
				Value:     val,
			})
		}

	inputabort:
		if err != nil {
			if err == sql.ErrNoRows {
				result.SetNotFound()
			} else {
				_ = result.SetRequestError(err)
			}
			q.reply <- result
			return
		}
		// not aborted: set the loaded attributes
		switch q.RequestType {
		case "repository":
			(*q.Repository.Repository.Properties)[0].Service.Attributes = attrs
		case "bucket":
			(*q.Bucket.Bucket.Properties)[0].Service.Attributes = attrs
		case "group":
			(*q.Group.Group.Properties)[0].Service.Attributes = attrs
		case "cluster":
			(*q.Cluster.Cluster.Properties)[0].Service.Attributes = attrs
		case "node":
			(*q.Node.Node.Properties)[0].Service.Attributes = attrs
		}
	}

	// store job in database
	log.Printf("R: jobsave/%s", q.Action)
	q.JobId = uuid.NewV4()
	j, _ = json.Marshal(q)
	res, err = g.jbsv_stmt.Exec(
		q.JobId.String(),
		"queued",
		"pending",
		q.Action,
		repoId,
		"00000000-0000-0000-0000-000000000000", // XXX user uuid
		"00000000-0000-0000-0000-000000000000", // XXX team uuid
		string(j),
	)
	if result.SetRequestError(err) {
		q.reply <- result
		return
	}
	rowCnt, _ := res.RowsAffected()
	switch {
	case rowCnt == 0:
		err = errors.New("No rows affected")
	case rowCnt > 1:
		err = fmt.Errorf("Too many rows affected: %d", rowCnt)
	case rowCnt < 0:
		err = fmt.Errorf("Space/Time Continuum broke, rows affected: %d", rowCnt)
	}
	if err != nil {
		switch q.RequestType {
		case "repository":
			result.Append(err, &somaRepositoryResult{})
		case "bucket":
			result.Append(err, &somaBucketResult{})
		case "group":
			result.Append(err, &somaGroupResult{})
		case "cluster":
			result.Append(err, &somaClusterResult{})
		case "node":
			result.Append(err, &somaNodeResult{})
		case "check":
			result.Append(err, &somaCheckConfigResult{})
		}
		q.reply <- result
		return
	}

	handler.input <- *q
	result.JobId = q.JobId.String()

	switch q.RequestType {
	case "repository":
		result.Append(nil, &somaRepositoryResult{
			Repository: q.Repository.Repository,
		})
	case "bucket":
		result.Append(nil, &somaBucketResult{
			Bucket: q.Bucket.Bucket,
		})
	case "group":
		result.Append(nil, &somaGroupResult{
			Group: q.Group.Group,
		})
	case "cluster":
		result.Append(nil, &somaClusterResult{
			Cluster: q.Cluster.Cluster,
		})
	case "node":
		result.Append(nil, &somaNodeResult{
			Node: q.Node.Node,
		})
	case "check":
		result.Append(nil, &somaCheckConfigResult{
			CheckConfig: q.CheckConfig.CheckConfig,
		})
	}
	q.reply <- result
}

// vim: ts=4 sw=4 sts=4 noet fenc=utf-8 ffs=unix
<|endoftext|>


<commit_before>/* This file is part of VoltDB.
 * Copyright (C) 2008-2016 VoltDB Inc.
 *
 * 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 VoltDB.  If not, see <http://www.gnu.org/licenses/>.
 */

// A simple example that demonstrates the use of the VoltDB database/sql/driver.
package main

import "database/sql"
import (
	"fmt"
	"log"
	"os"

	_ "github.com/VoltDB/voltdb-client-go/voltdbclient"
)

func main() {

	db, err := sql.Open("voltdb", "volt3g:21212/voltdb")
	if err != nil {
		fmt.Println("open")
		log.Fatal(err)
		os.Exit(-1)
	}
	defer db.Close()

	err = db.Ping()
	if err != nil {
		fmt.Println("open")
		log.Fatal(err)
		os.Exit(-1)
	}
	rows, err := db.Query("HELLOWORLD.select", "French")
	if err != nil {
		log.Fatal(err)
		os.Exit(-1)
	}
	printRows(rows)

	// with prepared statement
	stmt, err := db.Prepare("select * from HELLOWORLD where dialect = ?")
	if err != nil {
		log.Fatal(err)
		os.Exit(-1)
	}

	rows, err = stmt.Query("French")
	if err != nil {
		log.Fatal(err)
		os.Exit(-1)
	}
	printRows(rows)

}

func printRows(rows *sql.Rows) {
	for rows.Next() {
		var hello string
		var world string
		var dialect string
		err := rows.Scan(&hello, &world, &dialect)
		if err != nil {
			fmt.Println(err)
			break
		}
		fmt.Printf("SUCCESS: %s %s %s\n", hello, world, dialect)
	}
}
<commit_msg>ENG-10558 merge master down into ENG-10558<commit_after>/* This file is part of VoltDB.
 * Copyright (C) 2008-2016 VoltDB Inc.
 *
 * 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 VoltDB.  If not, see <http://www.gnu.org/licenses/>.
 */

// A simple example that demonstrates the use of the VoltDB database/sql/driver.
package main

import "database/sql"
import (
	"fmt"
	"log"
	"os"

	_ "github.com/VoltDB/voltdb-client-go/voltdbclient"
)

func main() {

	db, err := sql.Open("voltdb", "localhost:21212/voltdb")
	if err != nil {
		fmt.Println("open")
		log.Fatal(err)
		os.Exit(-1)
	}
	defer db.Close()

	err = db.Ping()
	if err != nil {
		fmt.Println("open")
		log.Fatal(err)
		os.Exit(-1)
	}
	rows, err := db.Query("HELLOWORLD.select", "French")
	if err != nil {
		log.Fatal(err)
		os.Exit(-1)
	}
	printRows(rows)

	// with prepared statement
	stmt, err := db.Prepare("select * from HELLOWORLD where dialect = ?")
	if err != nil {
		log.Fatal(err)
		os.Exit(-1)
	}

	rows, err = stmt.Query("French")
	if err != nil {
		log.Fatal(err)
		os.Exit(-1)
	}
	printRows(rows)

}

func printRows(rows *sql.Rows) {
	for rows.Next() {
		var hello string
		var world string
		var dialect string
		err := rows.Scan(&hello, &world, &dialect)
		if err != nil {
			fmt.Println(err)
			break
		}
		fmt.Printf("SUCCESS: %s %s %s\n", hello, world, dialect)
	}
}
<|endoftext|>


<commit_before>
package common

import (
  "encoding/base64"
  "strings"
  "github.com/ricallinson/forgery"
  "github.com/spacedock-io/index/models"
)

func UnpackAuth(raw string) (creds []string, err error) {
  auth := strings.Split(raw, " ")
  decoded, err := base64.StdEncoding.DecodeString(auth[1])
  if err != nil { return nil, err }

  creds = strings.Split(string(decoded), ":")
  return creds, nil
}

func CheckAuth(req *f.Request, res *f.Response, next func()) {
  auth := req.Get("authorization")

  if len(auth) == 0 {
    res.Send("No authorization provided.", 401)
    return
  }

  creds, err := UnpackAuth(auth)
  if err != nil {
    res.Send("Unauthorized", 401)
    return
  }

  u, ok := models.AuthUser(creds[0], creds[1])
  if !ok {
    res.Send("Unauthorized", 401)
  }
  req.Map["_uid"] = u.Id
}
<commit_msg>Add req.Map["admin"], clear req.Map security fields<commit_after>
package common

import (
  "encoding/base64"
  "strings"
  "github.com/ricallinson/forgery"
  "github.com/spacedock-io/index/models"
)

func UnpackAuth(raw string) (creds []string, err error) {
  auth := strings.Split(raw, " ")
  decoded, err := base64.StdEncoding.DecodeString(auth[1])
  if err != nil { return nil, err }

  creds = strings.Split(string(decoded), ":")
  return creds, nil
}

func CheckAuth(req *f.Request, res *f.Response, next func()) {
  auth := req.Get("authorization")
  req.Map["_uid"] = -1
  req.Map["_admin"] = false

  if len(auth) == 0 {
    res.Send("No authorization provided.", 401)
    return
  }

  creds, err := UnpackAuth(auth)
  if err != nil {
    res.Send("Unauthorized", 401)
    return
  }

  u, ok := models.AuthUser(creds[0], creds[1])
  if !ok {
    res.Send("Unauthorized", 401)
  }
  req.Map["_uid"] = u.Id
  req.Map["_admin"] = u.Admin
}
<|endoftext|>


<commit_before>/*
Copyright 2014 Google Inc. All rights reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package vagrant_cloud

import (
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	neturl "net/url"
	"sort"

	"github.com/GoogleCloudPlatform/kubernetes/pkg/api"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/cloudprovider"
)

// VagrantCloud is an implementation of Interface, TCPLoadBalancer and Instances for developer managed Vagrant cluster.
type VagrantCloud struct {
	saltURL  string
	saltUser string
	saltPass string
	saltAuth string
}

func init() {
	cloudprovider.RegisterCloudProvider("vagrant", func(config io.Reader) (cloudprovider.Interface, error) { return newVagrantCloud() })
}

// SaltToken is an authorization token required by Salt REST API.
type SaltToken struct {
	Token string `json:"token"`
	User  string `json:"user"`
	EAuth string `json:"eauth"`
}

// SaltLoginResponse is the response object for a /login operation against Salt REST API.
type SaltLoginResponse struct {
	Data []SaltToken `json:"return"`
}

// SaltMinion is a machine managed by the Salt service.
type SaltMinion struct {
	Roles []string `json:"roles"`
	IP    string   `json:"node_ip"`
	Host  string   `json:"host"`
}

// SaltMinions is a map of minion name to machine information.
type SaltMinions map[string]SaltMinion

// SaltMinionsResponse is the response object for a /minions operation against Salt REST API
type SaltMinionsResponse struct {
	Minions []SaltMinions `json:"return"`
}

// newVagrantCloud creates a new instance of VagrantCloud configured to talk to the Salt REST API.
func newVagrantCloud() (*VagrantCloud, error) {
	return &VagrantCloud{
		saltURL:  "http://127.0.0.1:8000",
		saltUser: "vagrant",
		saltPass: "vagrant",
		saltAuth: "pam",
	}, nil
}

func (aws *VagrantCloud) Clusters() (cloudprovider.Clusters, bool) {
	return nil, false
}

// TCPLoadBalancer returns an implementation of TCPLoadBalancer for Vagrant cloud.
func (v *VagrantCloud) TCPLoadBalancer() (cloudprovider.TCPLoadBalancer, bool) {
	return nil, false
}

// Instances returns an implementation of Instances for Vagrant cloud.
func (v *VagrantCloud) Instances() (cloudprovider.Instances, bool) {
	return v, true
}

// Zones returns an implementation of Zones for Vagrant cloud.
func (v *VagrantCloud) Zones() (cloudprovider.Zones, bool) {
	return nil, false
}

// IPAddress returns the address of a particular machine instance.
func (v *VagrantCloud) IPAddress(instance string) (net.IP, error) {
	token, err := v.saltLogin()
	if err != nil {
		return nil, err
	}
	minions, err := v.saltMinions(token)
	if err != nil {
		return nil, err
	}
	filteredMinions := v.saltMinionsByRole(minions, "kubernetes-pool")
	for _, minion := range filteredMinions {
		// Due to vagrant not running with a dedicated DNS setup, we return the IP address of a minion as its hostname at this time
		if minion.IP == instance {
			return net.ParseIP(minion.IP), nil
		}
	}
	return nil, fmt.Errorf("unable to find IP address for instance: %s", instance)
}

// saltMinionsByRole filters a list of minions that have a matching role.
func (v *VagrantCloud) saltMinionsByRole(minions []SaltMinion, role string) []SaltMinion {
	var filteredMinions []SaltMinion
	for _, value := range minions {
		sort.Strings(value.Roles)
		if pos := sort.SearchStrings(value.Roles, role); pos < len(value.Roles) {
			filteredMinions = append(filteredMinions, value)
		}
	}
	return filteredMinions
}

// saltMinions invokes the Salt API for minions using provided token.
func (v *VagrantCloud) saltMinions(token SaltToken) ([]SaltMinion, error) {
	var minions []SaltMinion

	url := v.saltURL + "/minions"
	req, err := http.NewRequest("GET", url, nil)
	req.Header.Add("X-Auth-Token", token.Token)

	client := &http.Client{}
	resp, err := client.Do(req)

	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return minions, err
	}

	var minionsResp SaltMinionsResponse
	if err = json.Unmarshal(body, &minionsResp); err != nil {
		return minions, err
	}

	for _, value := range minionsResp.Minions[0] {
		minions = append(minions, value)
	}

	return minions, nil
}

// saltLogin invokes the Salt API to get an authorization token.
func (v *VagrantCloud) saltLogin() (SaltToken, error) {
	url := v.saltURL + "/login"
	data := neturl.Values{
		"username": {v.saltUser},
		"password": {v.saltPass},
		"eauth":    {v.saltAuth},
	}

	var token SaltToken
	resp, err := http.PostForm(url, data)
	if err != nil {
		return token, err
	}

	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return token, err
	}

	var loginResp SaltLoginResponse
	if err := json.Unmarshal(body, &loginResp); err != nil {
		return token, err
	}

	if len(loginResp.Data) == 0 {
		return token, errors.New("No token found in response")
	}

	return loginResp.Data[0], nil
}

// List enumerates the set of minions instances known by the cloud provider.
func (v *VagrantCloud) List(filter string) ([]string, error) {
	token, err := v.saltLogin()
	if err != nil {
		return nil, err
	}

	minions, err := v.saltMinions(token)
	if err != nil {
		return nil, err
	}

	filteredMinions := v.saltMinionsByRole(minions, "kubernetes-pool")
	var instances []string
	for _, instance := range filteredMinions {
		// With no dedicated DNS setup in cluster, IP address is used as hostname
		instances = append(instances, instance.IP)
	}

	return instances, nil
}

func (v *VagrantCloud) GetNodeResources(name string) (*api.NodeResources, error) {
	return nil, nil
}
<commit_msg>Fix typos in vagrant.go<commit_after>/*
Copyright 2014 Google Inc. All rights reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package vagrant_cloud

import (
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	neturl "net/url"
	"sort"

	"github.com/GoogleCloudPlatform/kubernetes/pkg/api"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/cloudprovider"
)

// VagrantCloud is an implementation of Interface, TCPLoadBalancer and Instances for developer managed Vagrant cluster.
type VagrantCloud struct {
	saltURL  string
	saltUser string
	saltPass string
	saltAuth string
}

func init() {
	cloudprovider.RegisterCloudProvider("vagrant", func(config io.Reader) (cloudprovider.Interface, error) { return newVagrantCloud() })
}

// SaltToken is an authorization token required by Salt REST API.
type SaltToken struct {
	Token string `json:"token"`
	User  string `json:"user"`
	EAuth string `json:"eauth"`
}

// SaltLoginResponse is the response object for a /login operation against Salt REST API.
type SaltLoginResponse struct {
	Data []SaltToken `json:"return"`
}

// SaltMinion is a machine managed by the Salt service.
type SaltMinion struct {
	Roles []string `json:"roles"`
	IP    string   `json:"node_ip"`
	Host  string   `json:"host"`
}

// SaltMinions is a map of minion name to machine information.
type SaltMinions map[string]SaltMinion

// SaltMinionsResponse is the response object for a /minions operation against Salt REST API
type SaltMinionsResponse struct {
	Minions []SaltMinions `json:"return"`
}

// newVagrantCloud creates a new instance of VagrantCloud configured to talk to the Salt REST API.
func newVagrantCloud() (*VagrantCloud, error) {
	return &VagrantCloud{
		saltURL:  "http://127.0.0.1:8000",
		saltUser: "vagrant",
		saltPass: "vagrant",
		saltAuth: "pam",
	}, nil
}

func (v *VagrantCloud) Clusters() (cloudprovider.Clusters, bool) {
	return nil, false
}

// TCPLoadBalancer returns an implementation of TCPLoadBalancer for Vagrant cloud.
func (v *VagrantCloud) TCPLoadBalancer() (cloudprovider.TCPLoadBalancer, bool) {
	return nil, false
}

// Instances returns an implementation of Instances for Vagrant cloud.
func (v *VagrantCloud) Instances() (cloudprovider.Instances, bool) {
	return v, true
}

// Zones returns an implementation of Zones for Vagrant cloud.
func (v *VagrantCloud) Zones() (cloudprovider.Zones, bool) {
	return nil, false
}

// IPAddress returns the address of a particular machine instance.
func (v *VagrantCloud) IPAddress(instance string) (net.IP, error) {
	token, err := v.saltLogin()
	if err != nil {
		return nil, err
	}
	minions, err := v.saltMinions(token)
	if err != nil {
		return nil, err
	}
	filteredMinions := v.saltMinionsByRole(minions, "kubernetes-pool")
	for _, minion := range filteredMinions {
		// Due to vagrant not running with a dedicated DNS setup, we return the IP address of a minion as its hostname at this time
		if minion.IP == instance {
			return net.ParseIP(minion.IP), nil
		}
	}
	return nil, fmt.Errorf("unable to find IP address for instance: %s", instance)
}

// saltMinionsByRole filters a list of minions that have a matching role.
func (v *VagrantCloud) saltMinionsByRole(minions []SaltMinion, role string) []SaltMinion {
	var filteredMinions []SaltMinion
	for _, value := range minions {
		sort.Strings(value.Roles)
		if pos := sort.SearchStrings(value.Roles, role); pos < len(value.Roles) {
			filteredMinions = append(filteredMinions, value)
		}
	}
	return filteredMinions
}

// saltMinions invokes the Salt API for minions using provided token.
func (v *VagrantCloud) saltMinions(token SaltToken) ([]SaltMinion, error) {
	var minions []SaltMinion

	url := v.saltURL + "/minions"
	req, err := http.NewRequest("GET", url, nil)
	req.Header.Add("X-Auth-Token", token.Token)

	client := &http.Client{}
	resp, err := client.Do(req)

	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return minions, err
	}

	var minionsResp SaltMinionsResponse
	if err = json.Unmarshal(body, &minionsResp); err != nil {
		return minions, err
	}

	for _, value := range minionsResp.Minions[0] {
		minions = append(minions, value)
	}

	return minions, nil
}

// saltLogin invokes the Salt API to get an authorization token.
func (v *VagrantCloud) saltLogin() (SaltToken, error) {
	url := v.saltURL + "/login"
	data := neturl.Values{
		"username": {v.saltUser},
		"password": {v.saltPass},
		"eauth":    {v.saltAuth},
	}

	var token SaltToken
	resp, err := http.PostForm(url, data)
	if err != nil {
		return token, err
	}

	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return token, err
	}

	var loginResp SaltLoginResponse
	if err := json.Unmarshal(body, &loginResp); err != nil {
		return token, err
	}

	if len(loginResp.Data) == 0 {
		return token, errors.New("No token found in response")
	}

	return loginResp.Data[0], nil
}

// List enumerates the set of minions instances known by the cloud provider.
func (v *VagrantCloud) List(filter string) ([]string, error) {
	token, err := v.saltLogin()
	if err != nil {
		return nil, err
	}

	minions, err := v.saltMinions(token)
	if err != nil {
		return nil, err
	}

	filteredMinions := v.saltMinionsByRole(minions, "kubernetes-pool")
	var instances []string
	for _, instance := range filteredMinions {
		// With no dedicated DNS setup in cluster, IP address is used as hostname
		instances = append(instances, instance.IP)
	}

	return instances, nil
}

func (v *VagrantCloud) GetNodeResources(name string) (*api.NodeResources, error) {
	return nil, nil
}
<|endoftext|>


<commit_before>/*
 * This file is part of the KubeVirt project
 *
 * 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.
 *
 * Copyright 2017 Red Hat, Inc.
 *
 */

package containerdisk

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

	kubev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/resource"

	ephemeraldisk "kubevirt.io/kubevirt/pkg/ephemeral-disk"

	v1 "kubevirt.io/client-go/api/v1"
	diskutils "kubevirt.io/kubevirt/pkg/ephemeral-disk-utils"
	"kubevirt.io/kubevirt/pkg/util"
)

var containerDiskOwner = "qemu"

var podsBaseDir = util.KubeletPodsDir

var mountBaseDir = filepath.Join(util.VirtShareDir, "/container-disks")

type SocketPathGetter func(vmi *v1.VirtualMachineInstance, volumeIndex int) (string, error)

const ephemeralStorageOverheadSize = "50M"

func GetLegacyVolumeMountDirOnHost(vmi *v1.VirtualMachineInstance) string {
	return filepath.Join(mountBaseDir, string(vmi.UID))
}

func GetVolumeMountDirOnGuest(vmi *v1.VirtualMachineInstance) string {
	return filepath.Join(mountBaseDir, string(vmi.UID))
}

func GetVolumeMountDirOnHost(vmi *v1.VirtualMachineInstance) (string, bool, error) {
	basepath := ""
	foundEntries := 0
	foundBasepath := ""
	for podUID, _ := range vmi.Status.ActivePods {
		basepath = fmt.Sprintf("%s/%s/volumes/kubernetes.io~empty-dir/container-disks", podsBaseDir, string(podUID))
		exists, err := diskutils.FileExists(basepath)
		if err != nil {
			return "", false, err
		} else if exists {
			foundEntries++
			foundBasepath = basepath
		}
	}

	if foundEntries == 1 {
		return foundBasepath, true, nil
	} else if foundEntries > 1 {
		// Don't mount until outdated pod environments are removed
		// otherwise we might stomp on a previous cleanup
		return "", false, fmt.Errorf("Found multiple pods active for vmi %s/%s. Waiting on outdated pod directories to be removed", vmi.Namespace, vmi.Name)
	}
	return "", false, nil
}

func GetDiskTargetPathFromHostView(vmi *v1.VirtualMachineInstance, volumeIndex int) (string, error) {
	basepath, found, err := GetVolumeMountDirOnHost(vmi)
	if err != nil {
		return "", err
	} else if !found {
		return "", fmt.Errorf("container disk volume for vmi not found")
	}

	return fmt.Sprintf("%s/disk_%d.img", basepath, volumeIndex), nil
}

func GetDiskTargetPathFromLauncherView(volumeIndex int) string {
	return filepath.Join(mountBaseDir, fmt.Sprintf("disk_%d.img", volumeIndex))
}

func SetLocalDirectory(dir string) error {
	mountBaseDir = dir
	return os.MkdirAll(dir, 0755)
}

func SetKubeletPodsDirectory(dir string) {
	podsBaseDir = dir
}

// used for testing - we don't want to MkdirAll on a production host mount
func setPodsDirectory(dir string) error {
	podsBaseDir = dir
	return os.MkdirAll(dir, 0755)
}

// The unit test suite uses this function
func setLocalDataOwner(user string) {
	containerDiskOwner = user
}

// GetDiskTargetPartFromLauncherView returns (path to disk image, image type, and error)
func GetDiskTargetPartFromLauncherView(volumeIndex int) (string, error) {

	path := GetDiskTargetPathFromLauncherView(volumeIndex)
	exists, err := diskutils.FileExists(path)
	if err != nil {
		return "", err
	} else if exists {
		return path, nil
	}

	return "", fmt.Errorf("no supported file disk found for volume with index %d", volumeIndex)
}

// NewSocketPathGetter get the socket pat of a containerDisk. For testing a baseDir
// can be provided which can for instance point to /tmp.
func NewSocketPathGetter(baseDir string) SocketPathGetter {
	return func(vmi *v1.VirtualMachineInstance, volumeIndex int) (string, error) {
		for podUID, _ := range vmi.Status.ActivePods {
			basepath := fmt.Sprintf("%s/pods/%s/volumes/kubernetes.io~empty-dir/container-disks", baseDir, string(podUID))
			socketPath := filepath.Join(basepath, fmt.Sprintf("disk_%d.sock", volumeIndex))
			exists, _ := diskutils.FileExists(socketPath)
			if exists {
				return socketPath, nil
			}
		}
		return "", fmt.Errorf("container disk socket path not found for vmi")
	}
}

func GetImage(root string, imagePath string) (string, error) {
	fallbackPath := filepath.Join(root, DiskSourceFallbackPath)
	if imagePath != "" {
		imagePath = filepath.Join(root, imagePath)
		if _, err := os.Stat(imagePath); err != nil {
			if os.IsNotExist(err) {
				return "", fmt.Errorf("No image on path %s", imagePath)
			} else {
				return "", fmt.Errorf("Failed to check if an image exists at %s", imagePath)
			}
		}
	} else {
		files, err := ioutil.ReadDir(fallbackPath)
		if err != nil {
			return "", fmt.Errorf("Failed to check %s for disks: %v", fallbackPath, err)
		}
		if len(files) > 1 {
			return "", fmt.Errorf("More than one file found in folder %s, only one disk is allowed", DiskSourceFallbackPath)
		}
		imagePath = filepath.Join(fallbackPath, files[0].Name())
	}
	return imagePath, nil
}

func GenerateInitContainers(vmi *v1.VirtualMachineInstance, podVolumeName string, binVolumeName string) []kubev1.Container {
	return generateContainersHelper(vmi, podVolumeName, binVolumeName, true)
}

func GenerateContainers(vmi *v1.VirtualMachineInstance, podVolumeName string, binVolumeName string) []kubev1.Container {
	return generateContainersHelper(vmi, podVolumeName, binVolumeName, false)
}

// The controller uses this function to generate the container
// specs for hosting the container registry disks.
func generateContainersHelper(vmi *v1.VirtualMachineInstance, podVolumeName string, binVolumeName string, isInit bool) []kubev1.Container {
	var containers []kubev1.Container

	// Make VirtualMachineInstance Image Wrapper Containers
	for index, volume := range vmi.Spec.Volumes {
		if volume.ContainerDisk != nil {

			volumeMountDir := GetVolumeMountDirOnGuest(vmi)
			diskContainerName := fmt.Sprintf("volume%s", volume.Name)
			diskContainerImage := volume.ContainerDisk.Image
			resources := kubev1.ResourceRequirements{}
			if vmi.IsCPUDedicated() || vmi.WantsToHaveQOSGuaranteed() {
				resources.Limits = make(kubev1.ResourceList)
				resources.Limits[kubev1.ResourceCPU] = resource.MustParse("10m")
				resources.Limits[kubev1.ResourceMemory] = resource.MustParse("40M")
				resources.Requests = make(kubev1.ResourceList)
				resources.Requests[kubev1.ResourceCPU] = resource.MustParse("10m")
				resources.Requests[kubev1.ResourceMemory] = resource.MustParse("40M")
				resources.Requests[kubev1.ResourceEphemeralStorage] = resource.MustParse(ephemeralStorageOverheadSize)
			} else {
				resources.Limits = make(kubev1.ResourceList)
				resources.Limits[kubev1.ResourceCPU] = resource.MustParse("100m")
				resources.Limits[kubev1.ResourceMemory] = resource.MustParse("40M")
				resources.Requests = make(kubev1.ResourceList)
				resources.Requests[kubev1.ResourceCPU] = resource.MustParse("10m")
				resources.Requests[kubev1.ResourceMemory] = resource.MustParse("1M")
				resources.Requests[kubev1.ResourceEphemeralStorage] = resource.MustParse(ephemeralStorageOverheadSize)
			}
			var args []string
			var name string
			if isInit {
				name = diskContainerName + "-init"
				args = []string{"--no-op"}
			} else {
				name = diskContainerName
				args = []string{"--copy-path", volumeMountDir + "/disk_" + strconv.Itoa(index)}
			}
			container := kubev1.Container{
				Name:            name,
				Image:           diskContainerImage,
				ImagePullPolicy: volume.ContainerDisk.ImagePullPolicy,
				Command:         []string{"/usr/bin/container-disk"},
				Args:            args,
				VolumeMounts: []kubev1.VolumeMount{
					{
						Name:      podVolumeName,
						MountPath: volumeMountDir,
					},
					{
						Name:      binVolumeName,
						MountPath: "/usr/bin",
					},
				},
				Resources: resources,
			}

			containers = append(containers, container)
		}
	}
	return containers
}

func CreateEphemeralImages(vmi *v1.VirtualMachineInstance) error {
	// The domain is setup to use the COW image instead of the base image. What we have
	// to do here is only create the image where the domain expects it (GetDiskTargetPartFromLauncherView)
	// for each disk that requires it.

	for i, volume := range vmi.Spec.Volumes {
		if volume.VolumeSource.ContainerDisk != nil {
			if backingFile, err := GetDiskTargetPartFromLauncherView(i); err != nil {
				return err
			} else if err := ephemeraldisk.CreateBackedImageForVolume(volume, backingFile); err != nil {
				return err
			}
		}
	}

	return nil
}
<commit_msg>[container-disk]: Remove code duplications from generateContainersHelper function<commit_after>/*
 * This file is part of the KubeVirt project
 *
 * 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.
 *
 * Copyright 2017 Red Hat, Inc.
 *
 */

package containerdisk

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

	kubev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/resource"

	ephemeraldisk "kubevirt.io/kubevirt/pkg/ephemeral-disk"

	v1 "kubevirt.io/client-go/api/v1"
	diskutils "kubevirt.io/kubevirt/pkg/ephemeral-disk-utils"
	"kubevirt.io/kubevirt/pkg/util"
)

var containerDiskOwner = "qemu"

var podsBaseDir = util.KubeletPodsDir

var mountBaseDir = filepath.Join(util.VirtShareDir, "/container-disks")

type SocketPathGetter func(vmi *v1.VirtualMachineInstance, volumeIndex int) (string, error)

const ephemeralStorageOverheadSize = "50M"

func GetLegacyVolumeMountDirOnHost(vmi *v1.VirtualMachineInstance) string {
	return filepath.Join(mountBaseDir, string(vmi.UID))
}

func GetVolumeMountDirOnGuest(vmi *v1.VirtualMachineInstance) string {
	return filepath.Join(mountBaseDir, string(vmi.UID))
}

func GetVolumeMountDirOnHost(vmi *v1.VirtualMachineInstance) (string, bool, error) {
	basepath := ""
	foundEntries := 0
	foundBasepath := ""
	for podUID, _ := range vmi.Status.ActivePods {
		basepath = fmt.Sprintf("%s/%s/volumes/kubernetes.io~empty-dir/container-disks", podsBaseDir, string(podUID))
		exists, err := diskutils.FileExists(basepath)
		if err != nil {
			return "", false, err
		} else if exists {
			foundEntries++
			foundBasepath = basepath
		}
	}

	if foundEntries == 1 {
		return foundBasepath, true, nil
	} else if foundEntries > 1 {
		// Don't mount until outdated pod environments are removed
		// otherwise we might stomp on a previous cleanup
		return "", false, fmt.Errorf("Found multiple pods active for vmi %s/%s. Waiting on outdated pod directories to be removed", vmi.Namespace, vmi.Name)
	}
	return "", false, nil
}

func GetDiskTargetPathFromHostView(vmi *v1.VirtualMachineInstance, volumeIndex int) (string, error) {
	basepath, found, err := GetVolumeMountDirOnHost(vmi)
	if err != nil {
		return "", err
	} else if !found {
		return "", fmt.Errorf("container disk volume for vmi not found")
	}

	return fmt.Sprintf("%s/disk_%d.img", basepath, volumeIndex), nil
}

func GetDiskTargetPathFromLauncherView(volumeIndex int) string {
	return filepath.Join(mountBaseDir, fmt.Sprintf("disk_%d.img", volumeIndex))
}

func SetLocalDirectory(dir string) error {
	mountBaseDir = dir
	return os.MkdirAll(dir, 0755)
}

func SetKubeletPodsDirectory(dir string) {
	podsBaseDir = dir
}

// used for testing - we don't want to MkdirAll on a production host mount
func setPodsDirectory(dir string) error {
	podsBaseDir = dir
	return os.MkdirAll(dir, 0755)
}

// The unit test suite uses this function
func setLocalDataOwner(user string) {
	containerDiskOwner = user
}

// GetDiskTargetPartFromLauncherView returns (path to disk image, image type, and error)
func GetDiskTargetPartFromLauncherView(volumeIndex int) (string, error) {

	path := GetDiskTargetPathFromLauncherView(volumeIndex)
	exists, err := diskutils.FileExists(path)
	if err != nil {
		return "", err
	} else if exists {
		return path, nil
	}

	return "", fmt.Errorf("no supported file disk found for volume with index %d", volumeIndex)
}

// NewSocketPathGetter get the socket pat of a containerDisk. For testing a baseDir
// can be provided which can for instance point to /tmp.
func NewSocketPathGetter(baseDir string) SocketPathGetter {
	return func(vmi *v1.VirtualMachineInstance, volumeIndex int) (string, error) {
		for podUID, _ := range vmi.Status.ActivePods {
			basepath := fmt.Sprintf("%s/pods/%s/volumes/kubernetes.io~empty-dir/container-disks", baseDir, string(podUID))
			socketPath := filepath.Join(basepath, fmt.Sprintf("disk_%d.sock", volumeIndex))
			exists, _ := diskutils.FileExists(socketPath)
			if exists {
				return socketPath, nil
			}
		}
		return "", fmt.Errorf("container disk socket path not found for vmi")
	}
}

func GetImage(root string, imagePath string) (string, error) {
	fallbackPath := filepath.Join(root, DiskSourceFallbackPath)
	if imagePath != "" {
		imagePath = filepath.Join(root, imagePath)
		if _, err := os.Stat(imagePath); err != nil {
			if os.IsNotExist(err) {
				return "", fmt.Errorf("No image on path %s", imagePath)
			} else {
				return "", fmt.Errorf("Failed to check if an image exists at %s", imagePath)
			}
		}
	} else {
		files, err := ioutil.ReadDir(fallbackPath)
		if err != nil {
			return "", fmt.Errorf("Failed to check %s for disks: %v", fallbackPath, err)
		}
		if len(files) > 1 {
			return "", fmt.Errorf("More than one file found in folder %s, only one disk is allowed", DiskSourceFallbackPath)
		}
		imagePath = filepath.Join(fallbackPath, files[0].Name())
	}
	return imagePath, nil
}

func GenerateInitContainers(vmi *v1.VirtualMachineInstance, podVolumeName string, binVolumeName string) []kubev1.Container {
	return generateContainersHelper(vmi, podVolumeName, binVolumeName, true)
}

func GenerateContainers(vmi *v1.VirtualMachineInstance, podVolumeName string, binVolumeName string) []kubev1.Container {
	return generateContainersHelper(vmi, podVolumeName, binVolumeName, false)
}

// The controller uses this function to generate the container
// specs for hosting the container registry disks.
func generateContainersHelper(vmi *v1.VirtualMachineInstance, podVolumeName string, binVolumeName string, isInit bool) []kubev1.Container {
	var containers []kubev1.Container

	// Make VirtualMachineInstance Image Wrapper Containers
	for index, volume := range vmi.Spec.Volumes {
		if volume.ContainerDisk != nil {

			volumeMountDir := GetVolumeMountDirOnGuest(vmi)
			diskContainerName := fmt.Sprintf("volume%s", volume.Name)
			diskContainerImage := volume.ContainerDisk.Image
			resources := kubev1.ResourceRequirements{}
			resources.Limits = make(kubev1.ResourceList)
			resources.Requests = make(kubev1.ResourceList)
			resources.Limits[kubev1.ResourceMemory] = resource.MustParse("40M")
			resources.Requests[kubev1.ResourceCPU] = resource.MustParse("10m")
			resources.Requests[kubev1.ResourceEphemeralStorage] = resource.MustParse(ephemeralStorageOverheadSize)

			if vmi.IsCPUDedicated() || vmi.WantsToHaveQOSGuaranteed() {
				resources.Limits[kubev1.ResourceCPU] = resource.MustParse("10m")
				resources.Requests[kubev1.ResourceMemory] = resource.MustParse("40M")
			} else {
				resources.Limits[kubev1.ResourceCPU] = resource.MustParse("100m")
				resources.Requests[kubev1.ResourceMemory] = resource.MustParse("1M")
			}
			var args []string
			var name string
			if isInit {
				name = diskContainerName + "-init"
				args = []string{"--no-op"}
			} else {
				name = diskContainerName
				args = []string{"--copy-path", volumeMountDir + "/disk_" + strconv.Itoa(index)}
			}
			container := kubev1.Container{
				Name:            name,
				Image:           diskContainerImage,
				ImagePullPolicy: volume.ContainerDisk.ImagePullPolicy,
				Command:         []string{"/usr/bin/container-disk"},
				Args:            args,
				VolumeMounts: []kubev1.VolumeMount{
					{
						Name:      podVolumeName,
						MountPath: volumeMountDir,
					},
					{
						Name:      binVolumeName,
						MountPath: "/usr/bin",
					},
				},
				Resources: resources,
			}

			containers = append(containers, container)
		}
	}
	return containers
}

func CreateEphemeralImages(vmi *v1.VirtualMachineInstance) error {
	// The domain is setup to use the COW image instead of the base image. What we have
	// to do here is only create the image where the domain expects it (GetDiskTargetPartFromLauncherView)
	// for each disk that requires it.

	for i, volume := range vmi.Spec.Volumes {
		if volume.VolumeSource.ContainerDisk != nil {
			if backingFile, err := GetDiskTargetPartFromLauncherView(i); err != nil {
				return err
			} else if err := ephemeraldisk.CreateBackedImageForVolume(volume, backingFile); err != nil {
				return err
			}
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package golang

import (
	"bytes"
	"fmt"
	"github.com/dustin/go-humanize"
	"go/ast"
	"go/build"
	"io/ioutil"
	"margo.sh/golang/cursor"
	"margo.sh/mg"
	"margo.sh/mgutil"
	"os"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"time"
)

type GoCmd struct {
	mg.ReducerType

	Humanize bool
}

func (gc *GoCmd) Reduce(mx *mg.Ctx) *mg.State {
	switch act := mx.Action.(type) {
	case mg.QueryUserCmds:
		return gc.userCmds(mx)
	case mg.RunCmd:
		return gc.runCmd(mx, act)
	}
	return mx.State
}

func (gc *GoCmd) userCmds(mx *mg.Ctx) *mg.State {
	return mx.AddUserCmds(
		mg.UserCmd{
			Name:  "go.play",
			Title: "Go Play",
		},
		mg.UserCmd{
			Name:  "go.replay",
			Title: "Go RePlay (single instance)",
		},
	)
}

func (gc *GoCmd) runCmd(mx *mg.Ctx, rc mg.RunCmd) *mg.State {
	return mx.State.AddBuiltinCmds(
		mg.BuiltinCmd{
			Run:  gc.goBuiltin,
			Name: "go",
			Desc: "Wrapper around the go command, adding linter support",
		},
		mg.BuiltinCmd{
			Run:  gc.playBuiltin,
			Name: "go.play",
			Desc: "Automatically build and run go commands or run go test for packages with support for linting and unsaved files",
		},
		mg.BuiltinCmd{
			Run:  gc.replayBuiltin,
			Name: "go.replay",
			Desc: "Wrapper around go.play limited to a single instance",
		},
	)
}

func (gc *GoCmd) goBuiltin(bx *mg.CmdCtx) *mg.State {
	go gc.goTool(bx)
	return bx.State
}

func (gc *GoCmd) playBuiltin(bx *mg.CmdCtx) *mg.State {
	go gc.playTool(bx, "")
	return bx.State
}

func (gc *GoCmd) replayBuiltin(bx *mg.CmdCtx) *mg.State {
	v := bx.View
	cid := ""
	if v.Path == "" {
		cid = v.Name
	} else {
		cid = v.Dir()
	}
	go gc.playTool(bx, "go.replay`"+cid+"`")
	return bx.State
}

func (gc *GoCmd) goTool(bx *mg.CmdCtx) {
	gx := newGoCmdCtx(gc, bx, "go.builtin", "", "", "", bx.View)
	defer gx.Output.Close()
	gx.run(gx.View)
}

func (gc *GoCmd) playTool(bx *mg.CmdCtx, cancelID string) {
	origView := bx.View
	bx, tDir, tFn, err := gc.playTempDir(bx)
	gx := newGoCmdCtx(gc, bx, "go.play", cancelID, tDir, tFn, origView)
	defer gx.Output.Close()

	gx.Verbose = true

	if err != nil {
		fmt.Fprintf(gx.Output, "Error: %s\n", err)
	}
	if tDir == "" {
		return
	}
	defer os.RemoveAll(tDir)

	bld := BuildContext(gx.Ctx)
	pkg, err := bld.ImportDir(gx.pkgDir, 0)
	switch {
	case err != nil:
		fmt.Fprintln(gx.Output, "Error: cannot import package:", err)
	case !pkg.IsCommand() || strings.HasSuffix(bx.View.Filename(), "_test.go"):
		gc.playToolTest(gx, bld, origView)
	default:
		gc.playToolRun(gx, bld, origView)
	}
}

func (gc *GoCmd) playTempDir(bx *mg.CmdCtx) (newBx *mg.CmdCtx, tDir string, tFn string, err error) {
	tDir, err = mg.MkTempDir("go.play")
	if err != nil {
		return bx, "", "", fmt.Errorf("cannot MkTempDir: %s", err)
	}

	if !bx.LangIs(mg.Go) {
		return bx, tDir, "", nil
	}

	v := bx.View
	if v.Path != "" {
		return bx, tDir, tFn, nil
	}

	tFn = filepath.Join(tDir, v.Name)
	src, err := v.ReadAll()
	if err == nil {
		err = ioutil.WriteFile(tFn, src, 0600)
	}
	if err != nil {
		return bx, tDir, "", fmt.Errorf("cannot create temp file: %s", err)
	}

	bx = bx.Copy(func(bx *mg.CmdCtx) {
		bx.Ctx = bx.Ctx.Copy(func(mx *mg.Ctx) {
			mx.State = mx.State.Copy(func(st *mg.State) {
				st.View = st.View.Copy(func(v *mg.View) {
					v.Path = tFn
				})
			})
		})
	})

	return bx, tDir, tFn, nil
}

func (gc *GoCmd) playToolTest(gx *goCmdCtx, bld *build.Context, origView *mg.View) {
	argsPfx := []string{"test", "-test.run=."}
	cx := cursor.NewViewCurCtx(gx.Ctx)
	for _, n := range cx.Nodes {
		x, ok := n.(*ast.FuncDecl)
		if !ok || x.Name == nil {
			continue
		}
		nm := x.Name.String()
		if strings.HasPrefix(nm, "Benchmark") {
			argsPfx = append(argsPfx, "-test.bench=^"+nm+"$")
		}
	}
	gx.Args = append(argsPfx, gx.Args...)
	if origView.Path == "" {
		gx.Args = append(gx.Args, gx.tFn)
	}
	gx.run(origView)
}

func (gc *GoCmd) playToolRun(gx *goCmdCtx, bld *build.Context, origView *mg.View) {
	nm := filepath.Base(origView.Name)
	if origView.Path != "" {
		nm = filepath.Base(origView.Dir())
	}

	args := gx.Args
	exe := filepath.Join(gx.tDir, "margo.play~~"+nm+".exe")
	gx.CmdCtx = gx.CmdCtx.Copy(func(bx *mg.CmdCtx) {
		bx.Name = "go"
		bx.Args = []string{"build", "-o", exe}
		bx.Ctx = bx.Ctx.Copy(func(mx *mg.Ctx) {
			mx.State = mx.State.Copy(func(st *mg.State) {
				st.View = st.View.Copy(func(v *mg.View) {
					v.Wd = v.Dir()
				})
			})
		})
	})
	if err := gx.run(origView); err != nil {
		return
	}

	gx.CmdCtx = gx.CmdCtx.Copy(func(bx *mg.CmdCtx) {
		bx.Name = exe
		bx.Args = args
		bx.Ctx = bx.Ctx.Copy(func(mx *mg.Ctx) {
			mx.State = mx.State.Copy(func(st *mg.State) {
				st.View = origView
			})
		})
	})
	gx.RunProc()
}

type goCmdCtx struct {
	*mg.CmdCtx
	pkgDir string
	key    interface{}
	iw     *mg.IssueOut
	tDir   string
	tFn    string
}

func newGoCmdCtx(gc *GoCmd, bx *mg.CmdCtx, label, cancelID string, tDir, tFn string, origView *mg.View) *goCmdCtx {
	gx := &goCmdCtx{
		pkgDir: bx.View.Dir(),
		tDir:   tDir,
		tFn:    tFn,
	}

	output := bx.Output
	if gc.Humanize && len(bx.Args) > 0 && bx.Args[0] == "test" {
		output = &humanizeWriter{output}
	}
	if gx.tFn != "" {
		dir := filepath.Dir(gx.tFn)
		qDir := regexp.QuoteMeta(dir)
		qDirBase := regexp.QuoteMeta(filepath.Base(dir))
		qNm := regexp.QuoteMeta(filepath.Base(gx.tFn))
		output = &replWriter{
			OutputStream: output,
			old: []*regexp.Regexp{
				regexp.MustCompile(`(?:` + qDir + `|` + qDirBase + `)?[\\/.]+` + qNm),
				regexp.MustCompile(qDir),
			},
			new: [][]byte{
				[]byte(origView.Name),
				[]byte(`tmp~`),
			},
		}
	}
	output = mgutil.NewSplitStream(mgutil.SplitLineOrCR, output)

	type Key struct{ label string }
	gx.key = Key{label}

	gx.iw = &mg.IssueOut{
		Base:     mg.Issue{Label: label},
		Patterns: bx.CommonPatterns(),
		Dir:      gx.pkgDir,
	}

	gx.CmdCtx = bx.Copy(func(bx *mg.CmdCtx) {
		bx.Name = "go"
		bx.CancelID = cancelID
		bx.Output = mg.OutputStreams{
			output,
			gx.iw,
		}
	})

	return gx
}

func (gx *goCmdCtx) run(origView *mg.View) error {
	defer func() {
		gx.VFS.Invalidate(origView.Filename())
		gx.VFS.Invalidate(origView.Dir())
	}()

	p, err := gx.StartProc()
	if err == nil {
		err = p.Wait()
	}
	gx.iw.Flush()

	issues := gx.iw.Issues()

	for i, isu := range issues {
		if isu.Path == "" || (gx.tFn != "" && filepath.Base(isu.Path) == origView.Name) {
			isu.Name = origView.Name
			isu.Path = origView.Path
		}
		issues[i] = isu
	}

	ik := mg.IssueKey{Key: gx.key}
	if origView.Path == "" {
		ik.Name = origView.Name
	} else {
		ik.Dir = origView.Dir()
	}

	gx.Store.Dispatch(mg.StoreIssues{IssueKey: ik, Issues: issues})
	return err
}

func isWhiteSpace(c byte) bool {
	return c == ' ' || c == '\t' || c == '\n' || c == '\r'
}

func humanizeMetric(met string) string {
	i := 0
	for i < len(met) && isWhiteSpace(met[i]) {
		i++
	}
	j := i
	for j < len(met) && !isWhiteSpace(met[j]) {
		j++
	}
	pfx := met[:i]
	val := met[i:j]
	sfx := met[j:]
	num, err := strconv.ParseInt(val, 10, 64)
	if err != nil {
		return met
	}
	switch strings.TrimSpace(sfx) {
	case "ns/op":
		s := time.Duration(num).String()
		i := 0
		for i < len(s) {
			c := s[i]
			if (c >= '0' && c <= '9') || c == '.' {
				i++
			} else {
				break
			}
		}
		return pfx + s[:i] + " " + s[i:] + "/op"
	case "B/op":
		return pfx + humanize.IBytes(uint64(num)) + "/op"
	default:
		return pfx + humanize.Comma(num) + sfx
	}
}

type humanizeWriter struct {
	mg.OutputStream
}

func (w *humanizeWriter) Write(ln []byte) (int, error) {
	s := make([]byte, 0, len(ln)+42)
	for len(ln) != 0 {
		i := bytes.IndexByte(ln, '\t')
		if i < 0 {
			s = append(s, humanizeMetric(string(ln))...)
			break
		}
		i++
		s = append(s, humanizeMetric(string(ln[:i]))...)
		ln = ln[i:]
	}
	return w.OutputStream.Write(s)
}

type replWriter struct {
	mg.OutputStream
	old []*regexp.Regexp
	new [][]byte
}

func (w *replWriter) Write(ln []byte) (int, error) {
	for i, pat := range w.old {
		ln = pat.ReplaceAll(ln, w.new[i])
	}
	return w.OutputStream.Write(ln)
}
<commit_msg>sync margo<commit_after>package golang

import (
	"bytes"
	"fmt"
	"github.com/dustin/go-humanize"
	"go/ast"
	"go/build"
	"io/ioutil"
	"margo.sh/golang/cursor"
	"margo.sh/mg"
	"margo.sh/mgutil"
	"os"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"time"
)

type GoCmd struct {
	mg.ReducerType

	Humanize bool
}

func (gc *GoCmd) Reduce(mx *mg.Ctx) *mg.State {
	switch act := mx.Action.(type) {
	case mg.QueryUserCmds:
		return gc.userCmds(mx)
	case mg.RunCmd:
		return gc.runCmd(mx, act)
	}
	return mx.State
}

func (gc *GoCmd) userCmds(mx *mg.Ctx) *mg.State {
	return mx.AddUserCmds(
		mg.UserCmd{
			Name:  "go.play",
			Title: "Go Play",
		},
		mg.UserCmd{
			Name:  "go.replay",
			Title: "Go RePlay (single instance)",
		},
	)
}

func (gc *GoCmd) runCmd(mx *mg.Ctx, rc mg.RunCmd) *mg.State {
	return mx.State.AddBuiltinCmds(
		mg.BuiltinCmd{
			Run:  gc.goBuiltin,
			Name: "go",
			Desc: "Wrapper around the go command, adding linter support",
		},
		mg.BuiltinCmd{
			Run:  gc.playBuiltin,
			Name: "go.play",
			Desc: "Automatically build and run go commands or run go test for packages with support for linting and unsaved files",
		},
		mg.BuiltinCmd{
			Run:  gc.replayBuiltin,
			Name: "go.replay",
			Desc: "Wrapper around go.play limited to a single instance",
		},
	)
}

func (gc *GoCmd) goBuiltin(bx *mg.CmdCtx) *mg.State {
	go gc.goTool(bx)
	return bx.State
}

func (gc *GoCmd) playBuiltin(bx *mg.CmdCtx) *mg.State {
	go gc.playTool(bx, "")
	return bx.State
}

func (gc *GoCmd) replayBuiltin(bx *mg.CmdCtx) *mg.State {
	v := bx.View
	cid := ""
	if v.Path == "" {
		cid = v.Name
	} else {
		cid = v.Dir()
	}
	go gc.playTool(bx, "go.replay`"+cid+"`")
	return bx.State
}

func (gc *GoCmd) goTool(bx *mg.CmdCtx) {
	gx := newGoCmdCtx(gc, bx, "go.builtin", "", "", "", bx.View, len(bx.Args) > 0 && bx.Args[0] == "test")
	defer gx.Output.Close()
	gx.run(gx.View)
}

func (gc *GoCmd) playTool(bx *mg.CmdCtx, cancelID string) {
	bld := BuildContext(bx.Ctx)
	pkg, err := bld.ImportDir(bx.View.Dir(), 0)
	if err != nil {
		fmt.Fprintln(bx.Output, "Error: cannot import package:", err)
	}

	testMode := !pkg.IsCommand() ||
		strings.HasSuffix(bx.View.Filename(), "_test.go")

	origView := bx.View
	bx, tDir, tFn, err := gc.playTempDir(bx)
	if err != nil {
		fmt.Fprintf(bx.Output, "Error: %s\n", err)
	}
	defer os.RemoveAll(tDir)
	if tDir == "" {
		return
	}
	gx := newGoCmdCtx(gc, bx, "go.play", cancelID, tDir, tFn, origView, testMode)
	defer gx.Output.Close()

	gx.Verbose = true

	switch {
	case testMode:
		gc.playToolTest(gx, bld, origView)
	default:
		gc.playToolRun(gx, bld, origView)
	}
}

func (gc *GoCmd) playTempDir(bx *mg.CmdCtx) (newBx *mg.CmdCtx, tDir string, tFn string, err error) {
	tDir, err = mg.MkTempDir("go.play")
	if err != nil {
		return bx, "", "", fmt.Errorf("cannot MkTempDir: %s", err)
	}

	if !bx.LangIs(mg.Go) {
		return bx, tDir, "", nil
	}

	v := bx.View
	if v.Path != "" {
		return bx, tDir, tFn, nil
	}

	tFn = filepath.Join(tDir, v.Name)
	src, err := v.ReadAll()
	if err == nil {
		err = ioutil.WriteFile(tFn, src, 0600)
	}
	if err != nil {
		return bx, tDir, "", fmt.Errorf("cannot create temp file: %s", err)
	}

	bx = bx.Copy(func(bx *mg.CmdCtx) {
		bx.Ctx = bx.Ctx.Copy(func(mx *mg.Ctx) {
			mx.State = mx.State.Copy(func(st *mg.State) {
				st.View = st.View.Copy(func(v *mg.View) {
					v.Path = tFn
				})
			})
		})
	})

	return bx, tDir, tFn, nil
}

func (gc *GoCmd) playToolTest(gx *goCmdCtx, bld *build.Context, origView *mg.View) {
	argsPfx := []string{"test", "-test.run=."}
	cx := cursor.NewViewCurCtx(gx.Ctx)
	for _, n := range cx.Nodes {
		x, ok := n.(*ast.FuncDecl)
		if !ok || x.Name == nil {
			continue
		}
		nm := x.Name.String()
		if strings.HasPrefix(nm, "Benchmark") {
			argsPfx = append(argsPfx, "-test.bench=^"+nm+"$")
		}
	}
	gx.Args = append(argsPfx, gx.Args...)
	if origView.Path == "" {
		gx.Args = append(gx.Args, gx.tFn)
	}
	gx.run(origView)
}

func (gc *GoCmd) playToolRun(gx *goCmdCtx, bld *build.Context, origView *mg.View) {
	nm := filepath.Base(origView.Name)
	if origView.Path != "" {
		nm = filepath.Base(origView.Dir())
	}

	args := gx.Args
	exe := filepath.Join(gx.tDir, nm+".exe")
	gx.CmdCtx = gx.CmdCtx.Copy(func(bx *mg.CmdCtx) {
		bx.Name = "go"
		bx.Args = []string{"build", "-o", exe}
		bx.Ctx = bx.Ctx.Copy(func(mx *mg.Ctx) {
			mx.State = mx.State.Copy(func(st *mg.State) {
				st.View = st.View.Copy(func(v *mg.View) {
					v.Wd = v.Dir()
				})
			})
		})
	})
	if err := gx.run(origView); err != nil {
		return
	}

	gx.CmdCtx = gx.CmdCtx.Copy(func(bx *mg.CmdCtx) {
		bx.Name = exe
		bx.Args = args
		bx.Ctx = bx.Ctx.Copy(func(mx *mg.Ctx) {
			mx.State = mx.State.Copy(func(st *mg.State) {
				st.View = origView
			})
		})
	})
	gx.RunProc()
}

type goCmdCtx struct {
	*mg.CmdCtx
	pkgDir string
	key    interface{}
	iw     *mg.IssueOut
	tDir   string
	tFn    string
}

func newGoCmdCtx(gc *GoCmd, bx *mg.CmdCtx, label, cancelID string, tDir, tFn string, origView *mg.View, testMode bool) *goCmdCtx {
	gx := &goCmdCtx{
		pkgDir: bx.View.Dir(),
		tDir:   tDir,
		tFn:    tFn,
	}

	output := bx.Output
	if gc.Humanize && testMode {
		output = &humanizeWriter{output}
	}
	if gx.tFn != "" {
		dir := filepath.Dir(gx.tFn)
		qDir := regexp.QuoteMeta(dir)
		qDirBase := regexp.QuoteMeta(filepath.Base(dir))
		qNm := regexp.QuoteMeta(filepath.Base(gx.tFn))
		output = &replWriter{
			OutputStream: output,
			old: []*regexp.Regexp{
				regexp.MustCompile(`(?:` + qDir + `|` + qDirBase + `)?[\\/.]+` + qNm),
				regexp.MustCompile(qDir),
			},
			new: [][]byte{
				[]byte(origView.Name),
				[]byte(`tmp~`),
			},
		}
	}
	output = mgutil.NewSplitStream(mgutil.SplitLineOrCR, output)

	type Key struct{ label string }
	gx.key = Key{label}

	gx.iw = &mg.IssueOut{
		Base:     mg.Issue{Label: label},
		Patterns: bx.CommonPatterns(),
		Dir:      gx.pkgDir,
	}

	gx.CmdCtx = bx.Copy(func(bx *mg.CmdCtx) {
		bx.Name = "go"
		bx.CancelID = cancelID
		bx.Output = mg.OutputStreams{
			output,
			gx.iw,
		}
	})

	return gx
}

func (gx *goCmdCtx) run(origView *mg.View) error {
	defer func() {
		gx.VFS.Invalidate(origView.Filename())
		gx.VFS.Invalidate(origView.Dir())
	}()

	p, err := gx.StartProc()
	if err == nil {
		err = p.Wait()
	}
	gx.iw.Flush()

	issues := gx.iw.Issues()

	for i, isu := range issues {
		if isu.Path == "" || (gx.tFn != "" && filepath.Base(isu.Path) == origView.Name) {
			isu.Name = origView.Name
			isu.Path = origView.Path
		}
		issues[i] = isu
	}

	ik := mg.IssueKey{Key: gx.key}
	if origView.Path == "" {
		ik.Name = origView.Name
	} else {
		ik.Dir = origView.Dir()
	}

	gx.Store.Dispatch(mg.StoreIssues{IssueKey: ik, Issues: issues})
	return err
}

func isWhiteSpace(c byte) bool {
	return c == ' ' || c == '\t' || c == '\n' || c == '\r'
}

func humanizeMetric(met string) string {
	i := 0
	for i < len(met) && isWhiteSpace(met[i]) {
		i++
	}
	j := i
	for j < len(met) && !isWhiteSpace(met[j]) {
		j++
	}
	k := len(met)
	for k > j && isWhiteSpace(met[k-1]) {
		k--
	}
	pfx := met[:i]
	val := met[i:j]
	unit := met[j:k]
	sfx := met[k:]

	num, err := strconv.ParseInt(val, 10, 64)
	if err != nil {
		return met
	}
	switch strings.TrimSpace(unit) {
	case "ns/op":
		s := time.Duration(num).String()
		i := 0
		for i < len(s) {
			c := s[i]
			if (c >= '0' && c <= '9') || c == '.' {
				i++
			} else {
				break
			}
		}
		return pfx + s[:i] + " " + s[i:] + "/op" + sfx
	case "B/op":
		return pfx + humanize.IBytes(uint64(num)) + "/op" + sfx
	default:
		return pfx + humanize.Comma(num) + unit + sfx
	}
}

type humanizeWriter struct {
	mg.OutputStream
}

func (w *humanizeWriter) Write(ln []byte) (int, error) {
	s := make([]byte, 0, len(ln)+42)
	for len(ln) != 0 {
		i := bytes.IndexByte(ln, '\t')
		if i < 0 {
			s = append(s, humanizeMetric(string(ln))...)
			break
		}
		i++
		s = append(s, humanizeMetric(string(ln[:i]))...)
		ln = ln[i:]
	}
	return w.OutputStream.Write(s)
}

type replWriter struct {
	mg.OutputStream
	old []*regexp.Regexp
	new [][]byte
}

func (w *replWriter) Write(ln []byte) (int, error) {
	for i, pat := range w.old {
		ln = pat.ReplaceAll(ln, w.new[i])
	}
	return w.OutputStream.Write(ln)
}
<|endoftext|>