Method runLocalSubscriber

worker/import.go:217–297

(ctx context.Context, stream pb.Worker_StreamExtSnapshotServer)

}

func (ps *pubSub) runLocalSubscriber(ctx context.Context, stream pb.Worker_StreamExtSnapshotServer) error {
	defer func() {
		glog.Infof("[import] local subscriber stopped")
	}()

	buffer := ps.subscribe()
	defer ps.unsubscribe(buffer) // ensure publisher won't block on us if we exit

	// Defense in depth: never run the destructive Prepare()/dropAll() unless import mode has been
	// armed by an authorized UpdateExtSnapshotStreamingState(Start). On a follower receiving a
	// forwarded stream this may briefly lag the leader's Raft apply, so wait with a bounded
	// deadline rather than failing instantly.
	if err := waitForExtSnapshotArmed(ctx); err != nil {
		glog.Errorf("[import:flush] refusing to write external snapshot: %v", err)
		return err
	}

	glog.Infof("[import:flush] flushing external snapshot in badger db")

	sw := pstore.NewStreamWriter()
	defer sw.Cancel()

	// Prepare() calls Badger's dropAll(), which is destructive. Defer it until the first valid
	// data packet arrives so that a "Done-only" or otherwise empty stream completes as a no-op
	// and leaves the existing store intact.
	prepared := false

Loop:
	for {
		select {
		case <-ctx.Done():
			glog.Infof("[import] Context cancelled, stopping receive goroutine: %v", ctx.Err())
			return ctx.Err()

		default:
			msg, ok := <-buffer
			if !ok {
				break Loop
			}
			kvs := msg.GetPkt()
			if kvs == nil {
				continue
			}
			if kvs.Done {
				break
			}
			if len(kvs.Data) == 0 {
				continue
			}

			if !prepared {
				if err := sw.Prepare(); err != nil {
					return err
				}
				prepared = true
			}

			buf := z.NewBufferSlice(kvs.Data)