<commit_before>// Licensed to the Apache Software Foundation (ASF) under one or more
// contributor license agreements.  See the NOTICE file distributed with
// this work for additional information regarding copyright ownership.
// The ASF licenses this file to You 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 harness implements the SDK side of the Beam FnAPI.
package harness

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"io/ioutil"
	"runtime/pprof"
	"sync"
	"time"

	"github.com/apache/beam/sdks/go/pkg/beam/core/runtime/exec"
	"github.com/apache/beam/sdks/go/pkg/beam/core/runtime/graphx"
	"github.com/apache/beam/sdks/go/pkg/beam/log"
	fnpb "github.com/apache/beam/sdks/go/pkg/beam/model/fnexecution_v1"
	pb "github.com/apache/beam/sdks/go/pkg/beam/model/pipeline_v1"
	"github.com/apache/beam/sdks/go/pkg/beam/runners/direct"
	"github.com/apache/beam/sdks/go/pkg/beam/util/grpcx"
	"github.com/golang/protobuf/proto"
	"google.golang.org/grpc"
)

// TODO(herohde) 2/8/2017: for now, assume we stage a full binary (not a plugin).

// Main is the main entrypoint for the Go harness. It runs at "runtime" -- not
// "pipeline-construction time" -- on each worker. It is a Fn API client and
// ultimately responsible for correctly executing user code.
func Main(ctx context.Context, loggingEndpoint, controlEndpoint string) error {
	setupRemoteLogging(ctx, loggingEndpoint)
	setupDiagnosticRecording()

	recordHeader()

	// Connect to FnAPI control server. Receive and execute work.
	// TODO: setup data manager, DoFn register

	conn, err := dial(ctx, controlEndpoint, 60*time.Second)
	if err != nil {
		return fmt.Errorf("Failed to connect: %v", err)
	}
	defer conn.Close()

	client, err := fnpb.NewBeamFnControlClient(conn).Control(ctx)
	if err != nil {
		return fmt.Errorf("Failed to connect to control service: %v", err)
	}

	log.Debugf(ctx, "Successfully connected to control @ %v", controlEndpoint)

	// Each ProcessBundle is a sub-graph of the original one.

	var wg sync.WaitGroup
	respc := make(chan *fnpb.InstructionResponse, 100)

	wg.Add(1)

	// gRPC requires all writers to a stream be the same goroutine, so this is the
	// goroutine for managing responses back to the control service.
	go func() {
		defer wg.Done()
		for resp := range respc {
			log.Debugf(ctx, "RESP: %v", proto.MarshalTextString(resp))

			if err := client.Send(resp); err != nil {
				log.Errorf(ctx, "Failed to respond: %v", err)
			}
		}
	}()

	ctrl := &control{
		plans: make(map[string]*exec.Plan),
		data:  &DataManager{},
	}

	var cpuProfBuf bytes.Buffer
	for {
		req, err := client.Recv()
		if err != nil {
			close(respc)
			wg.Wait()

			if err == io.EOF {
				recordFooter()
				return nil
			}
			return fmt.Errorf("recv failed: %v", err)
		}

		// Launch a goroutine to handle the control message.
		// TODO(wcn): maybe only do this for 'heavy' messages?
		// TODO(wcn): implement a rate limiter for 'heavy' messages?
		go func() {
			log.Debugf(ctx, "RECV: %v", proto.MarshalTextString(req))
			recordInstructionRequest(req)

			if isEnabled("cpu_profiling") {
				cpuProfBuf.Reset()
				pprof.StartCPUProfile(&cpuProfBuf)
			}
			resp := ctrl.handleInstruction(ctx, req)

			if isEnabled("cpu_profiling") {
				pprof.StopCPUProfile()
				if err := ioutil.WriteFile(fmt.Sprintf("%s/cpu_prof%s", storagePath, req.InstructionId), cpuProfBuf.Bytes(), 0644); err != nil {
					log.Warnf(ctx, "Failed to write CPU profile for instruction %s: %v", req.InstructionId, err)
				}
			}

			recordInstructionResponse(resp)
			if resp != nil {
				respc <- resp
			}
		}()
	}
}

