public void completeSegmentOperations(String tableNameWithType, SegmentMetadata segmentMetadata,
      FileUploadType uploadType, @Nullable URI finalSegmentLocationURI, File segmentFile,
      @Nullable String sourceDownloadURIStr, String segmentDownloadURIStr, @Nullable String crypterName,
      long segmentSizeInBytes, boolean enableParallelPushProtection, boolean allowRefresh, HttpHeaders headers)
      throws Exception {
    String segmentName = segmentMetadata.getName();
    boolean refreshOnly =
        Boolean.parseBoolean(headers.getHeaderString(FileUploadDownloadClient.CustomHeaders.REFRESH_ONLY));

    ZNRecord existingSegmentMetadataZNRecord =
        _pinotHelixResourceManager.getSegmentMetadataZnRecord(tableNameWithType, segmentName);
    if (existingSegmentMetadataZNRecord != null && shouldProcessAsNewSegment(tableNameWithType, segmentName,
        existingSegmentMetadataZNRecord, enableParallelPushProtection)) {
      LOGGER.warn("Removing segment ZK metadata (recovering from previous upload failure) for table: {}, segment: {}",
          tableNameWithType, segmentName);
      Preconditions.checkState(_pinotHelixResourceManager.removeSegmentZKMetadata(tableNameWithType, segmentName),
          "Failed to remove segment ZK metadata for table: %s, segment: %s", tableNameWithType, segmentName);
      existingSegmentMetadataZNRecord = null;
    }

    if (existingSegmentMetadataZNRecord == null) {
      // Add a new segment
      if (refreshOnly) {
        throw new ControllerApplicationException(LOGGER,
            String.format("Cannot refresh non-existing segment: %s for table: %s", segmentName, tableNameWithType),
            Response.Status.GONE);
      }
      LOGGER.info("Adding new segment: {} to table: {}", segmentName, tableNameWithType);
      processNewSegment(tableNameWithType, segmentMetadata, uploadType, finalSegmentLocationURI, segmentFile,
          sourceDownloadURIStr, segmentDownloadURIStr, crypterName, segmentSizeInBytes, enableParallelPushProtection,
          headers);
    } else {
      // Refresh an existing segment
      if (!allowRefresh) {
        // We cannot perform this check up-front in UploadSegment API call. If a segment doesn't exist during the check
        // done up-front but ends up getting created before the check here, we could incorrectly refresh an existing
        // segment.
        throw new ControllerApplicationException(LOGGER,
            String.format("Segment: %s already exists in table: %s. Refresh not permitted.", segmentName,
                tableNameWithType), Response.Status.CONFLICT);
      }
      LOGGER.info("Segment: {} already exists in table: {}, refreshing it", segmentName, tableNameWithType);
      processExistingSegment(tableNameWithType, segmentMetadata, uploadType, existingSegmentMetadataZNRecord,
          finalSegmentLocationURI, segmentFile, sourceDownloadURIStr, segmentDownloadURIStr, crypterName,
          segmentSizeInBytes, enableParallelPushProtection, headers);
    }
  }
