@Override
    public void invoke() throws Exception {
        // --------------------------------------------------------------------
        // Initialize
        // --------------------------------------------------------------------
        initInputFormat();

        LOG.debug(getLogString("Start registering input and output"));

        try {
            initOutputs(getEnvironment().getUserCodeClassLoader());
        } catch (Exception ex) {
            throw new RuntimeException(
                    "The initialization of the DataSource's outputs caused an error: "
                            + ex.getMessage(),
                    ex);
        }

        LOG.debug(getLogString("Finished registering input and output"));

        // --------------------------------------------------------------------
        // Invoke
        // --------------------------------------------------------------------
        LOG.debug(getLogString("Starting data source operator"));

        RuntimeContext ctx = createRuntimeContext();

        final Counter numRecordsOut;
        {
            Counter tmpNumRecordsOut;
            try {
                InternalOperatorIOMetricGroup ioMetricGroup =
                        ((InternalOperatorMetricGroup) ctx.getMetricGroup()).getIOMetricGroup();
                ioMetricGroup.reuseInputMetricsForTask();
                if (this.config.getNumberOfChainedStubs() == 0) {
                    ioMetricGroup.reuseOutputMetricsForTask();
                }
                tmpNumRecordsOut = ioMetricGroup.getNumRecordsOutCounter();
            } catch (Exception e) {
                LOG.warn("An exception occurred during the metrics setup.", e);
                tmpNumRecordsOut = new SimpleCounter();
            }
            numRecordsOut = tmpNumRecordsOut;
        }

        Counter completedSplitsCounter = ctx.getMetricGroup().counter("numSplitsProcessed");

        if (RichInputFormat.class.isAssignableFrom(this.format.getClass())) {
            ((RichInputFormat) this.format).setRuntimeContext(ctx);
            LOG.debug(getLogString("Rich Source detected. Initializing runtime context."));
            ((RichInputFormat) this.format).openInputFormat();
            LOG.debug(getLogString("Rich Source detected. Opening the InputFormat."));
        }

        ExecutionConfig executionConfig = getExecutionConfig();

        boolean objectReuseEnabled = executionConfig.isObjectReuseEnabled();

        LOG.debug(
                "DataSourceTask object reuse: "
                        + (objectReuseEnabled ? "ENABLED" : "DISABLED")
                        + ".");

        final TypeSerializer<OT> serializer = this.serializerFactory.getSerializer();

        try {
            // start all chained tasks
            BatchTask.openChainedTasks(this.chainedTasks, this);

            // get input splits to read
            final Iterator<InputSplit> splitIterator = getInputSplits();

            // for each assigned input split
            while (!this.taskCanceled && splitIterator.hasNext()) {
                // get start and end
                final InputSplit split = splitIterator.next();

                LOG.debug(getLogString("Opening input split " + split.toString()));

                final InputFormat<OT, InputSplit> format = this.format;

                // open input format
                format.open(split);

                LOG.debug(getLogString("Starting to read input from split " + split.toString()));

                try {
                    final Collector<OT> output =
                            new CountingCollector<>(this.output, numRecordsOut);

                    if (objectReuseEnabled) {
                        OT reuse = serializer.createInstance();

                        // as long as there is data to read
                        while (!this.taskCanceled && !format.reachedEnd()) {

                            OT returned;
                            if ((returned = format.nextRecord(reuse)) != null) {
                                output.collect(returned);
                            }
                        }
                    } else {
                        // as long as there is data to read
                        while (!this.taskCanceled && !format.reachedEnd()) {
                            OT returned;
                            if ((returned = format.nextRecord(serializer.createInstance()))
                                    != null) {
                                output.collect(returned);
                            }
                        }
                    }

                    if (LOG.isDebugEnabled() && !this.taskCanceled) {
                        LOG.debug(getLogString("Closing input split " + split.toString()));
                    }
                } finally {
                    // close. We close here such that a regular close throwing an exception marks a
                    // task as failed.
                    format.close();
                }
                completedSplitsCounter.inc();
            } // end for all input splits

            // close all chained tasks letting them report failure
            BatchTask.closeChainedTasks(this.chainedTasks, this);

            // close the output collector
            this.output.close();
        } catch (Exception ex) {
            // close the input, but do not report any exceptions, since we already have another root
            // cause
            try {
                this.format.close();
            } catch (Throwable ignored) {
            }

            BatchTask.cancelChainedTasks(this.chainedTasks);

            ex = ExceptionInChainedStubException.exceptionUnwrap(ex);

            if (ex instanceof CancelTaskException) {
                // forward canceling exception
                throw ex;
            } else if (!this.taskCanceled) {
                // drop exception, if the task was canceled
                BatchTask.logAndThrowException(ex, this);
            }
        } finally {
            BatchTask.clearWriters(eventualOutputs);
            // --------------------------------------------------------------------
            // Closing
            // --------------------------------------------------------------------
            if (this.format != null
                    && RichInputFormat.class.isAssignableFrom(this.format.getClass())) {
                ((RichInputFormat) this.format).closeInputFormat();
                LOG.debug(getLogString("Rich Source detected. Closing the InputFormat."));
            }
        }

        if (!this.taskCanceled) {
            LOG.debug(getLogString("Finished data source operator"));
        } else {
            LOG.debug(getLogString("Data source operator cancelled"));
        }
    }
@Test
    void testFailingDataSourceTask() throws IOException {
        int keyCnt = 20;
        int valCnt = 10;

        this.outList = new NirvanaOutputList();
        File tempTestFile = new File(tempFolder.toFile(), UUID.randomUUID().toString());
        InputFilePreparator.prepareInputFile(
                new UniformRecordGenerator(keyCnt, valCnt, false), tempTestFile, false);

        super.initEnvironment(MEMORY_MANAGER_SIZE, NETWORK_BUFFER_SIZE);
        super.addOutput(this.outList);

        DataSourceTask<Record> testTask = new DataSourceTask<>(this.mockEnv);

        super.registerFileInputTask(
                testTask, MockFailingInputFormat.class, tempTestFile.toURI().toString(), "\n");

        boolean stubFailed = false;

        try {
            testTask.invoke();
        } catch (Exception e) {
            stubFailed = true;
        }
        assertThat(stubFailed).withFailMessage("Function exception was not forwarded.").isTrue();

        // assert that temp file was created
        assertThat(tempTestFile).withFailMessage("Temp output file does not exist").exists();
    }