001/*
002 * Licensed to the Apache Software Foundation (ASF) under one
003 * or more contributor license agreements.  See the NOTICE file
004 * distributed with this work for additional information
005 * regarding copyright ownership.  The ASF licenses this file
006 * to you under the Apache License, Version 2.0 (the
007 * "License"); you may not use this file except in compliance
008 * with the License.  You may obtain a copy of the License at
009 *
010 *     http://www.apache.org/licenses/LICENSE-2.0
011 *
012 * Unless required by applicable law or agreed to in writing, software
013 * distributed under the License is distributed on an "AS IS" BASIS,
014 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
015 * See the License for the specific language governing permissions and
016 * limitations under the License.
017 */
018package org.apache.hadoop.hbase.regionserver;
019
020import static org.junit.jupiter.api.Assertions.assertFalse;
021import static org.junit.jupiter.api.Assertions.assertThrows;
022import static org.junit.jupiter.api.Assertions.assertTrue;
023
024import java.io.IOException;
025import java.util.UUID;
026import org.apache.hadoop.conf.Configuration;
027import org.apache.hadoop.fs.FSDataOutputStream;
028import org.apache.hadoop.fs.FileSystem;
029import org.apache.hadoop.fs.Path;
030import org.apache.hadoop.hbase.CellBuilderType;
031import org.apache.hadoop.hbase.ExtendedCellBuilderFactory;
032import org.apache.hadoop.hbase.HBaseTestingUtil;
033import org.apache.hadoop.hbase.KeyValue;
034import org.apache.hadoop.hbase.io.hfile.HFile;
035import org.apache.hadoop.hbase.io.hfile.HFileContextBuilder;
036import org.apache.hadoop.hbase.testclassification.LargeTests;
037import org.apache.hadoop.hbase.testclassification.MiscTests;
038import org.apache.hadoop.hbase.util.Bytes;
039import org.apache.hadoop.hdfs.DistributedFileSystem;
040import org.junit.jupiter.api.AfterEach;
041import org.junit.jupiter.api.BeforeEach;
042import org.junit.jupiter.api.Tag;
043import org.junit.jupiter.api.Test;
044import org.junit.jupiter.api.TestInfo;
045
046/**
047 * Tests for failedBulkLoad logic to make sure staged files are returned to their original location
048 * if the bulkload have failed.
049 */
050@Tag(MiscTests.TAG)
051@Tag(LargeTests.TAG)
052public class TestSecureBulkloadListener {
053
054  private Configuration conf;
055  private HBaseTestingUtil htu;
056  private DistributedFileSystem dfs;
057  private final byte[] randomBytes = new byte[100];
058  private static final String host1 = "host1";
059  private static final String host2 = "host2";
060  private static final String host3 = "host3";
061  private static byte[] FAMILY = Bytes.toBytes("family");
062  private static final String STAGING_DIR = "staging";
063  private static final String CUSTOM_STAGING_DIR = "customStaging";
064  private String methodName;
065
066  @BeforeEach
067  public void setUp(TestInfo testInfo) throws Exception {
068    this.methodName = testInfo.getTestMethod().get().getName();
069    Bytes.random(randomBytes);
070    htu = new HBaseTestingUtil();
071    htu.getConfiguration().setInt("dfs.blocksize", 1024);// For the test with multiple blocks
072    htu.getConfiguration().setInt("dfs.replication", 3);
073    htu.startMiniDFSCluster(3, new String[] { "/r1", "/r2", "/r3" },
074      new String[] { host1, host2, host3 });
075
076    conf = htu.getConfiguration();
077    dfs = (DistributedFileSystem) FileSystem.get(conf);
078  }
079
080  @AfterEach
081  public void tearDown() throws Exception {
082    htu.shutdownMiniCluster();
083  }
084
085  @Test
086  public void testMovingStagedFile() throws Exception {
087    Path stagingDirPath = new Path(dfs.getWorkingDirectory(), new Path(methodName, STAGING_DIR));
088    if (!dfs.exists(stagingDirPath)) {
089      dfs.mkdirs(stagingDirPath);
090    }
091    SecureBulkLoadManager.SecureBulkLoadListener listener =
092      new SecureBulkLoadManager.SecureBulkLoadListener(dfs, stagingDirPath.toString(), conf);
093
094    // creating file to load
095    String srcFile = createHFileForFamilies(FAMILY);
096    Path srcPath = new Path(srcFile);
097    assertTrue(dfs.exists(srcPath));
098
099    Path stagedFamily = new Path(stagingDirPath, new Path(Bytes.toString(FAMILY)));
100    if (!dfs.exists(stagedFamily)) {
101      dfs.mkdirs(stagedFamily);
102    }
103
104    // moving file to staging
105    String stagedFile = listener.prepareBulkLoad(FAMILY, srcFile, false, null);
106    Path stagedPath = new Path(stagedFile);
107    assertTrue(dfs.exists(stagedPath));
108    assertFalse(dfs.exists(srcPath));
109
110    // moving files back to original location after a failed bulkload
111    listener.failedBulkLoad(FAMILY, stagedFile);
112    assertFalse(dfs.exists(stagedPath));
113    assertTrue(dfs.exists(srcPath));
114  }
115
116  @Test
117  public void testMovingStagedFileWithCustomStageDir() throws Exception {
118    Path stagingDirPath = new Path(dfs.getWorkingDirectory(), new Path(methodName, STAGING_DIR));
119    if (!dfs.exists(stagingDirPath)) {
120      dfs.mkdirs(stagingDirPath);
121    }
122    SecureBulkLoadManager.SecureBulkLoadListener listener =
123      new SecureBulkLoadManager.SecureBulkLoadListener(dfs, stagingDirPath.toString(), conf);
124
125    // creating file to load
126    String srcFile = createHFileForFamilies(FAMILY);
127    Path srcPath = new Path(srcFile);
128    assertTrue(dfs.exists(srcPath));
129
130    Path stagedFamily = new Path(stagingDirPath, new Path(Bytes.toString(FAMILY)));
131    if (!dfs.exists(stagedFamily)) {
132      dfs.mkdirs(stagedFamily);
133    }
134
135    Path customStagingDirPath =
136      new Path(dfs.getWorkingDirectory(), new Path(methodName, CUSTOM_STAGING_DIR));
137    Path customStagedFamily = new Path(customStagingDirPath, new Path(Bytes.toString(FAMILY)));
138    if (!dfs.exists(customStagedFamily)) {
139      dfs.mkdirs(customStagedFamily);
140    }
141
142    // moving file to staging using a custom staging dir
143    String stagedFile =
144      listener.prepareBulkLoad(FAMILY, srcFile, false, customStagingDirPath.toString());
145    Path stagedPath = new Path(stagedFile);
146    assertTrue(dfs.exists(stagedPath));
147    assertFalse(dfs.exists(srcPath));
148
149    // moving files back to original location after a failed bulkload
150    listener.failedBulkLoad(FAMILY, stagedFile);
151    assertFalse(dfs.exists(stagedPath));
152    assertTrue(dfs.exists(srcPath));
153  }
154
155  @Test
156  public void testCopiedStagedFile() throws Exception {
157    Path stagingDirPath = new Path(dfs.getWorkingDirectory(), new Path(methodName, STAGING_DIR));
158    if (!dfs.exists(stagingDirPath)) {
159      dfs.mkdirs(stagingDirPath);
160    }
161    SecureBulkLoadManager.SecureBulkLoadListener listener =
162      new SecureBulkLoadManager.SecureBulkLoadListener(dfs, stagingDirPath.toString(), conf);
163
164    // creating file to load
165    String srcFile = createHFileForFamilies(FAMILY);
166    Path srcPath = new Path(srcFile);
167    assertTrue(dfs.exists(srcPath));
168
169    Path stagedFamily = new Path(stagingDirPath, new Path(Bytes.toString(FAMILY)));
170    if (!dfs.exists(stagedFamily)) {
171      dfs.mkdirs(stagedFamily);
172    }
173
174    // copying file to staging
175    String stagedFile = listener.prepareBulkLoad(FAMILY, srcFile, true, null);
176    Path stagedPath = new Path(stagedFile);
177    assertTrue(dfs.exists(stagedPath));
178    assertTrue(dfs.exists(srcPath));
179
180    // should do nothing because the original file was copied to staging
181    listener.failedBulkLoad(FAMILY, stagedFile);
182    assertTrue(dfs.exists(stagedPath));
183    assertTrue(dfs.exists(srcPath));
184  }
185
186  @Test
187  public void testDeletedStagedFile() throws Exception {
188    Path stagingDirPath = new Path(dfs.getWorkingDirectory(), new Path(methodName, STAGING_DIR));
189    if (!dfs.exists(stagingDirPath)) {
190      dfs.mkdirs(stagingDirPath);
191    }
192    SecureBulkLoadManager.SecureBulkLoadListener listener =
193      new SecureBulkLoadManager.SecureBulkLoadListener(dfs, stagingDirPath.toString(), conf);
194
195    // creating file to load
196    String srcFile = createHFileForFamilies(FAMILY);
197    Path srcPath = new Path(srcFile);
198    assertTrue(dfs.exists(srcPath));
199
200    Path stagedFamily = new Path(stagingDirPath, new Path(Bytes.toString(FAMILY)));
201    if (!dfs.exists(stagedFamily)) {
202      dfs.mkdirs(stagedFamily);
203    }
204
205    // moving file to staging
206    String stagedFile = listener.prepareBulkLoad(FAMILY, srcFile, false, null);
207    Path stagedPath = new Path(stagedFile);
208    assertTrue(dfs.exists(stagedPath));
209    assertFalse(dfs.exists(srcPath));
210
211    dfs.delete(stagedPath, false);
212
213    // moving files back to original location after a failed bulkload
214    assertThrows(IOException.class, () -> listener.failedBulkLoad(FAMILY, stagedFile));
215  }
216
217  private String createHFileForFamilies(byte[] family) throws IOException {
218    HFile.WriterFactory hFileFactory = HFile.getWriterFactoryNoCache(conf);
219    Path testDir =
220      new Path(dfs.getWorkingDirectory(), new Path(methodName, Bytes.toString(family)));
221    if (!dfs.exists(testDir)) {
222      dfs.mkdirs(testDir);
223    }
224    Path hfilePath = new Path(testDir, generateUniqueName(null));
225    FSDataOutputStream out = dfs.createFile(hfilePath).build();
226    try {
227      hFileFactory.withOutputStream(out);
228      hFileFactory.withFileContext(new HFileContextBuilder().build());
229      HFile.Writer writer = hFileFactory.create();
230      try {
231        writer.append(new KeyValue(ExtendedCellBuilderFactory.create(CellBuilderType.DEEP_COPY)
232          .setRow(randomBytes).setFamily(family).setQualifier(randomBytes).setTimestamp(0L)
233          .setType(KeyValue.Type.Put.getCode()).setValue(randomBytes).build()));
234      } finally {
235        writer.close();
236      }
237    } finally {
238      out.close();
239    }
240    return hfilePath.toString();
241  }
242
243  private static String generateUniqueName(final String suffix) {
244    String name = UUID.randomUUID().toString().replaceAll("-", "");
245    if (suffix != null) name += suffix;
246    return name;
247  }
248}