very draft

cloud/tasks/persistence/health.go

@@ -1,6 +1,7 @@
 package persistence
 
 import (
+	"context"
 	"time"
 
 	"github.com/ydb-platform/nbs/cloud/tasks/metrics"
@@ -11,18 +12,19 @@ import (
 type HealthCheck struct {
 	queriesCount              uint64
 	successQueriesCount       uint64
-	storage                   *healthCheckStorage
+	storage                   HealthStorage
 	registry                  metrics.Registry
 	metricsCollectionInterval time.Duration
 }
 
-func (h *HealthCheck) reportSuccessRate() {
+func (h *HealthCheck) reportSuccessRate(ctx context.Context) {
 	if h.queriesCount == 0 {
 		h.registry.Gauge("successRate").Set(0)
 		return
 	}
 
 	h.registry.Gauge("successRate").Set(float64(h.successQueriesCount) / float64(h.queriesCount))
+	h.storage.HeartbeatNode(ctx, time.Now())
 }
 
 ////////////////////////////////////////////////////////////////////////////////
@@ -35,8 +37,9 @@ func (h *HealthCheck) AccountQuery(err error) {
 }
 
 func NewHealthCheck(
+	ctx context.Context,
 	componentName string,
-	storage *healthCheckStorage,
+	storage HealthStorage,
 	registry metrics.Registry,
 ) *HealthCheck {
 
@@ -55,7 +58,7 @@ func NewHealthCheck(
 		defer ticker.Stop()
 
 		for range ticker.C {
-			h.reportSuccessRate()
+			h.reportSuccessRate(ctx)
 		}
 	}()
 

cloud/tasks/persistence/health_storage.go

@@ -2,14 +2,12 @@ package persistence
 
 import (
 	"context"
+	"fmt"
+	"time"
 
 	"github.com/ydb-platform/nbs/cloud/tasks/logging"
 )
 
-// import (
-// 	"time"
-// )
-
 ////////////////////////////////////////////////////////////////////////////////
 
 func healthCheckTableDescription() CreateTableDescription {
@@ -42,10 +40,24 @@ func CreateYDBTables(
 }
 
 type healthCheckStorage struct {
+	db         *YDBClient
+	tablesPath string
 }
 
-func NewStorage(db *YDBClient) *healthCheckStorage {
-	return &healthCheckStorage{}
+type HealthStorage interface {
+	HeartbeatNode(
+		ctx context.Context,
+		// host string,
+		ts time.Time,
+		// value float64,
+	) error
+}
+
+func NewStorage(db *YDBClient) HealthStorage {
+	return &healthCheckStorage{
+		db:         db,
+		tablesPath: db.AbsolutePath("test"),
+	}
 }
 
 // type HealthCheck struct {
@@ -110,3 +122,44 @@ func NewStorage(db *YDBClient) *healthCheckStorage {
 ////////////////////////////////////////////////////////////////////////////////
 
 // Updates heartbeat timestamp and the current number of inflight tasks.
+func (s *healthCheckStorage) heartbeatNode(
+	ctx context.Context,
+	session *Session,
+	// host string,
+	ts time.Time,
+	// value float64,
+) error {
+
+	logging.Debug(
+		ctx,
+		"KEK IS",
+	)
+
+	_, err := session.ExecuteRW(ctx, fmt.Sprintf(`
+		--!syntax_v1
+		pragma TablePathPrefix = "%v";
+		declare $component as Utf8;
+		declare $update_at as Timestamp;
+
+		upsert into nodes (component, update_at)
+		values ($component, $update_at);
+	`, s.tablesPath),
+		ValueParam("$host", UTF8Value("BORIS_BG")),
+		ValueParam("$update_at", TimestampValue(ts)),
+	)
+	return err
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
+func (s *healthCheckStorage) HeartbeatNode(
+	ctx context.Context,
+	// host string,
+	ts time.Time,
+	// value float64,
+) error {
+
+	return s.db.Execute(ctx, func(ctx context.Context, session *Session) error {
+		return s.heartbeatNode(ctx, session, ts)
+	})
+}

cloud/tasks/persistence/health_storage_mocks.go

@@ -0,0 +1,38 @@
+package persistence
+
+import (
+	"context"
+	"time"
+
+	"github.com/stretchr/testify/mock"
+)
+
+////////////////////////////////////////////////////////////////////////////////
+
+type StorageMock struct {
+	mock.Mock
+}
+
+func (s *StorageMock) HeartbeatNode(
+	ctx context.Context,
+	// host string,
+	ts time.Time,
+	// inflightTaskCount uint32,
+) error {
+
+	args := s.Called(ctx, ts)
+	return args.Error(0)
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
+func NewStorageMock() *StorageMock {
+	return &StorageMock{}
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
+// // Ensure that StorageMock implements tasks_storage.Storage.
+// func assertStorageMockIsStorage(arg *StorageMock) tasks_storage.Storage {
+// 	return arg
+// }

cloud/tasks/persistence/health_test.go

@@ -18,7 +18,7 @@ func newStorage(
 	t *testing.T,
 	ctx context.Context,
 	db *YDBClient,
-) *healthCheckStorage {
+) HealthStorage {
 
 	err := CreateYDBTables(
 		ctx,
@@ -43,10 +43,10 @@ func TestHealthCheckMetric(t *testing.T) {
 	require.NoError(t, err)
 	defer db.Close(ctx)
 
-	storage := newStorage(t, ctx, db)
+	storage := NewStorageMock()
 
 	registry := mocks.NewRegistryMock()
-	healthCheck := NewHealthCheck("test", storage, registry)
+	healthCheck := NewHealthCheck(ctx, "test", storage, registry)
 
 	gaugeSetWg := sync.WaitGroup{}
 
@@ -59,6 +59,8 @@ func TestHealthCheckMetric(t *testing.T) {
 			gaugeSetWg.Done()
 		},
 	)
+
+	storage.On("HeartbeatNode", ctx, mock.Anything).Return(nil)
 	gaugeSetWg.Wait()
 
 	healthCheck.AccountQuery(nil)
@@ -78,4 +80,6 @@ func TestHealthCheckMetric(t *testing.T) {
 	gaugeSetWg.Wait()
 
 	registry.AssertAllExpectations(t)
+
+	require.NotNil(t, nil)
 }

cloud/tasks/persistence/s3.go

@@ -49,16 +49,18 @@ type S3Client struct {
 }
 
 func NewS3Client(
+	ctx context.Context,
 	endpoint string,
 	region string,
 	credentials S3Credentials,
 	callTimeout time.Duration,
 	s3MetricsRegistry metrics.Registry,
+	healthCheckStorage HealthStorage,
 	healthMetricsRegistry metrics.Registry,
 	maxRetriableErrorCount uint64,
 ) (*S3Client, error) {
 
-	s3Metrics := newS3Metrics(callTimeout, s3MetricsRegistry, healthMetricsRegistry)
+	s3Metrics := newS3Metrics(ctx, callTimeout, s3MetricsRegistry, healthCheckStorage, healthMetricsRegistry)
 
 	sessionConfig := &aws.Config{
 		Credentials: aws_credentials.NewStaticCredentials(
@@ -90,8 +92,10 @@ func NewS3Client(
 }
 
 func NewS3ClientFromConfig(
+	ctx context.Context,
 	config *persistence_config.S3Config,
 	s3MetricsRegistry metrics.Registry,
+	healthCheckStorage *healthCheckStorage,
 	healthMetricsRegistry metrics.Registry,
 ) (*S3Client, error) {
 
@@ -109,11 +113,13 @@ func NewS3ClientFromConfig(
 	}
 
 	return NewS3Client(
+		ctx,
 		config.GetEndpoint(),
 		config.GetRegion(),
 		credentials,
 		callTimeout,
 		s3MetricsRegistry,
+		healthCheckStorage,
 		healthMetricsRegistry,
 		config.GetMaxRetriableErrorCount(),
 	)

cloud/tasks/persistence/s3_metrics.go

@@ -123,13 +123,15 @@ func (m *s3Metrics) OnRetry(req *request.Request) {
 }
 
 func newS3Metrics(
+	ctx context.Context,
 	callTimeout time.Duration,
 	s3MetricsRegistry metrics.Registry,
+	healthCheckStorage HealthStorage,
 	healthMetricsRegistry metrics.Registry,
 ) *s3Metrics {
 	return &s3Metrics{
 		callTimeout:       callTimeout,
 		s3MetricsRegistry: s3MetricsRegistry,
-		healthCheck:       NewHealthCheck("s3", healthMetricsRegistry),
+		healthCheck:       NewHealthCheck(ctx, "s3", healthCheckStorage, healthMetricsRegistry),
 	}
 }

cloud/tasks/persistence/s3_test.go

@@ -7,6 +7,7 @@ import (
 
 	"github.com/stretchr/testify/require"
 	"github.com/ydb-platform/nbs/cloud/tasks/errors"
+	"github.com/ydb-platform/nbs/cloud/tasks/metrics/empty"
 	"github.com/ydb-platform/nbs/cloud/tasks/metrics/mocks"
 )
 
@@ -17,18 +18,22 @@ const maxRetriableErrorCount = 3
 ////////////////////////////////////////////////////////////////////////////////
 
 func newS3Client(
+	ctx context.Context,
 	callTimeout time.Duration,
 	s3MetricsRegistry *mocks.RegistryMock,
+	healthCheckStorage HealthStorage,
 	healthMetricsRegistry *mocks.RegistryMock,
 ) (*S3Client, error) {
 
 	credentials := NewS3Credentials("test", "test")
 	return NewS3Client(
+		ctx,
 		"endpoint",
 		"region",
 		credentials,
 		callTimeout,
 		s3MetricsRegistry,
+		healthCheckStorage,
 		healthMetricsRegistry,
 		maxRetriableErrorCount,
 	)
@@ -40,9 +45,15 @@ func TestS3ShouldSendErrorCanceledMetric(t *testing.T) {
 	ctx, cancel := context.WithCancel(newContext())
 
 	s3MetricsRegistry := mocks.NewRegistryMock()
+
+	db, err := newYDB(ctx, empty.NewRegistry())
+	require.NoError(t, err)
+	defer db.Close(ctx)
+
+	healthCheckStorage := newStorage(t, ctx, db)
 	healthMetricsRegistry := mocks.NewRegistryMock()
 
-	s3, err := newS3Client(10*time.Second /* callTimeout */, s3MetricsRegistry, healthMetricsRegistry)
+	s3, err := newS3Client(ctx, 10*time.Second /* callTimeout */, s3MetricsRegistry, healthCheckStorage, healthMetricsRegistry)
 	require.NoError(t, err)
 
 	cancel()
@@ -68,9 +79,15 @@ func TestS3ShouldSendErrorTimeoutMetric(t *testing.T) {
 	defer cancel()
 
 	s3MetricsRegistry := mocks.NewRegistryMock()
+
+	db, err := newYDB(ctx, empty.NewRegistry())
+	require.NoError(t, err)
+	defer db.Close(ctx)
+
+	healthCheckStorage := newStorage(t, ctx, db)
 	healthMetricsRegistry := mocks.NewRegistryMock()
 
-	s3, err := newS3Client(0 /* callTimeout */, s3MetricsRegistry, healthMetricsRegistry)
+	s3, err := newS3Client(ctx, 0 /* callTimeout */, s3MetricsRegistry, healthCheckStorage, healthMetricsRegistry)
 	require.NoError(t, err)
 
 	s3MetricsRegistry.GetCounter(
@@ -99,9 +116,15 @@ func TestS3ShouldRetryRequests(t *testing.T) {
 	defer cancel()
 
 	s3MetricsRegistry := mocks.NewRegistryMock()
+
+	db, err := newYDB(ctx, empty.NewRegistry())
+	require.NoError(t, err)
+	defer db.Close(ctx)
+
+	healthCheckStorage := newStorage(t, ctx, db)
 	healthMetricsRegistry := mocks.NewRegistryMock()
 
-	s3, err := newS3Client(10*time.Second /* callTimeout */, s3MetricsRegistry, healthMetricsRegistry)
+	s3, err := newS3Client(ctx, 10*time.Second /* callTimeout */, s3MetricsRegistry, healthCheckStorage, healthMetricsRegistry)
 	require.NoError(t, err)
 
 	s3MetricsRegistry.GetCounter(
@@ -125,9 +148,15 @@ func TestS3ShouldSendHealthMetric(t *testing.T) {
 	defer cancel()
 
 	s3MetricsRegistry := mocks.NewRegistryMock()
+
+	db, err := newYDB(ctx, empty.NewRegistry())
+	require.NoError(t, err)
+	defer db.Close(ctx)
+
+	healthCheckStorage := newStorage(t, ctx, db)
 	healthMetricsRegistry := mocks.NewRegistryMock()
 
-	s3, err := newS3Client(10*time.Second /* callTimeout */, s3MetricsRegistry, healthMetricsRegistry)
+	s3, err := newS3Client(ctx, 10*time.Second /* callTimeout */, s3MetricsRegistry, healthCheckStorage, healthMetricsRegistry)
 	require.NoError(t, err)
 
 	s3MetricsRegistry.GetCounter(

cloud/tasks/persistence/ya.make

@@ -9,6 +9,7 @@ SRCS(
 
     health.go
     health_storage.go
+    health_storage_mocks.go
 )
 
 GO_TEST_SRCS(