Examples with MetricsContainerImpl org.apache.beam.runners.core.metrics.MetricsContainerImpl used on opensource projects

Example 31 with MetricsContainerImpl

use of org.apache.beam.runners.core.metrics.MetricsContainerImpl in project beam by apache.

the class PTransformFunctionRegistry method register.

/**
 * Register the runnable to process the specific pTransformId and track its execution time.
 *
 * @param pTransformId
 * @param runnable
 */
public void register(String pTransformId, ThrowingRunnable runnable) {
    HashMap<String, String> labelsMetadata = new HashMap<String, String>();
    labelsMetadata.put(MonitoringInfoConstants.Labels.PTRANSFORM, pTransformId);
    String executionTimeUrn = "";
    if (executionStateName.equals(ExecutionStateTracker.START_STATE_NAME)) {
        executionTimeUrn = MonitoringInfoConstants.Urns.START_BUNDLE_MSECS;
    } else if (executionStateName.equals(ExecutionStateTracker.FINISH_STATE_NAME)) {
        executionTimeUrn = MonitoringInfoConstants.Urns.FINISH_BUNDLE_MSECS;
    }
    SimpleExecutionState state = new SimpleExecutionState(this.executionStateName, executionTimeUrn, labelsMetadata);
    executionStates.register(state);
    ThrowingRunnable wrapped = () -> {
        MetricsContainerImpl container = metricsContainerRegistry.getContainer(pTransformId);
        try (Closeable metricCloseable = MetricsEnvironment.scopedMetricsContainer(container)) {
            try (Closeable trackerCloseable = this.stateTracker.enterState(state)) {
                runnable.run();
            }
        }
    };
    runnables.add(wrapped);
}

Example 32 with MetricsContainerImpl

use of org.apache.beam.runners.core.metrics.MetricsContainerImpl in project beam by apache.

the class BigtableServiceImplTest method setup.

@Before
public void setup() {
    MockitoAnnotations.initMocks(this);
    BigtableOptions options = new BigtableOptions.Builder().setProjectId(PROJECT_ID).setInstanceId(INSTANCE_ID).build();
    when(mockSession.getOptions()).thenReturn(options);
    when(mockSession.createBulkMutation(eq(TABLE_NAME))).thenReturn(mockBulkMutation);
    when(mockSession.getDataClient()).thenReturn(mockBigtableDataClient);
    // Setup the ProcessWideContainer for testing metrics are set.
    MetricsContainerImpl container = new MetricsContainerImpl(null);
    MetricsEnvironment.setProcessWideContainer(container);
}

Example 33 with MetricsContainerImpl

use of org.apache.beam.runners.core.metrics.MetricsContainerImpl in project beam by apache.

the class DatastoreV1Test method setUp.

@Before
public void setUp() {
    MockitoAnnotations.initMocks(this);
    when(mockDatastoreFactory.getDatastore(any(PipelineOptions.class), any(String.class), any())).thenReturn(mockDatastore);
    when(mockDatastoreFactory.getQuerySplitter()).thenReturn(mockQuerySplitter);
    // Setup the ProcessWideContainer for testing metrics are set.
    MetricsContainerImpl container = new MetricsContainerImpl(null);
    MetricsEnvironment.setProcessWideContainer(container);
}

Example 34 with MetricsContainerImpl

use of org.apache.beam.runners.core.metrics.MetricsContainerImpl in project beam by apache.

the class BigQueryServicesImplTest method verifyRequestMetricWasSet.

private void verifyRequestMetricWasSet(String method, String projectId, String dataset, String table, String status, long count) {
    // Verify the metric as reported.
    HashMap<String, String> labels = new HashMap<String, String>();
    // TODO(ajamato): Add Ptransform label. Populate it as empty for now to prevent the
    // SpecMonitoringInfoValidator from dropping the MonitoringInfo.
    labels.put(MonitoringInfoConstants.Labels.PTRANSFORM, "");
    labels.put(MonitoringInfoConstants.Labels.SERVICE, "BigQuery");
    labels.put(MonitoringInfoConstants.Labels.METHOD, method);
    labels.put(MonitoringInfoConstants.Labels.RESOURCE, GcpResourceIdentifiers.bigQueryTable(projectId, dataset, table));
    labels.put(MonitoringInfoConstants.Labels.BIGQUERY_PROJECT_ID, projectId);
    labels.put(MonitoringInfoConstants.Labels.BIGQUERY_DATASET, dataset);
    labels.put(MonitoringInfoConstants.Labels.BIGQUERY_TABLE, table);
    labels.put(MonitoringInfoConstants.Labels.STATUS, status);
    MonitoringInfoMetricName name = MonitoringInfoMetricName.named(MonitoringInfoConstants.Urns.API_REQUEST_COUNT, labels);
    MetricsContainerImpl container = (MetricsContainerImpl) MetricsEnvironment.getProcessWideContainer();
    assertEquals(count, (long) container.getCounter(name).getCumulative());
}

Example 35 with MetricsContainerImpl

use of org.apache.beam.runners.core.metrics.MetricsContainerImpl in project beam by apache.

the class SpannerIOReadTest method setUp.

@Before
@SuppressWarnings("unchecked")
public void setUp() throws Exception {
    serviceFactory = new FakeServiceFactory();
    mockBatchTx = Mockito.mock(BatchReadOnlyTransaction.class);
    // Setup the ProcessWideContainer for testing metrics are set.
    MetricsContainerImpl container = new MetricsContainerImpl(null);
    MetricsEnvironment.setProcessWideContainer(container);
}