feat: support lookup offset by timestamp (#309)

* feat: support lookup offset by time

* update

* update

* Update internal/store/segment/server.go

Co-authored-by: James Yin <[email protected]>

* update due to review

* add todo

Co-authored-by: James Yin <[email protected]>

Author: wenfengwang

client/internal/vanus/eventlog/name_service.go

@@ -18,6 +18,7 @@ import (
 	// standard libraries.
 	"context"
 	"math"
+	"time"
 
 	// third-party libraries.
 	"go.opentelemetry.io/otel/trace"
@@ -26,7 +27,7 @@ import (
 	// first-party libraries.
 	"github.com/linkall-labs/vanus/observability/tracing"
 	"github.com/linkall-labs/vanus/pkg/controller"
-	ctlpb "github.com/linkall-labs/vanus/proto/pkg/controller"
+	ctrlpb "github.com/linkall-labs/vanus/proto/pkg/controller"
 	metapb "github.com/linkall-labs/vanus/proto/pkg/meta"
 
 	// this project.
@@ -42,15 +43,15 @@ func NewNameService(endpoints []string) *NameService {
 }
 
 type NameService struct {
-	client ctlpb.EventLogControllerClient
+	client ctrlpb.EventLogControllerClient
 	tracer *tracing.Tracer
 }
 
 func (ns *NameService) LookupWritableSegment(ctx context.Context, logID uint64) (*record.Segment, error) {
 	ctx, span := ns.tracer.Start(ctx, "LookupWritableSegment")
 	defer span.End()
 
-	req := &ctlpb.GetAppendableSegmentRequest{
+	req := &ctrlpb.GetAppendableSegmentRequest{
 		EventLogId: logID,
 		Limited:    1,
 	}
@@ -71,7 +72,7 @@ func (ns *NameService) LookupReadableSegments(ctx context.Context, logID uint64)
 	ctx, span := ns.tracer.Start(ctx, "LookupReadableSegments")
 	defer span.End()
 
-	req := &ctlpb.ListSegmentRequest{
+	req := &ctrlpb.ListSegmentRequest{
 		EventLogId:  logID,
 		StartOffset: 0,
 		EndOffset:   math.MaxInt64,
@@ -87,13 +88,13 @@ func (ns *NameService) LookupReadableSegments(ctx context.Context, logID uint64)
 	return segments, nil
 }
 
-func toSegments(segmentpbs []*metapb.Segment) []*record.Segment {
-	if len(segmentpbs) == 0 {
+func toSegments(pbs []*metapb.Segment) []*record.Segment {
+	if len(pbs) == 0 {
 		return make([]*record.Segment, 0)
 	}
-	segments := make([]*record.Segment, 0, len(segmentpbs))
-	for _, segmentpb := range segmentpbs {
-		segment := toSegment(segmentpb)
+	segments := make([]*record.Segment, 0, len(pbs))
+	for _, pb := range pbs {
+		segment := toSegment(pb)
 		segments = append(segments, segment)
 		// only return first working segment
 		if segment.Writable {
@@ -103,23 +104,22 @@ func toSegments(segmentpbs []*metapb.Segment) []*record.Segment {
 	return segments
 }
 
-func toSegment(segmentpb *metapb.Segment) *record.Segment {
-	blocks := make(map[uint64]*record.Block, len(segmentpb.Replicas))
-	for blockID, blockpb := range segmentpb.Replicas {
+func toSegment(segment *metapb.Segment) *record.Segment {
+	blocks := make(map[uint64]*record.Block, len(segment.Replicas))
+	for blockID, block := range segment.Replicas {
 		blocks[blockID] = &record.Block{
-			ID:       blockpb.Id,
-			Endpoint: blockpb.Endpoint,
+			ID:       block.Id,
+			Endpoint: block.Endpoint,
 		}
 	}
-	segment := &record.Segment{
-		ID:          segmentpb.GetId(),
-		StartOffset: segmentpb.GetStartOffsetInLog(),
-		// TODO align to server side
-		EndOffset: segmentpb.GetEndOffsetInLog() + 1,
-		// TODO: writable
-		Writable:      segmentpb.State == "working",
-		Blocks:        blocks,
-		LeaderBlockID: segmentpb.GetLeaderBlockId(),
+	return &record.Segment{
+		ID:               segment.GetId(),
+		StartOffset:      segment.GetStartOffsetInLog(),
+		EndOffset:        segment.GetEndOffsetInLog(),
+		FirstEventBornAt: time.UnixMilli(segment.FirstEventBornAtByUnixMs),
+		LastEventBornAt:  time.UnixMilli(segment.LastEvnetBornAtByUnixMs),
+		Writable:         segment.State == "working", // TODO: writable
+		Blocks:           blocks,
+		LeaderBlockID:    segment.GetLeaderBlockId(),
 	}
-	return segment
 }

client/internal/vanus/store/block_store.go

@@ -18,6 +18,7 @@ import (
 	// standard libraries
 	"context"
 	"strings"
+	"time"
 
 	"github.com/linkall-labs/vanus/observability/tracing"
 	"go.opentelemetry.io/otel/trace"
@@ -155,3 +156,24 @@ func (s *BlockStore) Read(
 
 	return []*ce.Event{}, err
 }
+
+func (s *BlockStore) LookupOffset(ctx context.Context, blockID uint64, t time.Time) (int64, error) {
+	ctx, span := s.tracer.Start(ctx, "LookupOffset")
+	defer span.End()
+
+	req := &segpb.LookupOffsetInBlockRequest{
+		BlockId: blockID,
+		Stime:   t.UnixMilli(),
+	}
+
+	client, err := s.client.Get(ctx)
+	if err != nil {
+		return -1, err
+	}
+
+	res, err := client.(segpb.SegmentServerClient).LookupOffsetInBlock(ctx, req)
+	if err != nil {
+		return -1, err
+	}
+	return res.Offset, nil
+}

client/pkg/eventlog/eventlog_impl.go

@@ -174,9 +174,40 @@ func (l *eventlog) Length(ctx context.Context) (int64, error) {
 	// TODO(kai.jiangkai)
 	return 0, nil
 }
+
 func (l *eventlog) QueryOffsetByTime(ctx context.Context, timestamp int64) (int64, error) {
-	// TODO(james.yin): lookup offset by timestamp
-	return 0, nil
+	t := time.UnixMilli(timestamp)
+	// get all segments
+	var target *segment
+	segs := l.fetchReadableSegments(ctx)
+
+	if len(segs) == 0 {
+		return -1, nil
+	}
+
+	if segs[0].firstEventBornAt.After(t) {
+		return segs[0].startOffset, nil
+	}
+
+	if segs[len(segs)-1].lastEventBornAt.Before(t) {
+		// the target offset maybe in newer segment, refresh immediately
+		l.refreshReadableSegments(ctx)
+		segs = l.fetchReadableSegments(ctx)
+	}
+
+	for idx := range segs {
+		s := segs[idx]
+		if !t.Before(s.firstEventBornAt) && !t.After(s.lastEventBornAt) {
+			target = s
+			break
+		}
+	}
+
+	if target == nil {
+		target = segs[len(segs)-1]
+	}
+
+	return target.LookupOffset(ctx, t)
 }
 
 func (l *eventlog) updateWritableSegment(ctx context.Context, r *record.Segment) {
@@ -286,12 +317,10 @@ func (l *eventlog) fetchReadableSegments(ctx context.Context) []*segment {
 	defer l.readableMu.RUnlock()
 
 	if len(l.readableSegments) == 0 {
+		l.readableMu.RUnlock()
 		// refresh
-		func() {
-			l.readableMu.RUnlock()
-			defer l.readableMu.RLock()
-			l.refreshReadableSegments(ctx)
-		}()
+		l.refreshReadableSegments(ctx)
+		l.readableMu.RLock()
 	}
 
 	return l.readableSegments

client/pkg/eventlog/log_segment.go

@@ -20,6 +20,7 @@ import (
 	"encoding/binary"
 	"math"
 	"sync"
+	"time"
 
 	"github.com/linkall-labs/vanus/observability/tracing"
 	"go.opentelemetry.io/otel/trace"
@@ -44,12 +45,14 @@ func newSegment(ctx context.Context, r *record.Segment, towrite bool) (*segment,
 	}
 
 	segment := &segment{
-		id:          r.ID,
-		startOffset: r.StartOffset,
-		endOffset:   atomic.Int64{},
-		writable:    atomic.Bool{},
-		prefer:      prefer,
-		tracer:      tracing.NewTracer("internal.eventlog.segment", trace.SpanKindClient),
+		id:               r.ID,
+		startOffset:      r.StartOffset,
+		endOffset:        atomic.Int64{},
+		writable:         atomic.Bool{},
+		firstEventBornAt: r.FirstEventBornAt,
+		lastEventBornAt:  r.LastEventBornAt,
+		prefer:           prefer,
+		tracer:           tracing.NewTracer("internal.eventlog.segment", trace.SpanKindClient),
 	}
 
 	if !r.Writable {
@@ -82,10 +85,12 @@ func newBlockExt(ctx context.Context, r *record.Segment, leaderOnly bool) (*bloc
 }
 
 type segment struct {
-	id          uint64
-	startOffset int64
-	endOffset   atomic.Int64
-	writable    atomic.Bool
+	id               uint64
+	startOffset      int64
+	endOffset        atomic.Int64
+	writable         atomic.Bool
+	firstEventBornAt time.Time
+	lastEventBornAt  time.Time
 
 	prefer *block
 	mu     sync.RWMutex
@@ -117,7 +122,9 @@ func (s *segment) Close(ctx context.Context) {
 }
 
 func (s *segment) Update(ctx context.Context, r *record.Segment, towrite bool) error {
-	// When a segment become read-only, the end offset needs to be set to the readlly value.
+	// When a segment become read-only, the end offset needs to be set to the real value.
+	// TODO(wenfeng) data race?
+	s.lastEventBornAt = r.LastEventBornAt
 	if s.Writable() && !r.Writable && s.writable.CAS(true, false) {
 		s.endOffset.Store(r.EndOffset)
 		return nil
@@ -219,3 +226,7 @@ func (s *segment) setPreferSegmentBlock(prefer *block) {
 	defer s.mu.Unlock()
 	s.prefer = prefer
 }
+
+func (s *segment) LookupOffset(ctx context.Context, t time.Time) (int64, error) {
+	return s.preferSegmentBlock().LookupOffset(ctx, t)
+}

client/pkg/eventlog/segment_block.go

@@ -17,6 +17,7 @@ package eventlog
 import (
 	// third-party libraries
 	"context"
+	"time"
 
 	ce "github.com/cloudevents/sdk-go/v2"
 
@@ -47,6 +48,10 @@ func (s *block) Close(ctx context.Context) {
 	store.Put(ctx, s.store)
 }
 
+func (s *block) LookupOffset(ctx context.Context, t time.Time) (int64, error) {
+	return s.store.LookupOffset(ctx, s.id, t)
+}
+
 func (s *block) Append(ctx context.Context, event *ce.Event) (int64, error) {
 	return s.store.Append(ctx, s.id, event)
 }

client/pkg/primitive/watcher.go

@@ -79,26 +79,21 @@ func (w *Watcher) Run() {
 func (w *Watcher) Refresh(ctx context.Context) error {
 	// batch multi-refresh into a group
 
-	wg := func() *sync.WaitGroup {
-		w.mu.RLock()
-		defer w.mu.RUnlock()
-		return w.wg
-	}()
+	w.mu.RLock()
+	wg := w.wg
+	w.mu.RUnlock()
 
 	isLeader := false
 	if wg == nil {
-		wg = func() *sync.WaitGroup {
-			w.mu.Lock()
-			defer w.mu.Unlock()
-
-			if w.wg == nil { // double check
-				w.wg = &sync.WaitGroup{}
-				w.wg.Add(1)
-				isLeader = true
-			}
-
-			return w.wg
-		}()
+		w.mu.Lock()
+		// double check
+		if w.wg == nil {
+			w.wg = &sync.WaitGroup{}
+			w.wg.Add(1)
+			isLeader = true
+		}
+		wg = w.wg
+		w.mu.Unlock()
 	}
 
 	if isLeader {
@@ -108,8 +103,8 @@ func (w *Watcher) Refresh(ctx context.Context) error {
 
 	ch := make(chan struct{})
 	go func() {
-		defer close(ch)
 		wg.Wait()
+		close(ch)
 	}()
 
 	select {

client/pkg/record/record.go

@@ -14,6 +14,8 @@
 
 package record
 
+import "time"
+
 type Eventbus struct {
 	Name string
 	Logs []*Eventlog
@@ -81,10 +83,12 @@ func allReadable(ls []*Eventlog) bool {
 }
 
 type Segment struct {
-	ID          uint64
-	StartOffset int64
-	EndOffset   int64
-	Writable    bool
+	ID               uint64
+	StartOffset      int64
+	EndOffset        int64
+	Writable         bool
+	FirstEventBornAt time.Time
+	LastEventBornAt  time.Time
 
 	Blocks        map[uint64]*Block
 	LeaderBlockID uint64

internal/controller/eventbus/eventlog/eventlog.go

@@ -319,7 +319,7 @@ func (mgr *eventlogManager) UpdateSegment(ctx context.Context, m map[string][]Se
 				})
 				continue
 			}
-			// TODO(wenfeng.wang) Don't update state in isNeedUpdate
+			// TODO(wenfeng.wang) Don't update state in isNeedUpdate, rename?
 			if seg.isNeedUpdate(newSeg) {
 				err := el.updateSegment(ctx, seg)
 				if err != nil {

internal/controller/eventbus/eventlog/segment.go

@@ -149,17 +149,19 @@ func Convert2ProtoSegment(ctx context.Context, ins ...*Segment) []*metapb.Segmen
 			}
 		}
 		segs[idx] = &metapb.Segment{
-			Id:                seg.ID.Uint64(),
-			PreviousSegmentId: seg.PreviousSegmentID.Uint64(),
-			NextSegmentId:     seg.NextSegmentID.Uint64(),
-			EventLogId:        seg.EventLogID.Uint64(),
-			StartOffsetInLog:  seg.StartOffsetInLog,
-			EndOffsetInLog:    seg.StartOffsetInLog + int64(seg.Number) - 1,
-			Size:              seg.Size,
-			Capacity:          seg.Capacity,
-			NumberEventStored: seg.Number,
-			Replicas:          blocks,
-			State:             string(seg.State),
+			Id:                       seg.ID.Uint64(),
+			PreviousSegmentId:        seg.PreviousSegmentID.Uint64(),
+			NextSegmentId:            seg.NextSegmentID.Uint64(),
+			EventLogId:               seg.EventLogID.Uint64(),
+			StartOffsetInLog:         seg.StartOffsetInLog,
+			EndOffsetInLog:           seg.StartOffsetInLog + int64(seg.Number),
+			Size:                     seg.Size,
+			Capacity:                 seg.Capacity,
+			NumberEventStored:        seg.Number,
+			Replicas:                 blocks,
+			State:                    string(seg.State),
+			FirstEventBornAtByUnixMs: seg.FirstEventBornTime.UnixMilli(),
+			LastEvnetBornAtByUnixMs:  seg.LastEventBornTime.UnixMilli(),
 		}
 		if seg.GetLeaderBlock() != nil {
 			segs[idx].LeaderBlockId = seg.GetLeaderBlock().ID.Uint64()