public/grafana
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Storage: streamline context handling (#84319)

Browse Source 
streamline context handling
This commit is contained in:
Dan Cech
2024-03-14 17:12:20 -04:00
committed by GitHub
parent 2795f9827a
commit 0fa0cede75
1 changed files with 18 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pkg/services/store/entity/sqlstash/sql_storage_server.go
+18 -16
View File
@@ -32,6 +32,7 @@ func ProvideSQLEntityServer(db db.EntityDBInterface /*, cfg *setting.Cfg */) (en
entityServer := &sqlEntityServer{ entityServer := &sqlEntityServer{
db: db, db: db,
log: log.New("sql-entity-server"), log: log.New("sql-entity-server"),
ctx: context.Background(),
} }
return entityServer, nil return entityServer, nil
@@ -44,6 +45,7 @@ type sqlEntityServer struct {
dialect migrator.Dialect dialect migrator.Dialect
snowflake *snowflake.Node snowflake *snowflake.Node
broadcaster Broadcaster[*entity.Entity] broadcaster Broadcaster[*entity.Entity]
ctx context.Context
} }
func (s *sqlEntityServer) Init() error { func (s *sqlEntityServer) Init() error {
@@ -80,7 +82,7 @@ func (s *sqlEntityServer) Init() error {
} }
// set up the broadcaster // set up the broadcaster
s.broadcaster, err = NewBroadcaster(context.Background(), func(stream chan *entity.Entity) error { s.broadcaster, err = NewBroadcaster(s.ctx, func(stream chan *entity.Entity) error {
// start the poller // start the poller
go s.poller(stream) go s.poller(stream)
@@ -132,7 +134,7 @@ func (s *sqlEntityServer) getReadSelect(r *entity.ReadEntityRequest) (string, er
return "SELECT " + strings.Join(quotedFields, ","), nil return "SELECT " + strings.Join(quotedFields, ","), nil
} }
func (s *sqlEntityServer) rowToEntity(ctx context.Context, rows *sql.Rows, r *entity.ReadEntityRequest) (*entity.Entity, error) { func rowToEntity(rows *sql.Rows, r *entity.ReadEntityRequest) (*entity.Entity, error) {
raw := &entity.Entity{ raw := &entity.Entity{
Origin: &entity.EntityOriginInfo{}, Origin: &entity.EntityOriginInfo{},
} }
@@ -247,7 +249,7 @@ func (s *sqlEntityServer) read(ctx context.Context, tx session.SessionQuerier, r
return &entity.Entity{}, nil return &entity.Entity{}, nil
} }
return s.rowToEntity(ctx, rows, r) return rowToEntity(rows, r)
} }
func (s *sqlEntityServer) BatchRead(ctx context.Context, b *entity.BatchReadEntityRequest) (*entity.BatchReadEntityResponse, error) { func (s *sqlEntityServer) BatchRead(ctx context.Context, b *entity.BatchReadEntityRequest) (*entity.BatchReadEntityResponse, error) {
@@ -293,7 +295,7 @@ func (s *sqlEntityServer) BatchRead(ctx context.Context, b *entity.BatchReadEnti
// TODO? make sure the results are in order? // TODO? make sure the results are in order?
rsp := &entity.BatchReadEntityResponse{} rsp := &entity.BatchReadEntityResponse{}
for rows.Next() { for rows.Next() {
r, err := s.rowToEntity(ctx, rows, req) r, err := rowToEntity(rows, req)
if err != nil { if err != nil {
return nil, err return nil, err
} }
@@ -964,7 +966,7 @@ func (s *sqlEntityServer) History(ctx context.Context, r *entity.EntityHistoryRe
Key: r.Key, Key: r.Key,
} }
for rows.Next() { for rows.Next() {
v, err := s.rowToEntity(ctx, rows, rr) v, err := rowToEntity(rows, rr)
if err != nil { if err != nil {
return nil, err return nil, err
} }
@@ -1183,7 +1185,7 @@ func (s *sqlEntityServer) List(ctx context.Context, r *entity.EntityListRequest)
ResourceVersion: s.snowflake.Generate().Int64(), ResourceVersion: s.snowflake.Generate().Int64(),
} }
for rows.Next() { for rows.Next() {
result, err := s.rowToEntity(ctx, rows, rr) result, err := rowToEntity(rows, rr)
if err != nil { if err != nil {
return rsp, err return rsp, err
} }
@@ -1218,13 +1220,13 @@ func (s *sqlEntityServer) Watch(r *entity.EntityWatchRequest, w entity.EntitySto
} }
// collect and send any historical events // collect and send any historical events
err = s.watchInit(w.Context(), r, w) err = s.watchInit(r, w)
if err != nil { if err != nil {
return err return err
} }
// subscribe to new events // subscribe to new events
err = s.watch(w.Context(), r, w) err = s.watch(r, w)
if err != nil { if err != nil {
s.log.Error("watch error", "err", err) s.log.Error("watch error", "err", err)
return err return err
@@ -1234,7 +1236,7 @@ func (s *sqlEntityServer) Watch(r *entity.EntityWatchRequest, w entity.EntitySto
} }
// watchInit is a helper function to send the initial set of entities to the client // watchInit is a helper function to send the initial set of entities to the client
func (s *sqlEntityServer) watchInit(ctx context.Context, r *entity.EntityWatchRequest, w entity.EntityStore_WatchServer) error { func (s *sqlEntityServer) watchInit(r *entity.EntityWatchRequest, w entity.EntityStore_WatchServer) error {
rr := &entity.ReadEntityRequest{ rr := &entity.ReadEntityRequest{
WithBody: r.WithBody, WithBody: r.WithBody,
WithStatus: r.WithStatus, WithStatus: r.WithStatus,
@@ -1325,7 +1327,7 @@ func (s *sqlEntityServer) watchInit(ctx context.Context, r *entity.EntityWatchRe
err = func() error { err = func() error {
query, args := entityQuery.toQuery() query, args := entityQuery.toQuery()
rows, err := s.sess.Query(ctx, query, args...) rows, err := s.sess.Query(w.Context(), query, args...)
if err != nil { if err != nil {
return err return err
} }
@@ -1340,7 +1342,7 @@ func (s *sqlEntityServer) watchInit(ctx context.Context, r *entity.EntityWatchRe
return nil return nil
} }
result, err := s.rowToEntity(ctx, rows, rr) result, err := rowToEntity(rows, rr)
if err != nil { if err != nil {
return err return err
} }
@@ -1382,14 +1384,14 @@ func (s *sqlEntityServer) poller(stream chan *entity.Entity) {
defer t.Stop() defer t.Stop()
for range t.C { for range t.C {
since, err = s.poll(context.Background(), since, stream) since, err = s.poll(since, stream)
if err != nil { if err != nil {
s.log.Error("watch error", "err", err) s.log.Error("watch error", "err", err)
} }
} }
} }
func (s *sqlEntityServer) poll(ctx context.Context, since int64, out chan *entity.Entity) (int64, error) { func (s *sqlEntityServer) poll(since int64, out chan *entity.Entity) (int64, error) {
s.log.Debug("watch poll", "since", since) s.log.Debug("watch poll", "since", since)
rr := &entity.ReadEntityRequest{ rr := &entity.ReadEntityRequest{
@@ -1416,7 +1418,7 @@ func (s *sqlEntityServer) poll(ctx context.Context, since int64, out chan *entit
query, args := entityQuery.toQuery() query, args := entityQuery.toQuery()
rows, err := s.sess.Query(ctx, query, args...) rows, err := s.sess.Query(s.ctx, query, args...)
if err != nil { if err != nil {
return err return err
} }
@@ -1429,7 +1431,7 @@ func (s *sqlEntityServer) poll(ctx context.Context, since int64, out chan *entit
return nil return nil
} }
result, err := s.rowToEntity(ctx, rows, rr) result, err := rowToEntity(rows, rr)
if err != nil { if err != nil {
return err return err
} }
@@ -1510,7 +1512,7 @@ func watchMatches(r *entity.EntityWatchRequest, result *entity.Entity) bool {
} }
// watch is a helper to get the next set of entities and send them to the client // watch is a helper to get the next set of entities and send them to the client
func (s *sqlEntityServer) watch(ctx context.Context, r *entity.EntityWatchRequest, w entity.EntityStore_WatchServer) error { func (s *sqlEntityServer) watch(r *entity.EntityWatchRequest, w entity.EntityStore_WatchServer) error {
s.log.Debug("watch started", "since", r.Since) s.log.Debug("watch started", "since", r.Since)
evts, err := s.broadcaster.Subscribe(w.Context()) evts, err := s.broadcaster.Subscribe(w.Context())