fix: verify branch metadata after snapshot creation (#646)

Author: agneum

engine/internal/provision/thinclones/zfs/branching.go

@@ -126,32 +126,33 @@ func (m *Manager) VerifyBranchMetadata() error {
 		return nil
 	}
 
-	latest := snapshots[0]
+	branchHeads := make(map[string]string)
 
-	brName, err := m.getProperty(branchProp, latest.ID)
-	if err != nil {
-		log.Dbg("cannot find branch for snapshot", latest.ID, err.Error())
-	}
+	for i := numberSnapshots; i > 0; i-- {
+		sn := snapshots[i-1]
+		log.Dbg(sn)
 
-	for i := numberSnapshots; i > 1; i-- {
-		if err := m.SetRelation(snapshots[i-1].ID, snapshots[i-2].ID); err != nil {
-			return fmt.Errorf("failed to set snapshot relations: %w", err)
+		if err := m.DeleteBranchProp(sn.Branch, sn.ID); err != nil {
+			return fmt.Errorf("failed to clean branch property: %w", err)
 		}
 
-		if brName == "" {
-			brName, err = m.getProperty(branchProp, snapshots[i-1].ID)
-			if err != nil {
-				log.Dbg("cannot find branch for snapshot", snapshots[i-1].ID, err.Error())
-			}
+		head, ok := branchHeads[sn.Branch]
+		if !ok {
+			branchHeads[sn.Branch] = sn.ID
+			continue
+		}
+
+		if err := m.SetRelation(head, sn.ID); err != nil {
+			return fmt.Errorf("failed to set snapshot relations: %w", err)
 		}
-	}
 
-	if brName == "" {
-		brName = branching.DefaultBranch
+		branchHeads[sn.Branch] = sn.ID
 	}
 
-	if err := m.AddBranchProp(brName, latest.ID); err != nil {
-		return fmt.Errorf("failed to add branch property: %w", err)
+	for brName, latestID := range branchHeads {
+		if err := m.AddBranchProp(brName, latestID); err != nil {
+			return fmt.Errorf("failed to add branch property: %w", err)
+		}
 	}
 
 	log.Msg("data branching has been verified")

engine/internal/provision/thinclones/zfs/zfs.go

@@ -373,11 +373,17 @@ func (m *Manager) CreateSnapshot(poolSuffix, dataStateAt string) (string, error)
 		return "", fmt.Errorf("failed to parse dataStateAt: %w", err)
 	}
 
+	branch := branching.ParseBranchNameFromSnapshot(snapshotName, poolName)
+	if branch == "" {
+		branch = branching.DefaultBranch
+	}
+
 	newSnapshot := resources.Snapshot{
 		ID:          snapshotName,
 		CreatedAt:   time.Now(),
 		DataStateAt: dataStateTime,
 		Pool:        m.config.Pool.Name,
+		Branch:      branch,
 	}
 
 	if !strings.HasSuffix(snapshotName, m.config.PreSnapshotSuffix) {

engine/internal/retrieval/engine/postgres/snapshot/physical.go

@@ -402,6 +402,10 @@ func (p *PhysicalInitial) run(ctx context.Context) (err error) {
 
 	p.tm.SendEvent(ctx, telemetry.SnapshotCreatedEvent, telemetry.SnapshotCreated{})
 
+	if err := p.cloneManager.VerifyBranchMetadata(); err != nil {
+		log.Warn("cannot verify branch metadata", err.Error())
+	}
+
 	if err := p.cleanupOldLogs(); err != nil {
 		log.Warn("cannot clean up old logs", err.Error())
 	}

engine/internal/srv/routes.go

@@ -55,7 +55,7 @@ func (s *Server) retrievalState(w http.ResponseWriter, r *http.Request) {
 		retrieving.NextRefresh = models.NewLocalTime(spec.Next(time.Now()))
 	}
 
-	retrieving.Activity = s.jobActivity(r.Context())
+	retrieving.Activity = s.jobActivity(context.Background())
 
 	if err := api.WriteJSON(w, http.StatusOK, retrieving); err != nil {
 		api.SendError(w, r, err)
@@ -182,7 +182,9 @@ func (s *Server) createSnapshot(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	// TODO: set branching metadata.
+	if err := fsManager.VerifyBranchMetadata(); err != nil {
+		log.Warn("cannot verify branch metadata", err.Error())
+	}
 
 	latestSnapshot := snapshotList[0]