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

package mapper

import (
	"fmt"
	"sync"

	"github.com/golang/protobuf/proto"
	"golang.org/x/net/context"

	"go.chromium.org/gae/service/datastore"

	"go.chromium.org/luci/common/clock"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/logging"
	"go.chromium.org/luci/common/retry/transient"

	"go.chromium.org/luci/appengine/tq"

	"go.chromium.org/luci/appengine/mapper/internal/tasks"
	"go.chromium.org/luci/appengine/mapper/splitter"
)

// ID identifies a mapper.
type ID string

// Mapper knows how to apply some function to a slice of entities.
//
// Mappers are supplied by the users of the library and registered in the
// controller via RegisterMapper call.
type Mapper interface {
	// MapperID uniquely identifies this mapper.
	//
	// It will be used as identifier of the mapper across process boundaries.
	MapperID() ID

	// Process applies some function to the given slice of entities, given by
	// their keys.
	Process(c context.Context, p Params, keys []*datastore.Key) error
}

// Controller is responsible for starting, progressing and finishing mapping
// jobs.
//
// It should be treated as a global singleton object. There's rarely a need to
// have more than one controller per application.
type Controller struct {
	// MapperQueue is a name of the GAE task queue to use for mapping jobs.
	//
	// This queue will perform all "heavy" tasks. It should be configured
	// appropriately to allow desired number of tasks to run in parallel.
	//
	// If empty, "default" is used.
	MapperQueue string

	// ControlQueue is a name of the GAE task queue to use for control signals.
	//
	// This queue is used very lightly when starting and stopping jobs.
	//
	// If empty, "default" is used.
	ControlQueue string

	l       sync.RWMutex
	mappers map[ID]Mapper
	disp    *tq.Dispatcher

	// Called in splitAndLaunchHandler during tests.
	testingRecordSplit func(rng splitter.Range)
}

// Install registers task queue task handlers in the given task queue
// dispatcher.
//
// This must be done before Controller is used.
//
// There can be at most one Controller installed into an instance of TQ
// dispatcher. Installing more will cause panics.
//
// If you need multiple different controllers for some reason, create multiple
// tq.Dispatchers (with different base URLs, so they don't conflict with each
// other) and install them all into the router.
func (ctl *Controller) Install(disp *tq.Dispatcher) {
	ctl.l.Lock()
	defer ctl.l.Unlock()

	if ctl.disp != nil {
		panic("mapper.Controller is already installed into a tq.Dispatcher")
	}
	ctl.disp = disp

	disp.RegisterTask(&tasks.SplitAndLaunch{}, ctl.splitAndLaunchHandler, ctl.ControlQueue, nil)
}

// tq returns a dispatcher set in Install or panics if not set yet.
//
// Grabs the reader lock inside.
func (ctl *Controller) tq() *tq.Dispatcher {
	ctl.l.RLock()
	defer ctl.l.RUnlock()
	if ctl.disp == nil {
		panic("mapper.Controller wasn't installed into tq.Dispatcher yet")
	}
	return ctl.disp
}

// RegisterMapper adds the given mapper to the internal registry.
//
// Intended to be used during init() time or early during the process
// initialization. Panics if a mapper with such ID has already been registered.
//
// The mapper ID will be used internally to identify which mapper a job should
// be using. If a mapper disappears while the job is running (e.g. if the
// service binary is updated and new binary doesn't have the mapper registered
// anymore), the job ends with a failure.
func (ctl *Controller) RegisterMapper(m Mapper) {
	id := m.MapperID()

	ctl.l.Lock()
	defer ctl.l.Unlock()

	if _, ok := ctl.mappers[id]; ok {
		panic(fmt.Sprintf("mapper %q is already registered", id))
	}

	if ctl.mappers == nil {
		ctl.mappers = make(map[ID]Mapper, 1)
	}
	ctl.mappers[id] = m
}

// getMapper returns a registered mapper or an error.
//
// Grabs the reader lock inside.
func (ctl *Controller) getMapper(id ID) (Mapper, error) {
	ctl.l.RLock()
	defer ctl.l.RUnlock()
	if m, ok := ctl.mappers[id]; ok {
		return m, nil
	}
	return nil, errors.Reason("no mapper with ID %q registered", id).Err()
}

