<commit_before>// Copyright 2020 The OPA Authors.  All rights reserved.
// Use of this source code is governed by an Apache2
// license that can be found in the LICENSE file.

package opa

import (
	"bytes"
	"context"
	"fmt"
	"sync"

	"github.com/open-policy-agent/opa/metrics"
)

// pool maintains a pool of WebAssemly VM instances.
type pool struct {
	available      chan struct{}
	mutex          sync.Mutex
	initialized    bool
	closed         bool
	policy         []byte
	parsedData     []byte // Parsed parsedData memory segment, used to seed new VM's
	parsedDataAddr int32  // Address for parsedData value root, used to seed new VM's
	memoryMinPages uint32
	memoryMaxPages uint32
	vms            []*vm // All current VM instances, acquired or not.
	acquired       []bool
	pendingReinit  *vm
	blockedReinit  chan struct{}
}

// newPool constructs a new pool with the pool and VM configuration provided.
func newPool(poolSize, memoryMinPages, memoryMaxPages uint32) *pool {
	available := make(chan struct{}, poolSize)
	for i := uint32(0); i < poolSize; i++ {
		available <- struct{}{}
	}

	return &pool{
		memoryMinPages: memoryMinPages,
		memoryMaxPages: memoryMaxPages,
		available:      available,
		vms:            make([]*vm, 0),
		acquired:       make([]bool, 0),
	}
}

// Acquire obtains a VM from the pool, waiting if all VMms are in use
// and building one as necessary. Returns either ErrNotReady or
// ErrInternal if an error.
func (p *pool) Acquire(ctx context.Context, metrics metrics.Metrics) (*vm, error) {
	metrics.Timer("wasm_pool_acquire").Start()
	defer metrics.Timer("wasm_pool_acquire").Stop()

	select {
	case <-ctx.Done():
		return nil, ctx.Err()
	case <-p.available:
	}

	p.mutex.Lock()
	defer p.mutex.Unlock()

	if !p.initialized || p.closed {
		return nil, ErrNotReady
	}

	for i, vm := range p.vms {
		if !p.acquired[i] {
			p.acquired[i] = true
			return vm, nil
		}
	}

	policy, parsedData, parsedDataAddr := p.policy, p.parsedData, p.parsedDataAddr

	p.mutex.Unlock()
	vm, err := newVM(vmOpts{
		policy:         policy,
		data:           nil,
		parsedData:     parsedData,
		parsedDataAddr: parsedDataAddr,
		memoryMin:      p.memoryMinPages,
		memoryMax:      p.memoryMaxPages,
	})
	p.mutex.Lock()

	if err != nil {
		p.available <- struct{}{}
		return nil, fmt.Errorf("%v: %w", err, ErrInternal)
	}

	p.acquired = append(p.acquired, true)
	p.vms = append(p.vms, vm)
	return vm, nil
}

// Release releases the VM back to the pool.
func (p *pool) Release(vm *vm, metrics metrics.Metrics) {
	metrics.Timer("wasm_pool_release").Start()
	defer metrics.Timer("wasm_pool_release").Stop()

	p.mutex.Lock()

	// If the policy data setting is waiting for this one, don't release it back to the general consumption.
	// Note the reinit is responsible for pushing to available channel once done with the VM.
	if vm == p.pendingReinit {
		p.mutex.Unlock()
		p.blockedReinit <- struct{}{}
		return
	}

	for i := range p.vms {
		if p.vms[i] == vm {
			p.acquired[i] = false
			p.mutex.Unlock()
			p.available <- struct{}{}
			return
		}
	}

	// VM instance not found anymore, hence pool reconfigured and can release the VM.

	p.mutex.Unlock()
	p.available <- struct{}{}

	vm.Close()
}