@Test
  public void testCompleteSegmentOperations()
      throws Exception {
    ZKOperator zkOperator = new ZKOperator(_resourceManager, mock(ControllerConf.class), mock(ControllerMetrics.class));

    SegmentMetadata segmentMetadata = mock(SegmentMetadata.class);
    when(segmentMetadata.getName()).thenReturn(SEGMENT_NAME);
    when(segmentMetadata.getCrc()).thenReturn("12345");
    when(segmentMetadata.getIndexCreationTime()).thenReturn(123L);
    HttpHeaders httpHeaders = mock(HttpHeaders.class);

    // Test if Zk segment metadata is removed if exception is thrown when moving segment to final location.
    try {
      // Create mock finalSegmentLocationURI and segmentFile.
      URI finalSegmentLocationURI =
          URIUtils.getUri("mockPath", OFFLINE_TABLE_NAME, URIUtils.encode(segmentMetadata.getName()));
      File segmentFile = new File(new File("foo/bar"), "mockChild");
      zkOperator.completeSegmentOperations(OFFLINE_TABLE_NAME, segmentMetadata, FileUploadType.SEGMENT,
          finalSegmentLocationURI, segmentFile, "downloadUrl", "downloadUrl", "crypter", 10, true, true, httpHeaders);
      fail();
    } catch (Exception e) {
      // Expected
    }

    // Wait for the segment Zk entry to be deleted.
    TestUtils.waitForCondition(aVoid -> {
      SegmentZKMetadata segmentZKMetadata = _resourceManager.getSegmentZKMetadata(OFFLINE_TABLE_NAME, SEGMENT_NAME);
      return segmentZKMetadata == null;
    }, 30_000L, "Failed to delete segmentZkMetadata.");

    zkOperator.completeSegmentOperations(OFFLINE_TABLE_NAME, segmentMetadata, FileUploadType.SEGMENT, null, null,
        "downloadUrl", "downloadUrl", "crypter", 10, true, true, httpHeaders);
    SegmentZKMetadata segmentZKMetadata = _resourceManager.getSegmentZKMetadata(OFFLINE_TABLE_NAME, SEGMENT_NAME);
    assertNotNull(segmentZKMetadata);
    assertEquals(segmentZKMetadata.getCrc(), 12345L);
    assertEquals(segmentZKMetadata.getCreationTime(), 123L);
    long pushTime = segmentZKMetadata.getPushTime();
    assertTrue(pushTime > 0);
    assertEquals(segmentZKMetadata.getRefreshTime(), Long.MIN_VALUE);
    assertEquals(segmentZKMetadata.getDownloadUrl(), "downloadUrl");
    assertEquals(segmentZKMetadata.getCrypterName(), "crypter");
    assertEquals(segmentZKMetadata.getSegmentUploadStartTime(), -1);
    assertEquals(segmentZKMetadata.getSizeInBytes(), 10);

    // Test if the same segment can be uploaded when the previous upload failed after segment ZK metadata is created but
    // before segment is assigned to the ideal state
    // Manually remove the segment from the ideal state
    IdealState idealState = _resourceManager.getTableIdealState(OFFLINE_TABLE_NAME);
    assertNotNull(idealState);
    idealState.getRecord().getMapFields().remove(SEGMENT_NAME);
    _resourceManager.getHelixAdmin()
        .setResourceIdealState(_resourceManager.getHelixClusterName(), OFFLINE_TABLE_NAME, idealState);
    // The segment should be uploaded as a new segment (push time should change, and refresh time shouldn't be set)
    zkOperator.completeSegmentOperations(OFFLINE_TABLE_NAME, segmentMetadata, FileUploadType.SEGMENT, null, null,
        "downloadUrl", "downloadUrl", "crypter", 10, true, true, httpHeaders);
    segmentZKMetadata = _resourceManager.getSegmentZKMetadata(OFFLINE_TABLE_NAME, SEGMENT_NAME);
    assertNotNull(segmentZKMetadata);
    assertEquals(segmentZKMetadata.getCrc(), 12345L);
    assertEquals(segmentZKMetadata.getCreationTime(), 123L);
    assertTrue(segmentZKMetadata.getPushTime() > pushTime);
    pushTime = segmentZKMetadata.getPushTime();
    assertTrue(pushTime > 0);
    assertEquals(segmentZKMetadata.getRefreshTime(), Long.MIN_VALUE);
    assertEquals(segmentZKMetadata.getDownloadUrl(), "downloadUrl");
    assertEquals(segmentZKMetadata.getCrypterName(), "crypter");
    assertEquals(segmentZKMetadata.getSegmentUploadStartTime(), -1);
    assertEquals(segmentZKMetadata.getSizeInBytes(), 10);

    // Upload the same segment with allowRefresh = false. Validate that an exception is thrown.
    try {
      zkOperator.completeSegmentOperations(OFFLINE_TABLE_NAME, segmentMetadata, FileUploadType.SEGMENT, null, null,
          "otherDownloadUrl", "otherDownloadUrl", "otherCrypter", 10, true, false, httpHeaders);
      fail();
    } catch (Exception e) {
      // Expected
    }

    // Refresh the segment with unmatched IF_MATCH field
    when(httpHeaders.getHeaderString(HttpHeaders.IF_MATCH)).thenReturn("123");
    try {
      zkOperator.completeSegmentOperations(OFFLINE_TABLE_NAME, segmentMetadata, FileUploadType.SEGMENT, null, null,
          "otherDownloadUrl", "otherDownloadUrl", "otherCrypter", 10, true, true, httpHeaders);
      fail();
    } catch (Exception e) {
      // Expected
    }

    // Refresh the segment with the same segment (same CRC) with matched IF_MATCH field but different creation time,
    // downloadURL and crypter
    when(httpHeaders.getHeaderString(HttpHeaders.IF_MATCH)).thenReturn("12345");
    when(segmentMetadata.getIndexCreationTime()).thenReturn(456L);
    zkOperator.completeSegmentOperations(OFFLINE_TABLE_NAME, segmentMetadata, FileUploadType.SEGMENT, null, null,
        "otherDownloadUrl", "otherDownloadUrl", "otherCrypter", 10, true, true, httpHeaders);

    segmentZKMetadata = _resourceManager.getSegmentZKMetadata(OFFLINE_TABLE_NAME, SEGMENT_NAME);
    assertNotNull(segmentZKMetadata);
    assertEquals(segmentZKMetadata.getCrc(), 12345L);
    // Push time should not change
    assertEquals(segmentZKMetadata.getPushTime(), pushTime);
    // Creation time and refresh time should change
    assertEquals(segmentZKMetadata.getCreationTime(), 456L);
    long refreshTime = segmentZKMetadata.getRefreshTime();
    assertTrue(refreshTime > 0);
    // Download URL should change. Refer: https://github.com/apache/pinot/issues/11535
    assertEquals(segmentZKMetadata.getDownloadUrl(), "otherDownloadUrl");
    // crypter should not be changed
    assertEquals(segmentZKMetadata.getCrypterName(), "crypter");
    assertEquals(segmentZKMetadata.getSegmentUploadStartTime(), -1);
    assertEquals(segmentZKMetadata.getSizeInBytes(), 10);

    // Refresh the segment with a different segment (different CRC)
    when(segmentMetadata.getCrc()).thenReturn("23456");
    when(segmentMetadata.getIndexCreationTime()).thenReturn(789L);
    // Add a tiny sleep to guarantee that refresh time is different from the previous round
    Thread.sleep(10);
    zkOperator.completeSegmentOperations(OFFLINE_TABLE_NAME, segmentMetadata, FileUploadType.SEGMENT, null, null,
        "otherDownloadUrl", "otherDownloadUrl", "otherCrypter", 100, true, true, httpHeaders);

    segmentZKMetadata = _resourceManager.getSegmentZKMetadata(OFFLINE_TABLE_NAME, SEGMENT_NAME);
    assertNotNull(segmentZKMetadata);
    assertEquals(segmentZKMetadata.getCrc(), 23456L);
    // Push time should not change
    assertEquals(segmentZKMetadata.getPushTime(), pushTime);
    // Creation time, refresh time, downloadUrl and crypter should change
    assertEquals(segmentZKMetadata.getCreationTime(), 789L);
    assertTrue(segmentZKMetadata.getRefreshTime() > refreshTime);
    assertEquals(segmentZKMetadata.getDownloadUrl(), "otherDownloadUrl");
    assertEquals(segmentZKMetadata.getCrypterName(), "otherCrypter");
    assertEquals(segmentZKMetadata.getSizeInBytes(), 100);
  }