// LaunchJob launches a new mapping job, returning its ID (that can be used to
// control it or query its status).
//
// Launches a datastore transaction inside.
func (ctl *Controller) LaunchJob(c context.Context, j *JobConfig) (JobID, error) {
	disp := ctl.tq()<commit_msg>[mapper] Reply to comments left in the previous CL after it was submitted.<commit_after>// Copyright 2018 The LUCI Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package mapper

import (
	"fmt"
	"sync"

	"github.com/golang/protobuf/proto"
	"golang.org/x/net/context"

	"go.chromium.org/gae/service/datastore"

	"go.chromium.org/luci/common/clock"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/logging"
	"go.chromium.org/luci/common/retry/transient"

	"go.chromium.org/luci/appengine/tq"

	"go.chromium.org/luci/appengine/mapper/internal/tasks"
	"go.chromium.org/luci/appengine/mapper/splitter"
)

// ID identifies a mapper registered in the controller.
//
// It will be passed across processes, so all processes that execute mapper jobs
// should register same mappers under same IDs.
//
// The safest approach is to keep mapper IDs in the app unique, e.g. do NOT
// reuse them when adding new mappers or significantly changing existing ones.
type ID string

// Mapper knows how to apply some function to a slice of entities.
//
// Mappers are supplied by the users of the library and registered in the
// controller via RegisterMapper call.
type Mapper interface {
	// MapperID uniquely identifies this mapper.
	//
	// It will be used as identifier of the mapper across process boundaries.
	MapperID() ID

	// Process applies some function to the given slice of entities, given by
	// their keys.
	//
	// May be called multiple times for same key (thus should be idempotent).
	//
	// Returning a transient error indicates that the processing of this batch of
	// keys should be retried (even if some keys were processed successfully).
	//
	// Returning a fatal error causes the entire shard (and eventually the entire
	// job) to be marked as failed. The processing of the failed shard stops right
	// away, but other shards are kept running until completion (or their own
	// failure).
	//
	// The function is called outside of any transactions, so it can start its own
	// if needed.
	Process(c context.Context, p Params, keys []*datastore.Key) error
}

// Controller is responsible for starting, progressing and finishing mapping
// jobs.
//
// It should be treated as a global singleton object. Having more than one
// controller in the production application is a bad idea (they'll collide with
// each other since they use global datastore namespace). It's still useful
// to instantiate multiple controllers in unit tests.
type Controller struct {
	// MapperQueue is a name of the GAE task queue to use for mapping jobs.
	//
	// This queue will perform all "heavy" tasks. It should be configured
	// appropriately to allow desired number of shards to run in parallel.
	//
	// For example, if the largest submitted job is expected to have 128 shards,
	// max_concurrent_requests setting of the mapper queue should be at least 128,
	// otherwise some shards will be stalled waiting for others to finish
	// (defeating the purpose of having large number of shards).
	//
	// If empty, "default" is used.
	MapperQueue string

	// ControlQueue is a name of the GAE task queue to use for control signals.
	//
	// This queue is used very lightly when starting and stopping jobs (roughly
	// 2*Shards tasks overall per job). A default queue.yaml settings for such
	// queue should be sufficient (unless you run a lot of different jobs at
	// once).
	//
	// If empty, "default" is used.
	ControlQueue string

	m       sync.RWMutex
	mappers map[ID]Mapper
	disp    *tq.Dispatcher

	// Called in splitAndLaunchHandler during tests.
	testingRecordSplit func(rng splitter.Range)
}

// Install registers task queue task handlers in the given task queue
// dispatcher.
//
// This must be done before Controller is used.
//
// There can be at most one Controller installed into an instance of TQ
// dispatcher. Installing more will cause panics.
//
// If you need multiple different controllers for some reason, create multiple
// tq.Dispatchers (with different base URLs, so they don't conflict with each
// other) and install them all into the router.
func (ctl *Controller) Install(disp *tq.Dispatcher) {
	ctl.m.Lock()
	defer ctl.m.Unlock()

	if ctl.disp != nil {
		panic("mapper.Controller is already installed into a tq.Dispatcher")
	}
	ctl.disp = disp

	disp.RegisterTask(&tasks.SplitAndLaunch{}, ctl.splitAndLaunchHandler, ctl.ControlQueue, nil)
}

// tq returns a dispatcher set in Install or panics if not set yet.
//
// Grabs the reader lock inside.
func (ctl *Controller) tq() *tq.Dispatcher {
	ctl.m.RLock()
	defer ctl.m.RUnlock()
	if ctl.disp == nil {
		panic("mapper.Controller wasn't installed into tq.Dispatcher yet")
	}
	return ctl.disp
}

// RegisterMapper adds the given mapper to the internal registry.
//
// Intended to be used during init() time or early during the process
// initialization. Panics if a mapper with such ID has already been registered.
//
// The mapper ID will be used internally to identify which mapper a job should
// be using. If a mapper disappears while the job is running (e.g. if the
// service binary is updated and new binary doesn't have the mapper registered
// anymore), the job ends with a failure.
func (ctl *Controller) RegisterMapper(m Mapper) {
	id := m.MapperID()

	ctl.m.Lock()
	defer ctl.m.Unlock()

	if _, ok := ctl.mappers[id]; ok {
		panic(fmt.Sprintf("mapper %q is already registered", id))
	}

	if ctl.mappers == nil {
		ctl.mappers = make(map[ID]Mapper, 1)
	}
	ctl.mappers[id] = m
}

// getMapper returns a registered mapper or an error.
//
// Grabs the reader lock inside.
func (ctl *Controller) getMapper(id ID) (Mapper, error) {
	ctl.m.RLock()
	defer ctl.m.RUnlock()
	if m, ok := ctl.mappers[id]; ok {
		return m, nil
	}
	return nil, errors.Reason("no mapper with ID %q registered", id).Err()
}

// LaunchJob launches a new mapping job, returning its ID (that can be used to
// control it or query its status).
//
// Launches a datastore transaction inside.
func (ctl *Controller) LaunchJob(c context.Context, j *JobConfig) (JobID, error) {
	disp := ctl.tq()
{'fraction_non_alphanumeric': 0.09736431420149125, 'fraction_numerical': 0.0021675047685104907, 'mean_word_length': 3.4382454790303965, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, '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': '5272580', 'n_tokens_mistral': 3522, 'n_tokens_neox': 3387, 'n_words': 1473}