Skip to content

Commit

Permalink
feat(bigquery/storage/managedwriter): add opencensus instrumentation (#…
Browse files Browse the repository at this point in the history
…4512)

This starts to plumb in oc instrumentation, as we're already using it in
bigquery proper and other veneers.

Testing instrumentation helped catch another double-close in the recv processor, so this
addresses that as well.

Towards #4366
  • Loading branch information
shollyman committed Jul 30, 2021
1 parent 95a27f7 commit 73b6f5e
Show file tree
Hide file tree
Showing 5 changed files with 225 additions and 43 deletions.
1 change: 1 addition & 0 deletions bigquery/go.mod
Expand Up @@ -8,6 +8,7 @@ require (
github.com/golang/protobuf v1.5.2
github.com/google/go-cmp v0.5.6
github.com/googleapis/gax-go/v2 v2.0.5
go.opencensus.io v0.23.0
golang.org/x/sync v0.0.0-20210220032951-036812b2e83c
golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1
google.golang.org/api v0.51.0
Expand Down
4 changes: 3 additions & 1 deletion bigquery/storage/managedwriter/client.go
Expand Up @@ -74,7 +74,6 @@ func (c *Client) NewManagedStream(ctx context.Context, opts ...WriterOption) (*M
}

func (c *Client) buildManagedStream(ctx context.Context, streamFunc streamClientFunc, skipSetup bool, opts ...WriterOption) (*ManagedStream, error) {

ctx, cancel := context.WithCancel(ctx)

ms := &ManagedStream{
Expand Down Expand Up @@ -122,6 +121,9 @@ func (c *Client) buildManagedStream(ctx context.Context, streamFunc streamClient
}
}
if ms.streamSettings != nil {
if ms.ctx != nil {
ms.ctx = keyContextWithStreamID(ms.ctx, ms.streamSettings.streamID)
}
ms.fc = newFlowController(ms.streamSettings.MaxInflightRequests, ms.streamSettings.MaxInflightBytes)
} else {
ms.fc = newFlowController(0, 0)
Expand Down
111 changes: 111 additions & 0 deletions bigquery/storage/managedwriter/instrumentation.go
@@ -0,0 +1,111 @@
// Copyright 2021 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// https://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 managedwriter

import (
"context"
"log"
"sync"

"go.opencensus.io/stats"
"go.opencensus.io/stats/view"
"go.opencensus.io/tag"
)

var (
// Metrics on a stream are tagged with the stream ID.
keyStream = tag.MustNewKey("streamID")
)

const statsPrefix = "cloud.google.com/go/bigquery/storage/managedwriter/"

var (
// AppendRequests is a measure of the number of append requests sent.
// It is EXPERIMENTAL and subject to change or removal without notice.
AppendRequests = stats.Int64(statsPrefix+"append_requests", "Number of append requests sent", stats.UnitDimensionless)

// AppendResponses is a measure of the number of append responses received.
// It is EXPERIMENTAL and subject to change or removal without notice.
AppendResponses = stats.Int64(statsPrefix+"append_responses", "Number of append responses sent", stats.UnitDimensionless)

// FlushRequests is a measure of the number of FlushRows requests sent.
// It is EXPERIMENTAL and subject to change or removal without notice.
FlushRequests = stats.Int64(statsPrefix+"flush_requests", "Number of FlushRows requests sent", stats.UnitDimensionless)

// AppendClientOpenCount is a measure of the number of times the AppendRowsClient was opened.
// It is EXPERIMENTAL and subject to change or removal without notice.
AppendClientOpenCount = stats.Int64(statsPrefix+"stream_open_count", "Number of times AppendRowsClient was opened", stats.UnitDimensionless)

// AppendClientOpenRetryCount is a measure of the number of times the AppendRowsClient open was retried.
// It is EXPERIMENTAL and subject to change or removal without notice.
AppendClientOpenRetryCount = stats.Int64(statsPrefix+"stream_open_retry_count", "Number of times AppendRowsClient open was retried", stats.UnitDimensionless)
)

var (
// AppendRequestsView is a cumulative sum of AppendRequests.
// It is EXPERIMENTAL and subject to change or removal without notice.
AppendRequestsView *view.View

// AppendResponsesView is a cumulative sum of AppendResponses.
// It is EXPERIMENTAL and subject to change or removal without notice.
AppendResponsesView *view.View

// FlushRequestsView is a cumulative sum of FlushRequests.
// It is EXPERIMENTAL and subject to change or removal without notice.
FlushRequestsView *view.View

// AppendClientOpenView is a cumulative sum of AppendClientOpenCount.
// It is EXPERIMENTAL and subject to change or removal without notice.
AppendClientOpenView *view.View

// AppendClientOpenRetryView is a cumulative sum of AppendClientOpenRetryCount.
// It is EXPERIMENTAL and subject to change or removal without notice.
AppendClientOpenRetryView *view.View
)

func init() {
AppendRequestsView = createCountView(stats.Measure(AppendRequests), keyStream)
AppendResponsesView = createCountView(stats.Measure(AppendResponses), keyStream)
FlushRequestsView = createCountView(stats.Measure(FlushRequests), keyStream)
AppendClientOpenView = createCountView(stats.Measure(AppendClientOpenCount), keyStream)
AppendClientOpenRetryView = createCountView(stats.Measure(AppendClientOpenRetryCount), keyStream)
}

func createCountView(m stats.Measure, keys ...tag.Key) *view.View {
return &view.View{
Name: m.Name(),
Description: m.Description(),
TagKeys: keys,
Measure: m,
Aggregation: view.Sum(),
}
}

var logOnce sync.Once

// keyContextWithStreamID returns a new context modified with the streamID tag.
func keyContextWithStreamID(ctx context.Context, streamID string) context.Context {
ctx, err := tag.New(ctx, tag.Upsert(keyStream, streamID))
if err != nil {
logOnce.Do(func() {
log.Printf("managedwriter: error creating tag map for 'stream' key: %v", err)
})
}
return ctx
}

func recordStat(ctx context.Context, m *stats.Int64Measure, n int64) {
stats.Record(ctx, m.M(n))
}
146 changes: 104 additions & 42 deletions bigquery/storage/managedwriter/integration_test.go
Expand Up @@ -17,6 +17,7 @@ package managedwriter
import (
"context"
"fmt"
"math"
"testing"
"time"

Expand All @@ -25,6 +26,7 @@ import (
"cloud.google.com/go/bigquery/storage/managedwriter/testdata"
"cloud.google.com/go/internal/testutil"
"cloud.google.com/go/internal/uid"
"go.opencensus.io/stats/view"
"google.golang.org/api/option"
"google.golang.org/protobuf/encoding/protojson"
"google.golang.org/protobuf/proto"
Expand All @@ -45,6 +47,14 @@ var testSimpleSchema = bigquery.Schema{
{Name: "value", Type: bigquery.IntegerFieldType, Required: true},
}

var testSimpleData = []*testdata.SimpleMessage{
{Name: "one", Value: 1},
{Name: "two", Value: 2},
{Name: "three", Value: 3},
{Name: "four", Value: 1},
{Name: "five", Value: 2},
}

func getTestClients(ctx context.Context, t *testing.T, opts ...option.ClientOption) (*Client, *bigquery.Client) {
if testing.Short() {
t.Skip("Integration tests skipped in short mode")
Expand Down Expand Up @@ -162,12 +172,15 @@ func TestIntegration_ManagedWriter(t *testing.T) {
t.Parallel()
testPendingStream(ctx, t, mwClient, bqClient, dataset)
})
t.Run("Instrumentation", func(t *testing.T) {
// Don't run this in parallel, we only want to collect stats from this subtest.
testInstrumentation(ctx, t, mwClient, bqClient, dataset)
})
})

}

func testDefaultStream(ctx context.Context, t *testing.T, mwClient *Client, bqClient *bigquery.Client, dataset *bigquery.Dataset) {

// prep a suitable destination table.
testTable := dataset.Table(tableIDs.New())
if err := testTable.Create(ctx, &bigquery.TableMetadata{Schema: testSimpleSchema}); err != nil {
t.Fatalf("failed to create test table %q: %v", testTable.FullyQualifiedName(), err)
Expand All @@ -188,17 +201,9 @@ func testDefaultStream(ctx context.Context, t *testing.T, mwClient *Client, bqCl
}
validateRowCount(ctx, t, bqClient, testTable, 0, "before send")

testData := []*testdata.SimpleMessage{
{Name: "one", Value: 1},
{Name: "two", Value: 2},
{Name: "three", Value: 3},
{Name: "four", Value: 1},
{Name: "five", Value: 2},
}

// First, send the test rows individually.
var results []*AppendResult
for k, mesg := range testData {
for k, mesg := range testSimpleData {
b, err := proto.Marshal(mesg)
if err != nil {
t.Errorf("failed to marshal message %d: %v", k, err)
Expand All @@ -211,12 +216,12 @@ func testDefaultStream(ctx context.Context, t *testing.T, mwClient *Client, bqCl
}
// wait for the result to indicate ready, then validate.
results[0].Ready()
wantRows := int64(len(testData))
wantRows := int64(len(testSimpleData))
validateRowCount(ctx, t, bqClient, testTable, wantRows, "after first send")

// Now, send the test rows grouped into in a single append.
var data [][]byte
for k, mesg := range testData {
for k, mesg := range testSimpleData {
b, err := proto.Marshal(mesg)
if err != nil {
t.Errorf("failed to marshal message %d: %v", k, err)
Expand Down Expand Up @@ -313,16 +318,8 @@ func testBufferedStream(ctx context.Context, t *testing.T, mwClient *Client, bqC

validateRowCount(ctx, t, bqClient, testTable, 0, "before send")

testData := []*testdata.SimpleMessage{
{Name: "one", Value: 1},
{Name: "two", Value: 2},
{Name: "three", Value: 3},
{Name: "four", Value: 1},
{Name: "five", Value: 2},
}

var expectedRows int64
for k, mesg := range testData {
for k, mesg := range testSimpleData {
b, err := proto.Marshal(mesg)
if err != nil {
t.Errorf("failed to marshal message %d: %v", k, err)
Expand Down Expand Up @@ -368,16 +365,8 @@ func testCommittedStream(ctx context.Context, t *testing.T, mwClient *Client, bq
}
validateRowCount(ctx, t, bqClient, testTable, 0, "before send")

testData := []*testdata.SimpleMessage{
{Name: "one", Value: 1},
{Name: "two", Value: 2},
{Name: "three", Value: 3},
{Name: "four", Value: 1},
{Name: "five", Value: 2},
}

var results []*AppendResult
for k, mesg := range testData {
for k, mesg := range testSimpleData {
b, err := proto.Marshal(mesg)
if err != nil {
t.Errorf("failed to marshal message %d: %v", k, err)
Expand All @@ -390,7 +379,7 @@ func testCommittedStream(ctx context.Context, t *testing.T, mwClient *Client, bq
}
// wait for the result to indicate ready, then validate.
results[0].Ready()
wantRows := int64(len(testData))
wantRows := int64(len(testSimpleData))
validateRowCount(ctx, t, bqClient, testTable, wantRows, "after send")
}

Expand All @@ -413,17 +402,9 @@ func testPendingStream(ctx context.Context, t *testing.T, mwClient *Client, bqCl
}
validateRowCount(ctx, t, bqClient, testTable, 0, "before send")

testData := []*testdata.SimpleMessage{
{Name: "one", Value: 1},
{Name: "two", Value: 2},
{Name: "three", Value: 3},
{Name: "four", Value: 1},
{Name: "five", Value: 2},
}

// Send data.
var results []*AppendResult
for k, mesg := range testData {
for k, mesg := range testSimpleData {
b, err := proto.Marshal(mesg)
if err != nil {
t.Errorf("failed to marshal message %d: %v", k, err)
Expand All @@ -435,7 +416,7 @@ func testPendingStream(ctx context.Context, t *testing.T, mwClient *Client, bqCl
}
}
results[0].Ready()
wantRows := int64(len(testData))
wantRows := int64(len(testSimpleData))

// Mark stream complete.
trackedOffset, err := ms.Finalize(ctx)
Expand All @@ -457,3 +438,84 @@ func testPendingStream(ctx context.Context, t *testing.T, mwClient *Client, bqCl
}
validateRowCount(ctx, t, bqClient, testTable, wantRows, "after send")
}

func testInstrumentation(ctx context.Context, t *testing.T, mwClient *Client, bqClient *bigquery.Client, dataset *bigquery.Dataset) {
testedViews := []*view.View{
AppendRequestsView,
AppendResponsesView,
AppendClientOpenView,
}

if err := view.Register(testedViews...); err != nil {
t.Fatalf("couldn't register views: %v", err)
}

testTable := dataset.Table(tableIDs.New())
if err := testTable.Create(ctx, &bigquery.TableMetadata{Schema: testSimpleSchema}); err != nil {
t.Fatalf("failed to create test table %q: %v", testTable.FullyQualifiedName(), err)
}

m := &testdata.SimpleMessage{}
descriptorProto := protodesc.ToDescriptorProto(m.ProtoReflect().Descriptor())

// setup a new stream.
ms, err := mwClient.NewManagedStream(ctx,
WithDestinationTable(fmt.Sprintf("projects/%s/datasets/%s/tables/%s", testTable.ProjectID, testTable.DatasetID, testTable.TableID)),
WithType(DefaultStream),
WithSchemaDescriptor(descriptorProto),
)
if err != nil {
t.Fatalf("NewManagedStream: %v", err)
}

var results []*AppendResult
for k, mesg := range testSimpleData {
b, err := proto.Marshal(mesg)
if err != nil {
t.Errorf("failed to marshal message %d: %v", k, err)
}
data := [][]byte{b}
results, err = ms.AppendRows(ctx, data, NoStreamOffset)
if err != nil {
t.Errorf("single-row append %d failed: %v", k, err)
}
}
// wait for the result to indicate ready.
results[0].Ready()
// Ick. Stats reporting can't force flushing, and there's a race here. Sleep to give the recv goroutine a chance
// to report.
time.Sleep(time.Second)

for _, tv := range testedViews {
metricData, err := view.RetrieveData(tv.Name)
if err != nil {
t.Errorf("view %q RetrieveData: %v", tv.Name, err)
}
if len(metricData) > 1 {
t.Errorf("%q: only expected 1 row, got %d", tv.Name, len(metricData))
}
if len(metricData[0].Tags) != 1 {
t.Errorf("%q: only expected 1 tag, got %d", tv.Name, len(metricData[0].Tags))
}
entry := metricData[0].Data
sum, ok := entry.(*view.SumData)
if !ok {
t.Errorf("unexpected metric type: %T", entry)
}
got := sum.Value
var want int64
switch tv {
case AppendRequestsView:
want = int64(len(testSimpleData))
case AppendResponsesView:
want = int64(len(testSimpleData))
case AppendClientOpenView:
want = 1
}

// float comparison; diff more than error bound is error
if math.Abs(got-float64(want)) > 0.1 {
t.Errorf("%q: metric mismatch, got %f want %d", tv.Name, got, want)
}
}
}

0 comments on commit 73b6f5e

Please sign in to comment.