type control struct {
	plans  map[string]*exec.Plan // protected by mu
	active map[string]*exec.Plan // protected by mu
	mu     sync.Mutex
	data   *DataManager
}

func (c *control) handleInstruction(ctx context.Context, req *fnpb.InstructionRequest) *fnpb.InstructionResponse {
	id := req.GetInstructionId()
	ctx = context.WithValue(ctx, instKey, id)

	switch {
	case req.GetRegister() != nil:
		msg := req.GetRegister()

		for _, desc := range msg.GetProcessBundleDescriptor() {
			var roots []string
			for id := range desc.GetTransforms() {
				roots = append(roots, id)
			}
			p := &pb.Pipeline{
				Components: &pb.Components{
					Transforms:          desc.Transforms,
					Pcollections:        desc.Pcollections,
					Coders:              desc.Coders,
					WindowingStrategies: desc.WindowingStrategies,
					Environments:        desc.Environments,
				},
				RootTransformIds: roots,
			}

			g, err := graphx.Unmarshal(p)
			if err != nil {
				return fail(id, "Invalid bundle desc: %v", err)
			}
			log.Debugf(ctx, "Added subgraph %v:\n %v", desc.GetId(), g)

			edges, _, err := g.Build()
			if err != nil {
				return fail(id, "invalid pipeline: %v", err)
			}
			plan, err := direct.Compile(edges)
			if err != nil {
				return fail(id, "translation failed: %v", err)
			}
			log.Debugf(ctx, "Plan %v: %v", desc.GetId(), plan)

			c.mu.Lock()
			c.plans[desc.GetId()] = plan
			c.mu.Unlock()
		}
<commit_msg>Remove GetId() call from under lock.<commit_after>// Licensed to the Apache Software Foundation (ASF) under one or more
// contributor license agreements.  See the NOTICE file distributed with
// this work for additional information regarding copyright ownership.
// The ASF licenses this file to You 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 harness implements the SDK side of the Beam FnAPI.
package harness

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"io/ioutil"
	"runtime/pprof"
	"sync"
	"time"

	"github.com/apache/beam/sdks/go/pkg/beam/core/runtime/exec"
	"github.com/apache/beam/sdks/go/pkg/beam/core/runtime/graphx"
	"github.com/apache/beam/sdks/go/pkg/beam/log"
	fnpb "github.com/apache/beam/sdks/go/pkg/beam/model/fnexecution_v1"
	pb "github.com/apache/beam/sdks/go/pkg/beam/model/pipeline_v1"
	"github.com/apache/beam/sdks/go/pkg/beam/runners/direct"
	"github.com/apache/beam/sdks/go/pkg/beam/util/grpcx"
	"github.com/golang/protobuf/proto"
	"google.golang.org/grpc"
)

// TODO(herohde) 2/8/2017: for now, assume we stage a full binary (not a plugin).

// Main is the main entrypoint for the Go harness. It runs at "runtime" -- not
// "pipeline-construction time" -- on each worker. It is a Fn API client and
// ultimately responsible for correctly executing user code.
func Main(ctx context.Context, loggingEndpoint, controlEndpoint string) error {
	setupRemoteLogging(ctx, loggingEndpoint)
	setupDiagnosticRecording()

	recordHeader()

	// Connect to FnAPI control server. Receive and execute work.
	// TODO: setup data manager, DoFn register

	conn, err := dial(ctx, controlEndpoint, 60*time.Second)
	if err != nil {
		return fmt.Errorf("Failed to connect: %v", err)
	}
	defer conn.Close()

	client, err := fnpb.NewBeamFnControlClient(conn).Control(ctx)
	if err != nil {
		return fmt.Errorf("Failed to connect to control service: %v", err)
	}

	log.Debugf(ctx, "Successfully connected to control @ %v", controlEndpoint)

	// Each ProcessBundle is a sub-graph of the original one.

	var wg sync.WaitGroup
	respc := make(chan *fnpb.InstructionResponse, 100)

	wg.Add(1)

	// gRPC requires all writers to a stream be the same goroutine, so this is the
	// goroutine for managing responses back to the control service.
	go func() {
		defer wg.Done()
		for resp := range respc {
			log.Debugf(ctx, "RESP: %v", proto.MarshalTextString(resp))

			if err := client.Send(resp); err != nil {
				log.Errorf(ctx, "Failed to respond: %v", err)
			}
		}
	}()

	ctrl := &control{
		plans: make(map[string]*exec.Plan),
		data:  &DataManager{},
	}

	var cpuProfBuf bytes.Buffer
	for {
		req, err := client.Recv()
		if err != nil {
			close(respc)
			wg.Wait()

			if err == io.EOF {
				recordFooter()
				return nil
			}
			return fmt.Errorf("recv failed: %v", err)
		}

		// Launch a goroutine to handle the control message.
		// TODO(wcn): maybe only do this for 'heavy' messages?
		// TODO(wcn): implement a rate limiter for 'heavy' messages?
		go func() {
			log.Debugf(ctx, "RECV: %v", proto.MarshalTextString(req))
			recordInstructionRequest(req)

			if isEnabled("cpu_profiling") {
				cpuProfBuf.Reset()
				pprof.StartCPUProfile(&cpuProfBuf)
			}
			resp := ctrl.handleInstruction(ctx, req)

			if isEnabled("cpu_profiling") {
				pprof.StopCPUProfile()
				if err := ioutil.WriteFile(fmt.Sprintf("%s/cpu_prof%s", storagePath, req.InstructionId), cpuProfBuf.Bytes(), 0644); err != nil {
					log.Warnf(ctx, "Failed to write CPU profile for instruction %s: %v", req.InstructionId, err)
				}
			}

			recordInstructionResponse(resp)
			if resp != nil {
				respc <- resp
			}
		}()
	}
}

type control struct {
	plans  map[string]*exec.Plan // protected by mu
	active map[string]*exec.Plan // protected by mu
	mu     sync.Mutex
	data   *DataManager
}

func (c *control) handleInstruction(ctx context.Context, req *fnpb.InstructionRequest) *fnpb.InstructionResponse {
	id := req.GetInstructionId()
	ctx = context.WithValue(ctx, instKey, id)

	switch {
	case req.GetRegister() != nil:
		msg := req.GetRegister()

		for _, desc := range msg.GetProcessBundleDescriptor() {
			var roots []string
			for id := range desc.GetTransforms() {
				roots = append(roots, id)
			}
			p := &pb.Pipeline{
				Components: &pb.Components{
					Transforms:          desc.Transforms,
					Pcollections:        desc.Pcollections,
					Coders:              desc.Coders,
					WindowingStrategies: desc.WindowingStrategies,
					Environments:        desc.Environments,
				},
				RootTransformIds: roots,
			}

			g, err := graphx.Unmarshal(p)
			if err != nil {
				return fail(id, "Invalid bundle desc: %v", err)
			}
			log.Debugf(ctx, "Added subgraph %v:\n %v", desc.GetId(), g)

			edges, _, err := g.Build()
			if err != nil {
				return fail(id, "invalid pipeline: %v", err)
			}
			plan, err := direct.Compile(edges)
			if err != nil {
				return fail(id, "translation failed: %v", err)
			}

			pid := desc.GetId()
			log.Debugf(ctx, "Plan %v: %v", pid, plan)

			c.mu.Lock()
			c.plans[pid] = plan
			c.mu.Unlock()
		}

{'fraction_non_alphanumeric': 0.12809579969155402, 'fraction_numerical': 0.003991653814750975, 'mean_word_length': 3.383300198807157, 'pattern_counts': {'":': 0, '<': 5, '<?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': '5037428', 'n_tokens_mistral': 4080, 'n_tokens_neox': 3627, 'n_words': 1081}