// SetPolicyData re-initializes the vms within the pool with the new policy
// and data. The re-initialization takes place atomically: all new vms
// are constructed in advance before touching the pool.  Returns
// either ErrNotReady, ErrInvalidPolicy or ErrInternal if an error
// occurs.
func (p *pool) SetPolicyData(policy []byte, data []byte) error {
	p.mutex.Lock()

	if !p.initialized {
		var err error
		var parsedData []byte
		var parsedDataAddr int32
		for i := 0; i < len(p.available); i++ {
			var vm *vm
			vm, err = newVM(vmOpts{
				policy:         policy,
				data:           data,
				parsedData:     parsedData,
				parsedDataAddr: parsedDataAddr,
				memoryMin:      p.memoryMinPages,
				memoryMax:      p.memoryMaxPages,
			})

			if err == nil {
				if parsedData == nil {
					parsedDataAddr, parsedData = vm.cloneDataSegment()
					p.memoryMinPages = pages(vm.memory.Length())
				}
				p.vms = append(p.vms, vm)
				p.acquired = append(p.acquired, false)
			} else {
				err = fmt.Errorf("%v: %w", err, ErrInvalidPolicyOrData)
				break
			}
		}

		if err == nil {
			p.initialized = true
			p.policy, p.parsedData, p.parsedDataAddr = policy, parsedData, parsedDataAddr
		}

		p.mutex.Unlock()
		return err
	}

	if p.closed {<commit_msg>internal/wasm/sdk: Revert full pool initialization<commit_after>// Copyright 2020 The OPA Authors.  All rights reserved.
// Use of this source code is governed by an Apache2
// license that can be found in the LICENSE file.

package opa

import (
	"bytes"
	"context"
	"fmt"
	"sync"

	"github.com/open-policy-agent/opa/metrics"
)

// pool maintains a pool of WebAssemly VM instances.
type pool struct {
	available      chan struct{}
	mutex          sync.Mutex
	initialized    bool
	closed         bool
	policy         []byte
	parsedData     []byte // Parsed parsedData memory segment, used to seed new VM's
	parsedDataAddr int32  // Address for parsedData value root, used to seed new VM's
	memoryMinPages uint32
	memoryMaxPages uint32
	vms            []*vm // All current VM instances, acquired or not.
	acquired       []bool
	pendingReinit  *vm
	blockedReinit  chan struct{}
}

// newPool constructs a new pool with the pool and VM configuration provided.
func newPool(poolSize, memoryMinPages, memoryMaxPages uint32) *pool {
	available := make(chan struct{}, poolSize)
	for i := uint32(0); i < poolSize; i++ {
		available <- struct{}{}
	}

	return &pool{
		memoryMinPages: memoryMinPages,
		memoryMaxPages: memoryMaxPages,
		available:      available,
		vms:            make([]*vm, 0),
		acquired:       make([]bool, 0),
	}
}

// Acquire obtains a VM from the pool, waiting if all VMms are in use
// and building one as necessary. Returns either ErrNotReady or
// ErrInternal if an error.
func (p *pool) Acquire(ctx context.Context, metrics metrics.Metrics) (*vm, error) {
	metrics.Timer("wasm_pool_acquire").Start()
	defer metrics.Timer("wasm_pool_acquire").Stop()

	select {
	case <-ctx.Done():
		return nil, ctx.Err()
	case <-p.available:
	}

	p.mutex.Lock()
	defer p.mutex.Unlock()

	if !p.initialized || p.closed {
		return nil, ErrNotReady
	}

	for i, vm := range p.vms {
		if !p.acquired[i] {
			p.acquired[i] = true
			return vm, nil
		}
	}

	policy, parsedData, parsedDataAddr := p.policy, p.parsedData, p.parsedDataAddr

	p.mutex.Unlock()
	vm, err := newVM(vmOpts{
		policy:         policy,
		data:           nil,
		parsedData:     parsedData,
		parsedDataAddr: parsedDataAddr,
		memoryMin:      p.memoryMinPages,
		memoryMax:      p.memoryMaxPages,
	})
	p.mutex.Lock()

	if err != nil {
		p.available <- struct{}{}
		return nil, fmt.Errorf("%v: %w", err, ErrInternal)
	}

	p.acquired = append(p.acquired, true)
	p.vms = append(p.vms, vm)
	return vm, nil
}

// Release releases the VM back to the pool.
func (p *pool) Release(vm *vm, metrics metrics.Metrics) {
	metrics.Timer("wasm_pool_release").Start()
	defer metrics.Timer("wasm_pool_release").Stop()

	p.mutex.Lock()

	// If the policy data setting is waiting for this one, don't release it back to the general consumption.
	// Note the reinit is responsible for pushing to available channel once done with the VM.
	if vm == p.pendingReinit {
		p.mutex.Unlock()
		p.blockedReinit <- struct{}{}
		return
	}

	for i := range p.vms {
		if p.vms[i] == vm {
			p.acquired[i] = false
			p.mutex.Unlock()
			p.available <- struct{}{}
			return
		}
	}

	// VM instance not found anymore, hence pool reconfigured and can release the VM.

	p.mutex.Unlock()
	p.available <- struct{}{}

	vm.Close()
}

// SetPolicyData re-initializes the vms within the pool with the new policy
// and data. The re-initialization takes place atomically: all new vms
// are constructed in advance before touching the pool.  Returns
// either ErrNotReady, ErrInvalidPolicy or ErrInternal if an error
// occurs.
func (p *pool) SetPolicyData(policy []byte, data []byte) error {
	p.mutex.Lock()

	if !p.initialized {
		vm, err := newVM(vmOpts{
			policy:         policy,
			data:           data,
			parsedData:     nil,
			parsedDataAddr: 0,
			memoryMin:      p.memoryMinPages,
			memoryMax:      p.memoryMaxPages,
		})

		if err == nil {
			parsedDataAddr, parsedData := vm.cloneDataSegment()
			p.memoryMinPages = pages(vm.memory.Length())
			p.vms = append(p.vms, vm)
			p.acquired = append(p.acquired, false)
			p.initialized = true
			p.policy, p.parsedData, p.parsedDataAddr = policy, parsedData, parsedDataAddr
		} else {
			err = fmt.Errorf("%v: %w", err, ErrInvalidPolicyOrData)
		}

		p.mutex.Unlock()
		return err
	}

	if p.closed {
{'fraction_non_alphanumeric': 0.11688014638609331, 'fraction_numerical': 0.004574565416285453, 'mean_word_length': 3.080727951469902, 'pattern_counts': {'":': 0, '<': 20, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6174936', 'n_tokens_mistral': 3148, 'n_tokens_neox': 2899, 'n_words': 915}